Entity-Centric Reinforcement Learning for Object Manipulation from Pixels
Authors: Dan Haramati
Tal Daniel
Aviv Tamar
Venue: ICLR 2024 spotlight (2024)
Keywords: deep reinforcement learning
visual reinforcement learning
object-centric
robotic object manipulation
compositional generalization
Primary area: reinforcement learning
TL;DR: We introduce a structured visual goal-conditioned RL approach for multi-object manipulation, showcasing agents that generalize from 3 to over 10 objects.

Abstract: Manipulating objects is a hallmark of human intelligence, and an important task in domains such as robotics. In principle, Reinforcement Learning (RL) offers a general approach to learn object manipulation. In practice, however, domains with more than a few objects are difficult for RL agents due to the curse of dimensionality, especially when learning from raw image observations. In this work we propose a structured approach for visual RL that is suitable for representing multiple objects and their interaction, and use it to learn goal-conditioned manipulation of several objects. Key to our method is the ability to handle goals with dependencies between the objects (e.g., moving objects in a certain order). We further relate our architecture to the generalization capability of the trained agent, based on a theoretical result for compositional generalization, and demonstrate agents that learn with 3 objects but generalize to similar tasks with over 10 objects. Videos and code are available on the project website: https://sites.google.com/view/entity-centric-rl

Published as a conference paper at ICLR 2024
ENTITY-CENTRIC REINFORCEMENT LEARNING FOR OBJECT MANIPULATION FROM PIXELS
Dan Haramati, Tal Daniel, Tamar
Department of Electrical and Computer Engineering, Technion Israel Institute of Technology (haramati, taldant 1m)mcampus techni on.iiin..e ac. i1; avivt@techiiiiniooo on. ac.
ABSTRACT
Manipulating objects is a hallmark of human intelligence, and an important task in domains such as roboties. In principle, Reinforcement Learning (RL) offers a general approach to learn object manipulation. In practice, however. domains with more than a few objects are difficult for RL agents due to the curse of dimen- sionality, especially when learning from raw image observations. In this work we propose structured approach for visual RL that is suitable for represent- ing multiple objects and their interaction, and use it to learn goal-conditioned manipulation of several objects. Key to our method goals with dependencies between the objects (e.g., moving objects in cer- tain order). We further relate our architecture to the generalization capability of the trained agent, based on theoretical result for compositional generaliza- tion, and demonstrate agents that learn with objects but generalize to similar tasks with over 10 objects. Videos and code are available on the project website: https: google. ott.rr/reeooooeoooo.ereeeeeerrrr
the ability to handle
/etttttttttttttteeeee centric-t
INTRODUCTION
Deep Reinforcement Learning (RL) has been successfully applied various domains such as video games (Mnih et al. 2013) and robotic manipulation (Kalashnikov et al. 2018). While some studies focus on developing general agents that can solve wide range of tasks (Schulman al. 2017). the difficulty of particular problems has motivated studying agents that incorporate structure into the learning algorithm (Mohan et al., (023). Object manipulation our focus in this work is clear example for the necessity of structure: as the number of degrees of freedom of the system grows exponentially with the number objects, the curse of dimensionality inhibits standard approaches from learning. Indeed, off-the shelf RL algorithms struggle with learning to manipulate even modest number of objects (Zhou et al. 2022).
The factored Markov decision process formulation (factored MDP, Guestrin et al. 2003) factorizes the full state of the environment into the individual states of each object, or entity, in the system S;: S= & S2 @ ... 3 Sy. A key observation is that if the state transition depends only subset of the entities (e.g., only objects that are physically near by), this structure can be exploited to simplify learning. In the context of deep RI, several studies suggested structured representations based permutation invariant neural network architectures for the policy and Q-value functions (Li et al., 2020 Zhou al. 2022). These approaches require access the true factored state of the system.
For problems such as robotic manipulation with image inputs, however, how to factor the state (images of robot and objects) into individual entities and their attributes (positions, orientation, etc.) is not trivial. Even the knowledge of which attributes exist and are relevant to the task is typically not given in advance. This problem setting therefore calls for learning- based approach that can pick up the relevant structure from data. As with any learning-basee method, the main performance indicator generalization. In our setting, generalization may be measured with respect to 3 different factors of variation: (1) the states of the objects in the system, (2) different types of objects, and (3) different number objects than in training, known as compositional generalization (Lin et al. 2023).
Our main contribution in this work is goal-conditioned RL framework for multi-object manipu- lation from pixels. Our approach consists of two components. The first is an unsupervised object- centric image representation (OCR), which extracts entities and their attributes from image data. The second and key component is a Transformer baseed architecture for the policy and Q-function neural networks that we name Entity Interaction Transformer (EIT). Different from previous work

---

Published as conference paper at ICLR 2024
such Zadaianchuk et al. (2020), EIT is structured such that it can easily model not only lations between goal and state entities, but also entity-entity interactions in the current state. This allows us to learn tasks where interactions between objects are important for achieving the goal, for example, moving objects particular order.
Furthermore, the EIT does not require explicit matching between entities in different images, and can therefore handle multiple images from different viewpoints seamlessly. As we find out, multi- view perception is crucial the RL agent for constructing an internal "understanding" of the 3D scene dynamics from 2D observations in a sample efficient manner. Combined with our choice of Deep Latent Particles IDLP,Damice & Tamar 2022) image representations, we demonstrate what is, to the best of our knowledge, the most accurate object manipulation from pixels involving more than two objects, or with goals that require interactions between objects.
Finally, we investigate the generalization ability of our proposed framework. Starting with formal definition of compositional generalization in RL, we show that self-attention based Q-value func- tions have fundamental capability to generalize, under suitable conditions the task structure. This result provides sound basis for our Transformer-based approach. Empirically, we demon- strate that an EIT trained on manipulating up to 3 objects can perform well on tasks with up to 6 objects, and in certain tasks we show generalization to over 10 objects.
Environment Overview
Agent Vieu 1 Front
Agent Vieu Side
Figure I: The environment we used for our experiments (left) and how the agent perceives it (middle, right), colored keypoints are the position altribute Zp of particles from the DLP representation.
2 RELATED WORK
Latent-based Visual RL: Unstructured latent representations consisting of single latent vector have been widely employed, both in model-free (Levine al.) 2016 Nair al. 2018 Pong al.)
2019 Yarats et al.] 2021) and model-based (Hafner et al., 2023 Mendonca et al. [021] settings. However, in manipulation tasks involving multiple objects. this approach falls short compared to
models based on structured representations (Gmelin et al. 2023 Zadaianchuk et al. 2020).
Object-centric RL: Several recent works explored network architectures for structured representa- tion of the state in model free (Li et al. 2020 Zhou et al. 2022, Mambelli et al. 2022 Zadaianchuk
al. 2018) RL as well as model- planning methods (Sancaktar et al.) (022). Compared to our approach, the aforementioned methods assume access to ground-trut trunt states, while we learn representations from images.
al. 2022) and model-based (Sanchez-Coonalle
Object-centric RL from Pixels: Several works have explored adopting ideas from state-based structured representation methods to learning from visual inputs. COBRA (Watters et al.) 2019
STOVE (Kossen al. 2019) NCS (Chang et al.] 2023). FOCUS (Ferraro al. 2023) DAFT-
RL (Fngg& & Magliacane| 2023), HOWM (Zhao et al. 2022) and Driess et al. (2023) leam object- centric world models which they for planning or policy learning to solve multi-object tasks. In contrast to these methods, our method is trained in a model-free setting, makes less assump- tions on the problem and less complex, allowing simple integration with various object-centric representation models and standard online or offline Q-learning algorithms. OCRL (Yoon et al.) 2023), Heravi et al. (2023) and yan Bergen & Lanillos (2022) have investigated slot-based rep- resentations (Greff et al.] 2019 Locatello et al. 2020, [Singh et al.] 2022) Traub et al., 2022) for manipulation tasks model-free, imitation learning and active inference settings, respectively. The above methods have demonstrated the clear advantages of object-centric representations over non- object-centric alternatives. In this work, we utilize particle-basee based image representations (Daniel & Tamar 2022) and extend these findings by tackling more complex manipulation tasks and show- casing generalization capabilities. Closely related to our work. SMORL (Zadaianchuk et al. 2020) employs SCALOR (Jiang al. 2020). patch-based image representation for goal-conditioned
2

---

Published as a conference paper at ICLR 2024
manipulation tasks. The key assumption in SMORL is that goal-conditioned multi-object tasks can be addressed sequentially (object by object) and independently, overlooking potential interactions among objects that might influence reaching goals. Our model, the other hand, considers inter- action between entities, leading to improved performance and generalization, as we demonstrate in thorough comparison with SMORL (see Section
3 BACKGROUND
Cond-Condi Reinforcement Learning (GCRL): RL considers Markov Decision Process (MDP, Puterrann (2014)) defined as M = (S, A, P.r, Y, Po), where S represents the state space, A the action space, P the environment transition dynamics, the reward function, Y the discount factor and Po the initial state distribution. GCRL additionally includes a goal space G. The agent seeks learn policy x 8
A that maximizes the expected return E., [- 7'rr], where rs r (St, g) S x g - IR is the immediate reward at time when the state and goal are 8, and g.
Deep Latent Particles (DLP): DI.P (Daniel & Tamar 2022 2023) is an unsupervised object-centric model for image representation. DLP provides a disentangled latent space structured as a set of (6++()))) where K is the number of particles, 2p E R² particles 2 = {(Zp, Za, 2d, 2t, 2f) ) K-1 E &R is the position of the particle as (Ix, y) coordinates in Euclidean pixel space. 2g E R2 is a scale attribute containing the (z. y) dimensions of the bounding-boo around the particle, 2d R is pixel space "depth" attribute used to signify which particle is in front of the other in case there an overlap, 2z R is a transparency attribute and 2, E TR' are the latent features that encode the visual appearance of a region surrounding the particle, where / is the dimension of learned visual features. See Appendix A for an extended background.
4
METHOD
We propose an approach to solving goal-conditioned multi-object manipulation tasks from images. Our approach is entity-centric each represented by a latent feature vector, and learn the relationships between them. Our method consists of the following 2 components: (1) Object-Centrit Representation (OCR) of Images[4.
is structured to decompose the input into individual entities,
We extract a representation of state and goal images consisting of set of latent vectors using a pretrained model; (2) Entity Interaction Transformer (EIT) 4.2 We feed the sets of latent vectors, extracted from multiple viewpoints, to a Transformer- based architecture for the RL policy and Q-function neural networks. These two components can be used with standard RL algorithms to optimize a given reward function. We additionally propose a novel image-based reward that is based on the OCR and the Chamfer distance, and corresponds to moving objects to goal configurations. We term this Chamfer Reward, and it enables learning entirely from pixels. We begin with a general reasoning that underlies our approach.
The complexity tradeoff between representation learning and decision making: An important observation is that the representation learning (OCR) and decision making (EIT) problems de- pendent. Consider for example the task of moving objects with a robot arm, as in our experiments. Ideally, the OCR should output the physical state of the robot and each object, which is sufficient for optimal control. However, identifying that the robot is a single entity with several degrees of free- pertains to
dom, while the objects are separate entities, is difficult to learn just from image data, the dynamics of the system. The relevant properties of each object can also be task dependent - for example, the color of the objects may only matter if the task's goal depends on it. Alternatively, one may settle for a much leaner OCR component that does not understand the dynamics of the objects nor their relevance to the task, and delegate the learning of this information to the EIT. Our design choice in this paper is the latter, i.e., a lean OCR and an expressive EIT. We posit that this design allows to (1) easily acquire an OCR, and (2) handle multiple views and mismatches between the visible objects in the current state and the goal seamlessly. We next detail our design.
4.1 OBJECT- CENTRIC REPRESENTATION OF IMAGES
The first step in our method requires extracting a compact disentangled OCR from raw pixel obser- vations. Given a tuple of image observations from K different viewpoints of the state (Ti.-. ....... ... I&x) and goal (I (Ie...., 1% (DLP) DDanil & Tamar 2022 2023) model, extracting a set of M vectors {pr, LM 1. k indexing the viewpoint of the source image, which will refer to as (latent) particles. We denote particle
we process each image separately using a pretrained Deep Latent Particles
3

---

Published as a conference paper at ICLR 2024
m of state image I by pi: and of goal image 1% by gm We emphasize that there is no alignment between particles from different views (e.g- Pm: Pi can correspond to different objects) or between state and goal of the same view (c.g., between p,, and g..). The vectors pm. 9 E RR contain different attributes detailed in the DLP section of the Background 3 In contrast previous object- centric approaches that utilize patch-based (Zadaianchuk et al. 2020) or slot-based (Yoon et al.) 2023) representations,
adopt DLP, which has recently demonstrated state-of- state-of-the-ar the-hee-rrt performance in single-image decomposition and various downstream tasks. The input to the goal-conditioned JM m m==1 {
... { For the Q-function, Q(s, a, 9), an action particle Pa E RR+1/ is added to the input, obtained by learning a projection from the action dimension da to the latent particle dimension.
policy is tuple of the 2K sets: (PmJm==:>>
}M
Pretraining the DLP: In this work, we pretrain the DLP from dataset of images collected by a random policy. We found that in all our experiments, this simple pretraining was sufficient to obtain well performing policies even though the image trajectories in the pretraining data are different from trajectories collected by an optimal policy. We attribute this to the tradeoff described above the lean single-imaee based DLP OCR is complemented by a strong EIT that can account for dynamics necessary to solve the task.
4.2 ENTITY-CENTRIC ARCHITECTURE
We next describe the EIT, which processes the OCR entities into an action or Q-value. As men- tioned above, our choice of a lean OCR requires the EIT account for the dynamics of the entities and their relation to the task. In particular, we have the following desiderata from the EIT: (1) Per- mutation Invariance: (2) Handle goal-based RL; (3) Handle multiple views; (4) Compositional generalization. Proper use of the attention mechanism provides us with (1). The main difficulty in (2) and (3) that particles in different views and the goal are not necessarily matched. Thus, we designed our EIT to seamlessly handle unmatched entities. For (4), we compose the EIT using )
Transformer (Vaswani et al. 2017) blocks. As we shall show in Section 4.4 this architecture has a fundamental capacity to generalize. An outline of the architecture is presented in Figure ? Com- pared to previous goal-conditioned methods' use the attention mechanism, we use it to explicitly model relationships between entities from both state and goal across multiple viewpoints and do not assume privileged entity-entity matching information. A more detailed comparison can be found in Appendix
E
We now describe the EIT in detail:
Input- The EIT policy receives the latent particles extracted from both the current state images and the goal images input. We inject information the source viewpoint of each state and goal parti- cle with an additive encoding which is learned concurrently with the rest of the network parameters. Forward State particles (1p)) MM
/ .... {e) } are processed by a sequence of Transformer blocks: SA + CA + SA + AA, denoting Self-Attention SA, Cross- Aggregation-Attention AA, followed by an MLP.
CA,, and
Goal-conditioning We condition on the goal particles k.......... (1.. N k==) using the CA block between the state (provide the queries) and goal (provide the keys and values) particles. Permutation Invariant Output- The AA block reduces the set to single-vector output and is implemented with a CA block between a single particle with learned features (provides the query) and the output particles from the previous block (provide the keys and values). This permutation in- variant operation on the processed state particles, preceded by permutation equivariant Transformer blocks, results in output that is invariant to permutations of the particles in each set. The aggre- gated particle is input to an MLP, producing the final output (action/value). Action Entity The EIT Q-function, in addition the state and goal particles, receives action as input. The action is projected to the dimension of the particles and added to the input set. Treating the action as an individual entity proved to be significant design choice, see ablation study in C.4
4.3 CHAMFER REWARD
We define an image-based reward from the DLP representations of images as the Generalized Density-Aware Chamfer (GDAC) distance between state and goal particles, which we term Cham- fer reward. The standard Chamfer distance is defined between two sets and measures the average
We provide detailed exposition attention in Appendix A.3
4

---

Published as a conference paper at ICLR 2024
View Encoding
Self-Atrention (SA)
Cross-Attention (CA)
Q-val/Action
Token
Particies View 1
SA
SA
CA
Input Tokens
Multi Mult-Head
Add B Norm
20l Norm
press min
poy
LuAo i pole
Particles View
Aitertion
Foneard
Food
pad
Action
Goal Tokens
MLP
Encoding
Goal
Goal
Particles
Particles
View
Figure 2: Outline of the Entity Interaction Transformer (EIT) Sets of state and goal particles from multiple views with additive view encoding input sequence of Transformer blocks. For the Q-function, an action particle added. We condition on goals with cross-attention. Attention-baser aggregation reduces the set single vector, followed by MLP that produces the final output.
distance between each entity in one set to the closest entity in the other. The Density-Aware Cham- fer distance (Wu et al., 202T) takes into account the fact that multiple entities from one set can be mapped to the same entity in the other set by reweighing their contribution to the overall distance accordingly. The Generalized Density-Awes
Chamfer distance, decouples the distance function that is used to match between entities and the used calculate the distance between them. For space considerations, we elaborate on the Chamfer reward Appendix B
4.4 COMPOSITIONAL GENERALIZATION
In context, Compositional Generalization (CG) refers to the ability of policy to perform multi-object task with a different number of objects than it was trained on. In this section, we formally define notion of CG for RL, and show that under sufficient conditions, self-attention based Q-function architecture can obtain it. This result motivates our choice of the Transformer- based EIT, and shall investigate further in experiments. We give main result here, and provide an in-depth discussion and full proofs in Appendix F
Let Š denote the state space single object, and consider an N-object MDP as an MDP with a factored state space S S₁ @ S₂ & ... @ Sv, Vi S, S, action space A, reward function r and discount factor y. Without loss of generality, assume that
Rmnr 1.
Naturally, tasks which are completely different for every N would not be amenable for CG. We constrain the set of tasks by assuming certain structure of the optimal Q-function. We say that class of functions admits compositional generalization if, after training a Q-function from this class on M objects, the error on M + k objects grows most linearly k. The more expressive the function class that admits CG is. the greater the chance that it would apply for the task at hand. In the following, we prove our main result CG for the class self-attention functions. This class is suitable for tasks where the optimal policy involves a similar procedure that must be applied to each object, such as the tasks we consider Section 5
Assumption 1. AS E S Va E A, VN E N we have: Q' (51 a) ΣN Q (S1,..., SN, a), where
Qi (S1, u) ΣN (x. Σj= (8,, 8), ()),0), ER.
The following result shows that when obtaining e-optimal f-uuctio for up M objects where the attention weights are 8-optimll, the oub-optimality w.r.t. M + objects grows at most linearly in k.
5

---

Published as a conference paper at ICLR 2024
Theorem 2. Let Assumption hold. Let Q be approximation of Q with the same structure.
Assume that Vs E SN va E A, VN E [1, M] we have Q (81, .... SN, a) Q" ($1, SN, <s,
0(861311)) and by a(861851a)
(a11)))) Ejsr a (81185ia)
< 6. Then, E SM+h Va E A, Ak [1, M - 11:
Q" (s1, SM+k. 33 3 + k) +2 6.
(S1...,, SM+ a)
4.5 TRAINING AND IMPLEMENTATION DETAILS
We developed our method with the off- policy algorithm TD3 (Fujimoto et al. 2018) along with hind- sight experience replay (HER, Andrychowicz et al. In principal, our approach actor-critic algorithms or to the online setting, and can be used with any deep Q-learning algorithm, online or offline. pre-train a single DLP model on images from multiple viewpoints of rollouts collected with random policy. We convert state and goal images to object-centric latent represen- tations with the DLP encoder before inserting them to the replay buffer. We use our EIT architecture for all policy and Q-function neural networks. Further details and hyper-parameters can be found in Appendix D Our code publicly available on http::///gg.bb....o///tttttttt//
limited to
com Dttp//nnnnrmttttmmmm
5 EXPERIMENTS
We design experimental setup address the following aspects: (1) benchmarking method on multi-object manipulation tasks from pixels; (2) assessing the significance of accounting for interac- tions between entities for the RL agent's performance; (3) evaluating the scalability of our approach to increasing number of objects; (4) analyzing the generalization capabilities of our method.
Environments We evaluate method on several simulated tabletop robotic object manipulation environments implemented with IsaacGym (Makoviychuk et al. 2021 The environment includes robotic set in front of table with varying number of cubes in different colors. The agent observes the state of the system through a number of cameras in fixed locations, and performs actions in the form of deltas in the end effector coordinates a (Areer Agee, Azee). At the beginning of each episode. the cube positions are randomly initialized on the table, and goal configuration is sampled similarly. The goal of the agent push the cubes match the goal configuration. We categorize suite of tasks follows (see Figure 3):
N- N-Cubes: Push N different-colored cubes their goal location. Ad edjacent-Goal all cubes are adjacent. This task requires accounting for interactions between objects.
A 3-Cubes setting where goals are sampled randomly on the table such that
Small-Tables A 3-Cubes setting where the table is substantially smaller. This task requires to accurately account for all objects in the scene at all times, to avoid pushing blocks off the table.
Ordered-Push. A 2-Cubes setting where narrow corridor is set top of the table such that its width can only fit a single cube. We consider two possible goal configurations: red cube in the rear of the corridor and green cube in the front, or vice versa. This task requires fulfill the goals in certain order, otherwise the agent fails (pulling a block out of the corridor is not possible). Push-2T: Push T-shaped blocks to a single goal orientation.
E-Cubes
Adjocent-Goola Drdered-Push
Secll-Table
Puah-2T
Figure 3: The simulated environments used for experiments in this work.
Reward The reward calculated from the ground-truth state of the system, which refer to as the ground-truth (GT) reward, the mean negative L2 distance between each cube and its desired goal position the table. The image-based reward calculated from the DLP OCR for method is the negative GDAC distance (see Eq. T between state and goal sets of particles, averaged over viewpoints. Further reward details be found in the Appendix D.I
Evaluation Metries We evaluate the performance the agents on several metrics. In this environ- ment, define episode success if all N objects threshold distance from their desired

---

Published as a conference paper at ICLR 2024
goal. The metric most closely captures task success, but does not capture intermediate success or timestep efficiency. To this end, we additionally evaluate based on success fraction, maximum object distance, average object distance and average return. For a formal definition of these metrics Appendix D.1 All results show means and standard deviations across 3 random seeds.
Baselines We compare our method with the following baselines: Unstructured - A single-vector latent representation of images using a pre-trained VAE is ex- tracted from multiple viewpoints from both state and goal images and then concatenated and fed to an MLP architecture for the policy and Q-function neural networks. This baseline corresponds to methods such as Nair et al. (2018), with the additional multi-view view data as in our method.
SMORL We re- re-implement SMORL (Zadaianchuk et al. (2020), extending it to multiview inputs and tune its paree-rarameeers for the environments in this work. Re-implementation details are available in Appendix [D.4] We use DLP as the pre-trained OCR for this method for a fair comparison. Note that image-based SMORL cannot utilize GT reward since it requires matching the particle selected from the goal image at the beginning of each training episode to the corresponding object in the environment. We therefore do not present such experiments.
Object-centric Pretraining All image-based methods in this work utilize pre-trained unsupervised image representations trained on data collected with a random policy. For the object-centric meth- ods, train a single DLP model on data collected from the 6-Cubes environment. We found it generalizes well to fewer objects and changing backgrounds (e.g., smaller table). For the Push-2T T environment we trained DLP on data collected from Push-T (single block) which generalized well to 2 blocks. Figure [4] illustrates the DLP decomposition of single training image. For the non- OCR baselines, we use mixture of data from the 3--ubbs environments to learn latent representation with a B-VAE (Higgins et al. 2017). More information on the architectures and hyper-parameters is available in Appendix D.3
Experiment Outline We separate our investigation into 3 parts. In the first part, we focus on the design of our OCR and EIT. and how it handles complex interactions between objects. We study this using the GT reward, to concentrate on the representation learning question. When comparing with baselines, we experiment with both the OCR and a ground-truth state representation") To prove that our method indeed handles complex interactions, we shall show that our method with the OCR outperforms baselines with GT state on complex tasks. In the second part, we study compositional the GT reward, with similar motivation as above. Finally, in the third part we evaluate our method using the Chamfer reward. The Chamfer reward is calculated by filtering out particles that do not correspond objects of interest such the agent. Due to lack of space this part is detailed in Appendix B.T An ablation study of key design choices such as incorporating multi-view inputs can be found in Appendix C.4
generalization. In this case we also
5.1 MULTI-OBJECT MANIPULATION
We evaluate the different methods with GT rewards on the environments detailed above. Results are presented in Figure and Table We observe that with single object, all methods succeed, yet the unstructured baselines are less sample efficient. With more than 1 object, the image-based unstruc- tured baseline not able to learn at all, while the unstructured state-based baseline significantly outperformed by the structured methods. Our method and SMORL reach similar performance in the state-based setting, SMORL being more sample efficient. This is expected as SMORL essentially learns single object manipulation, regardless of the number of cubes in the environment. Notably, on 3-Cubes, our image-base method surpasses the unstructured state-based method.
In environments that require interaction between objects adjacent-Goall Small-Table and Ordered-Push - our method outperforms SMORL using state input. Moreover, with image inputs our method outperforms SMORL. with state inputs (significantly on Ordered-Push, yet marginally on Small-Table, bdjacent-Goalle demonstrating that SMORL is fundamen- tally limited in performing these more complex tasks.
We present preliminary results on the Push- 2T environment. For a visualization of the task and performance results see Figure 3 The success in this task demonstrates the following additional ca- pabilities of our proposed method: (1) Handling objects that have more complex physical properties
Thee ground truth state is 8, (a,, v.). gi = (z. v the (z, v) coordinates of the state and goal of entity i respectively. We detail how this state is input the networks in Appendix D.I
7

---

Published as a conference paper at ICLR 2024
100
75
Succes Rate
75
Success Bute
75
30
25
25
0
403000
0.0
0.5
Enwitomment mesteps
Envrronmntt Timesteos
Envrronmnnt Timesteas
(a) 1-Cube
(b) 2-Cubes
(c) 3-Uubes
Method
Success Rate
80
Unstructured
Success Rate
80
SMORL
40
Ours
40
Input Type
State
8.0
0.5
2.0 Environment Timesteps
3.0
Image
Environment Timesteps
(d) Ordered-Push
(c) Legend
(f) Small-Table
Figure 4: Success Rate vs. Environment Timesteps Values calculated on 96 randomly sampled goals. Methods with imput type 'State' are presented in dashed lines and learn from GT state ob- servations, otherwise, from images. Our method performs better than or equivalently to the best performing baseline each category (state/image-ban -useed. the environments requiring object teraction ((d), (f)), our method achieves significantly better performance than SMORL Notably, our image-based based method matches/surpasses state state-based SMORL.
Method
Success Rate
Success Fraction
Max Obj Dist
Avg Obj Dist
Avg Return
Ours (State)
0.963 + 0.005
0.982 = 0.005
0.022 + 0.002
0.014 = 0.002
-0.140 0.008
SMORL (State)
0.716 + 0.006
0.863 + 0.005
0.063 + 0.003
0.031 + 0.001
-0.233 + 0.004
Ours (Image)
0.710 + 0.016
0.883 + 0.005
0.044 + 0.003
0.027 + 0.001
-0.202 + 0.007
Table I: Performance Metrics for Adjacent Methods trained on 3-Cubes and eval- uated on Adjacent -Goa 1s. Values calculated on 400 random gouls per random seed.
which affect dynamics. (2) The ability of the EIT to infer object properties that not explicit in the latent representation (i.e. inferring orientation from latent particle visual attributes), and accurately manipulate them in order to achieve desired goals.
Figure 5: Left Rollout of agent trained the Push-27 task. Right Distribution of object angle difference (radians) from goal. Values of 400 episodes with randomly initialized goal and initial configurations.
5.2 COMPOSITIONAL GENERALIZATION
In this section, investigate method's 's ability to achieve zero-shot compositional generalization. Agents were trained from images with our method using GT rewards and require purely image inputs during inference. We present several inference scenarios requiring compositional generalization:
Different Number of Cubes than in Training We train an agent on the 3-Cubes environment and deploy the obtained policy on the N-Cubes environment for N [1. []. Colors sampled uniformly out of 6 options and are distinct for each cube. Visual results on 6 cubes are presented in Figure 6 (left) and evaluation metrics in Table (Appendix). We see that agent generalizes
8

---

Published as a conference paper at ICLR 2024
-0.12
t
80
100
Gool
0.30
4..0
00.00
Number of Cubes
Figure 6: Left Left-Rollout of agent trained on 3 cubes generalizing to cubes. Right The average return of an agent trained on 3 cubes vs. the number of cubes in the environment it was deployed in during inference. Values are averaged over 400 episodes with randomly initialized goal and initial configurations. Note that the graph is approximately linear, corresponding with Theorem
changing number of objects with some decay performance as the number of objects grows. Notably, the decay in the average return, plotted in Figure 6 (right), is approximately linear in the number cubes, which corresponds with Theorem
Cube Sorting We train an agent on the 3 3-Cubes environment with constant cube colors (red, green, blue). During inference, we provide goal image containing X S cubes of different colors and then deploy the policy the cubes around each goal cube position with matching color, and is also able to perform the task with cube colors unseen during RL training. Visual results 12 cubes in 3 colors are presented in Figure 7 We find these results exceptional, as they require compositional generalization from both the EIT policy (trained on cubes) and the DLP model (trained on 6 cubes) to significantly more cubes, occupying a large portion the table's surface.
environment containing 4XX cubes, 4 of each color. The agent sorts
t=1
t = 40
80
= 100
= 120
Gool
Figure 7: Rollout of an agent trained on 3 cubes, then provided a goal image containing 3 different colored cubes and deployed in environment with 12 cubes, 4 of each color. The agent sorts the cubes around each goal position with matching color.
Further results and analysis of our agent's generalization capabilities such generalizing to cube properties not seen during training are detailed Appendix C Videos are available on our website.
CONCLUSION AND FUTURE WORK
In this work we proposed an RL framework for object manipulation from images, composed of an off- tff-the-shelf object-centric image representation, and novel Transformer-bast basee policy architecture that can account for multiple viewpoints and interactions between entities. We have also investigated compositional generalization, starting with a formal motivation, and concluding with experiments that demonstrate non-trivial generalization behavior of our trained policies.
Limitations: Our approach requires pretrained DLP. In this work, our DLP was pretrained from data collected using random policy, which worked well on all of our domains. However, more complex tasks may require more sophisticated pretraining, or an online approach that mixes training the DI.P with the policy. Our results with Chamfer rewards show worse performance than with ground truth reward. This hints that directly running our method on real robots may be more difficult than in simulation due to the challenges of reward design.
Future work: One interesting direction is understanding what environments are solvable using our approach. While the environments investigated here more complex than in previous studies, environments with more complex interactions between objects, such as with interlocking objects or articulated objects, may be difficult solve due to exploration challenges. Other interesting direc- tions include multi-modal goal specification (e.g., language), and more expressive sensing (depth cameras. force sensors, etc.), which could be integrated as additional input entities the EIT.

---

Published as a conference paper at ICLR 2024
7 REPRODUCIBILITY STATEMENT
We are committed to ensuring the reproducibility of our work and have taken several measures to facilitate this. Appendix [ contains detailed implementation notes and our experiments. Furthermore, we make our code openly available to the community to facilitate further research in following repository: /tttps://githhe codebase includes the implementation of the environments and proposed method, well as scripts for reproducing the experiments reported in this paper. In addition, Appendix [F contains in depth details of the theoretical portion of this paper which include the full proof for Theorem 2
used in
com CommDannrmtt i /ECRL The
ACKNOWLEDGMENTS AND DISCLOSURE OF FUNDING
This research was Funded by the European Union (ERC, Bayes-RL, 101041250). Views and opin- ions expressed 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 Euro- pean Union the granting authority can be held responsible for them.
REFERENCES
Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAl Pieter Abbeel, and Wojeiech Zaremba. Hindsight experience play. Advances in neural information processing systems, 30, 2017.
Dzmitry Bahdanau, Kyung Hyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In 3rd International Conference on Learning Representations, ICLR 2015, 2015.
Michael Chang, Alyssa Li Dayan, Franziska Meier, Thomas L. Griffiths, Sergey Levine, and Amy Zhang. Hierarchical abstraction for combinatorial generalization in object rearrangement. In The Eleventh International Conference Learning Reprexentations, 2023. URL https://
openreview net feeet//frrmmddd=ff6ooo
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. arXiv preprint arXiv:2306.055957 2023.
Danny Driess, Zhiao Huang, Yunzhu Li, Russ Tedrake, and Mare Toussaint. Learning multi-object dynamics with compositional neural radiance fields. In Conference Robot Learning, pp. 1755- 1768. PMLR, 2023.
Fan Feng and Sara Magliacane. Learning dynamic attribute-factored world models for efficient multi-object reinforcement learning. arXiv preprint arXiv:2307.092005 2023.
Stefano Ferraro, Pictro Mazzaglia, Tim Verbelen, and Bart Dhoedt. Focus: Object-centric world models for robotics manipulation. arXiv preprint arXiv:2307.02427, 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. IEBE, 2016.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor- critic methods. In International conference machine learning, pp. 1587-1596. PMLR, 2018.
Kevin Gmelin, Shikhar Bahl, Russell Mendonca, and Deepak Pathak. Efficient RL via disentangled environment and agent representations. In Proceedings of the 40th International Conference Machine Learning, Proceedings of Machine Learning Research, pp. 11525-11545. PMLR, 23-29 Jul 2023.
10

---

Published as a conference paper at ICLR 2024
Klaus Greff, Raphaël Lopez Kaufman, Rishabh Kabra, Nick Watters, Christopher Burgess, Daniel Zoran, Loic Matthey, Matthew Botvinick, and Alexander Lerchner. Multi-object representation learning with iterative variational inference. In International Conference on Machine Learning, pp. 2433-2433. PMLR, 2019.
Carlos Guestrin, Daphne Koller, Ronald Parr, and Shobha Venkataraman. Efficient solution algo- rithms for factored mdps. Journal of Artificial Intelligence Research, 19:399 46899446,,20 2003.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv: 2301.02300.0004 2023.
Negin Heravi, Ayzaan Wahid, Corey Lynch, Pete Florence, Travis Armstrong, Jonathan Tompson, Pierre Sermanet, Jeannette Bohg, and Debidatta Dwibedi. Visuomotor control in multi-object scenes using object-aware representations. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pp. 9515-95522 IEEE, 2023.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerehner. beta- Learning basic visual concepts with a constrained variational framework. In International Conference Learning Representations, 2017.
Max Jaderberg, Karen Simonyan, Andrew Zisserman, 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 Con- ference on Neural Information Processing Systems, pp. 4020-4031, 2018.
Jindong Jiang, Sepehr Janghorbani. Gerard de Melo, and Sungjin Ahn. Scalor: Generative world models with scalable object representations. In Proceedings of ICLR 2020. OpenReview.net, 2020.
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, and Sergey Levine. Scalable deep reinforcement learning for ussionbbssed robotic manipulation. In Aude Billard, Anca Dra- gan, Jan Peters, and Jun Morimoto (eds.). Proceedings of The 2nd Conference on Robot Learning, volume 87 of Proceedings of Machine Learning Research, pp.
PMLR, 29-31 Oct 2018 8.
Andrej Karpathy.
htt.ttua..coo/kapppttt//iiiiaaaa 2021.
Jannik Kossen, Karl Stelzner, Marcel Hussing, Claas Voelcker, and Kristian Kersting. Struc- tured object-aware physics prediction for video modeling and planning. arXiv preprint arXiv:1910.02425, 2019.
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuo- motor policies. The Journal of Machine Learning Research, 17(1):133-1373,, 2016.
Richard Li, Allan Jabri, Trevor Darrell, and Pulkit Agrawal. Towards practical multi-object manip- ulation using relational reinforcement learning. In 2020 ieee international conference on robotics and automation (icra), pp. 4051 1051- IEEE, 2020.
Baihan Lin, Djallel Bouneffouf, and Irina Rish. A survey on compositional generalization in appli- cations. 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 leamiigg with slot atten- tion. Advances in Neural Information Processing Systems, 33:1 1eetl1125-11133 2020.
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 gau-based physics simulation for robot learning, 2021.
11

---

Published as a conference paper at ICLR 2024
Davide Mambelli, Frederik Träuble, Stefan Bauer, Bernhard Schölkopf, and Francesco Locatello. Compositional multi-object reinforcement learning with linear relation networks. arXiv preprint MrXi1:2201.13388, 2022.
Russell Mendonca, Oleh Rybkin, Kostas Daniilidis, Danijar Hafner, and Deepak Pathak. Discover- ing and achieving goals via world models. Advances in Neural Information Processing Systems, 34:24379-2493 2021.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Joannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:
2013.
Aditya Mohan, Amy Zhang, and Marius Lindauer. Structure reinforcement learning: A survey and open problems. arXiv preprint arXiv:2306.602// 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.
Vitchyr H Pong, Murtaza Dalal, Steven Lin, Ashvin Nair, Shikhar Bahl, and Sergey Levine. Skew- fit: State-covering self-supervised reinforcement learning. arXiv preprint urXiv:1903.000 03699, 2019.
Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dor- mann. Stable-baselines:: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268):1 2021.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Cansu Sancaktar, Sebastian Blaes, and Georg Martius. Curious exploration via structured world models yields zero-shot object manipulation. Advances in Neural Information Processing Sys- tems, 35:24170-241 2022.
Alvaro Sonzhez-Gonzalez, Nicolas Heess, Jost Tobias Springenberg, Josh Merel, Martin Riedmiller, Raia Hadsell, and Peter Battaglia. Graph networks as learnable physics engines for inference and control. In International Conference on Machine Learning, pp. 4470-4479, PMLR, 2018.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 2017.
Gautam Singh, Sungjin Ahn, and Fei Deng. Illiterate dall-e learns to compose. In The Tenth Inter- national Conference on Learning Representations, ICLR2022. The International Conference on Learning Representations (ICLR), 2022.
Richard Sutton. Learning to predict by the methods of temporal differences. Machine learning. 3:9 44, 1988.
Manuel Traub, Sebastian Otte, Tobias Menge, Matthias Karlbauer, Jannik Thuemmel, and Martin V Butz. Learning what and where: Disentangling location and identity tracking without supervision. In The Eleventh International Conference on Learning Representations, 2022.
Ruben S van Bergen and Pablo Lanillos. Object- based active inference. In International Workshop on Active Inference, pp. 60-64. Springer, 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- fion processing systems, 30, 2017.
12

---

Published as a conference paper at ICLR 2024
Nicholas Watters, Loic Matthey, Matko Bosnjak, Christopher P Burgess, and Alexander Lerchner. Cobra: Data-efficient model-based rl through unsupervised object discovery and curiosity-driven 2019.
exploration. arXiv preprint
Tong Wu, Liang Pan, Junzhe Zhang, Tai Wang, Ziwei Liu, and Dahua Lin. Density-aware cham- fer distance as comprehensive metric for point cloud completion. In In Advances in Neural Information Processing Systems (NeurIPS), 2021, 2021.
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 Reprexentations, 2022.
Denis Yarats, Amy Zhang. Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Improv- ing sample efficiency in model-free reinforcement learning from images. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 10674-10681, 2021.
Jaesik Yoon, Yi-Fu Wu, Heechul Bac, and Sungjin Ahn. An investigation into pre-training object-
centric representations for reinforcement learning. arXiv preprint arXiv:2302.04449, 2023.
Andrii Zadaianchuk, Maximilian Seitzer, and Georg Martius. Self-supervised visual reinforcement learning with object-centric representations. arXiv preprint arXiv:2011.14381, 2020.
Andrii Zadaianchuk, Georg Martius, and Fanny Yang. Self-super reinforcement learning with independently controllable subgoals. In Conference on Robot Learning, pp. 384-394. PMLR, 2022.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander Smola. Deep sets. Advances in neural information processing systems, 30, 2017.
Linfeng Zhao, Lingzhi Kong, Robin Walters, and Lawson LS Wong. Toward compositional gener- alization in object-oriented world modeling. International Conference on Machine Learning. pp.
PMLR, 2022.
Allan Zhou, Vikash Kumar, Chelsea Finn, and Aravind Rajeswaran. Policy architectures for com- positional generalization in control. arXiv preprint arXiv:2203.05960, 2022.
13

---

Published as a conference paper at ICLR 2024
A EXTENDED BACKGROUND
A.1 DEEP Q-LEARNING
Goal-conditioned Deep Q-tearning approaches learn a Q-function Q" as, : S x A x E parameterized by deep neural network Qo (s, a, g) with parameters 0, which ap- proximates the expected return given goal when taking action a at state & and then follow- ing the policy и. Qo (s, a, g) is learned via minimization of the temporal difference (TD, SSt- ton (1988) objective given a state, action, next state, goal tuple "s.,as'.). CID(O) gr(s,g) + +.20 "(''ss 9) Qa (s.a,g) a, g)]² where rameters which are constant under the TD objective. Specifically in off-policy actor-critic al- gorithms (Fujimoto et al.) 2018). a policy network To (s, g) (actor) with parameters Φ is learned concurrently with the Q-function network (critic) with the objective of maximizing it with respect to the action: (Ø) = Qo (s, Th (s,g).g). Alternating between data collection via deploying The in the environment and updating Qe. with this data is expected to converge to an approximately optimal Q-function Qo ~ Q" and policy To 2 n².
a', g) is target network with pa-
A.2 DEEP LATENT PARTICLES (DLP)
DLP (Daniel & Tamar 2022) is a VAE-based unsupervised object-centric model for images. The key idea in DLP is that the latent space of the VAE is structured as a set of K particles 2 = [zy, xpl e Rxx(2), where 2 E k'l is a latent feature vector that encodes the visual appearance of each particle, and Zp E IR Kx2 encodes the position of each particle as (x, y) coordinates in Euclidean pixel-space. This requires several structural modifications to the standard VAE as described below.
The prior p(2[z) in DLP is conditioned on the image x, and has a different struc- ture for 2 and žp. p(=)(z)) p(=f) is modeled by set of standard zero-mean Gaussians. p(=plar) consists of Gaussians centered on a set of keypoint proposals which are produced by a CNN applied to individual patches of the image, followed by a spatial- (SSM, Jakab et al. 2018 Finn al. 2016).
Prior Modeling:
Encoder: DLP employs a CNN-based encoder that maps the image to set of means and log- variances for the keypoint positions 2p. The appearance features 2 are encoded from a region around each keypoint (termed glimpse) using Spatial Transformer Network (Jaderherg et al. 2015).
KL. Loss Term: As the posterior keypoints S1 and the prior keypoint proposals S2 are unordered sets of Gaussian distributions, the KL term for the position latents is replaced with the Chamfer-Ki KL: dCH-KI (S ( -,,,,, min es, K (=pll>() +L min K Lzpp l1
Decoder: Each particle decoded separately to reconstruct its glimpse RGBA patch. The glimpses are then composed with respect to their encoded positions to stitch the final image.
All components of the DLP model are learned end-to-end in an unsupervised fashion, by maximizing the ELBO (i.e., minimizing the reconstruction loss and the (Chamfer) KL-divergence between the posterior and prior distributions).
DLPv2: Daniel & Tamar (2023) expands upon the original DLP's definition of latent particle, as described above, by incorporating additional attributes. DLPv2 provides a disentangled latent space structured as set of K foreground particles z = ((pps Za, zd, 2t, 24)x 1 -1 E R x(6++)) 2p E R² and m, E R remain unchanged, 2, E IR2 is a scale attribute containing the (ar dimensions of the bounding-bor box around the particle, zu E R is pixel space "depth" attribute used to signify which particle is in front of the other in case there is an overlap and 2, E R is a transparency attribute. Moreover, it assigns single abstract particle for the background that is always located in the center of the image and described only by mag latent background visual features, 2ng ~ NN ogg) E Te We discard the background particle from the latent representation after pre-training the DLP for RL purposes. Training of DLPv2 is similar to the standard DLP with modifications to the encoding and decoding that take into account the finer control over inference and generation due to the additional attributes.
14

---

Published as a conference paper at ICLR 2024
A.3 THE ATTENTION MECHANISM
Attention (Bahdanau et al. 2015 Vaswani et al. 2017) denoted A(:,)) is an operator between two sets of vectors, X [Z5 i= and Y Uy ز=1 producing third set of vectors Z {zi J For simplicity, we describe the case were all input, output and intermediate vectors are in R Denote the key, query and value projection functions q(-), k(.), v(.)) Rd R= respectively. The attention operator is defined as A(X,Y) = Z where:
N
q (2x) (Us) dd
2j = Σ
a (T., (5) (5) (y5), a (1;, U5) = softmax,
E TR.
Namely, each element of the output set Z is a weighted average of the projected input set Y: {u(2,)) 5- where the attention weights Osj a (Iz Vj) express the "relevance" of yj for com- puting output 2, corresponding to Σ₁. An important property of A(X,Y) to permutations of X (permutation of elements in X results in the same permutation in the output elements in Z, with no change in individual elements' values) and invariant to permutations of Y (permutation of elements in Y does not change the output Z). In the special case where X YY, the operation is termed self-attention (SA), and otherwise, cross-attention (CA).
that it is equivariant
B CHAMFER REWARD
We desire reward that captures the task of moving objects to goal configurations, However, because particles in different images are not aligned, and some particles may be occluded or missing, we cannot directly construct a reward based on distances between the particles. Instead, we define a reward from the DLP representations of images as the Generalized Density-Aware Chamfer (GDAC) distance between state and goal particles, which we term Chamfer reward. The GDAC distance is defined between two sets X {x:) Y = {/
Izi Vi E R2 in the following manner:
Distanse (X,Y)= E,(({, 0) Σ Tix/te (1r E1(E> 1II+E =xx,
1
Σ EDI Di
BEY,
where Xj faril arg min EY (D2 (a,, 9x)) - jj.. Y, {y5l arg ming,ex (D2 (v, xx)) Dj (x, y) and D2(x, y) are two distance functions between entities. The standard Chamfer distance obtained by setting Di (1( y) D2(x, v) || - =112 and substituting Σ, (([x)>>0) [Xy++e*
E. 70))))) Mil+e with 1 1 | respectively, removing the inner sums in both terms. |
The Chamfer distance measures the average distance between each entity in X and the closest en- tity to in Y and vice versa. The Density-Aware Chamfer distance (Wu et al., (202T) takes into account the fact that multiple entities from one set can be mapped to the same entity in the other set, and re-weights their contribution to the overall distance accordingly. The Generalized Density- Aware Chamfer distance, decouples the distance function that is used to match between entities D2(x, y) and the one used to calculate the distance between them D (z. y). Decoupling these two allows using entity-identifying attributes for matching while calculating the actual distances between matching entities based on localization features. For example, we can use the DLP visual features
2, match between objects in the current and goal images, and then measure their distance using the (x, y) coordinate attributes Zp-
B. FOCUSED CHAMFER REWARD
In many robotic object manipulation settings, we do not care about the robot in our goal specification as long the objects reach the desired configuration. In order to consider only subset of the entities for the image-based reward (e.g. particles corresponding to objects and not the agent), we train simple multi-layer perceptron (MLP) binary classifier on the latent visual features of the DLP representation, differentiating objects of interest from the rest of the particles. We train this classifier annotated particles extracted from 20 images of the environment. Annotation required 5 minutes of our time and training the classifier itself took just a few seconds. We then filter out particles based on the classifier output before inputting them the Chamfer reward. We emphasize
15

---

Published as a conference paper at ICLR 2024
100
100
75
Success Rute
75
Succesn Bute hate
75
50
25
25
Durs SMORL
25
Durs SMORL
Ours
0
100000 200000 30000C 403000 Enwiromment Teeeteps
0.0
1,0
2.0
3.0
Environment Timesteps
Enviromment Timesteps
(a) 1-Cube
(b) 2-Cubes
(c) 3-Cubes
Figure 8: Success Rate vs. Environment Timesteps (Image-Based Rewards) Values calculated
based on 96 randomly sampled goals.
that this supervision is only required for training the classifier which is used for the image-based reward exclusively during RL training, and is very simple to acquire both in simulation and the real world.
B.2 TRAINING WITH THE CHAMFER REWARD
We compare our method SMORL, trained entirely from images. Results on N-Cubes for N {1. 3.3 are presented in Figure 8 and Table 2 obtains lower success rates compared to training with the GT reward. While this is expected, we believe the large differences are due to noise originated in the DLP representation and occlusions, which make the reward signal less consistent and harder to learn from. This especially hard with increasing number of objects, as the chances of at least one object being occluded are very high. This is highlighted by the drop in performance from 1 to 2 cubes, compared to the GT reward. Image- based reward calculation for single object manipulation, as in SMORL, is slightly more consistent as occlusions in a single view will not affect the overall reward as much. Adding more viewpoints for the reward calculation might improve these results without increasing inference complexity. We see that in the 3- CuCbes environment, our method surpasses SMORL, although SMORL's reward is based on single object regardless of the number of objects in the environment. This could be as a result of object-object interactions being more significant in this case.
Training our method with the image-based reward
Method
Success Rate
Success Fraction
Max Obj Dist
Avg Obj Dist
Avg Return
Ours
0.765 + 0.025
0.875 + 0.015
0.037 = 0.002
0.026 + 0.001
-0,210 = 0.009
SMORL
0.838 + 0.016
0,911 + 0.008
0.038 + 0.004
0.025 + 0.002
-0.320 = 0.007
Ours
0.580 + 0.093
0.822 + 0.052
0.063 + 0.008
0.035 + 0.005
-0.251
0.022
SMORL
0.509 + 0.044
0.794 + 0.024
0.092 + 0.006
0.047 + 0.004
-0.451 = 0.031
Table 2: Performance Metries: Image-Based Rewards Methods trained and evaluated on the
2-Cubes (top) and 3-Cubes (bottom) environments. Values calculated on 400 random goals per random seed.
C ADDITIONAL RESULTS
C.I MULTI-OBIECT MANIPULATION
Performance metrics for the 2-Cubes and 3-Cubes are presented in Table
Method
Success Rate
Success Fraction
Max Obj Dist
Avg Obj Dist
Avg Return
Ours (State)
0.991 + 0.004
0.995 + 0.003
0.014 0.001
0.010 = 0.001
-0.129 + 0.006
SMORL (State)
0.980 1 0.006
0.989 i 0.005
0.014 0.002
0.009 + 0.002
-0.142 0.016
Ours (Image)
0.968 1 0.019
0.983 1 0.009
0.020
0.015 1 0.001
-0.150 + 0.008
Ours (State)
0.978 1 0.006
0.991 + 0.002
0.016:00000
0.010 0.001
-0.124 0.007
SMORL (State)
0.932 1 0.022
0.974 1 0.009
0.028 1 0.005
0.015 1 0.002
-0,201 0.011
Ours (Image)
0.919 + 0.008
0.969 + 0.004
0.026 + 0.002
0.016 + 0.001
-0.157 + 0.007
Table 3: Performance Metrics: GT Reward Methods trained and evaluated on the 2 2-Cubes (top)
and 3 -Cubes (bottom) environments. Values calculated on 400 random goals per random seed.
16

---

Published as a conference paper at ICLR 2024
C.2 GENERALIZATION
Performance metrics for the compositional generalization to different numbers of objects of an agent trained the 3- Cubes environment are presented in Table Additionally, we compare compositional generalization performance with respect to the number of objects seen during training in Table[5] We see that when learning with 3 objects our agent able to generalize reasonably well to larger amount of objects. This is not the case with agents trained on 1 or 2 objects, where there is a sharp decay in performance starting from a single additional object. When training on 1 object, the policy lacks the need to perform reasoning between multiple objects in the state, thus it not surprising it does not generalize to more than a single object. While training on objects does require this type of reasoning, we believe training on 3 objects has such an increase in generalization abilities because the agent encounters more scenarios during training where modeling object interaction and interference is necessary.
Number of Objects:
Number of Cubes
Success Rate
Success Fraction
Max Obj Dist
Avg Obj Dist
Avg Return
0.973
0,973
0.016
0.016
-0. 162
2
0.963
0,981
0.023
0.017
-0.154
0,838
0.942
0.034
0.02
-0.175
4
0.723
0,912
0.051
0.027
-0.213
5
0.57
0,876
0.068
0.031
-0.245
6
0.398
0.826
0.09
0.036
-0.294
Table 4: Performance Metrics for Different Number of Cubes than in Training Our method's performance on different numbers of cubes in the -ccubes environment, trained on the 3 3-cubes environment (results in bold) with cubes of 6 different colors. Values are averaged over 400 episodes with randomly initialized goal and initial configurations.
Cubes in Test / Train
3
2
0.016
0.013
0.017
0.023
0.024
0.256
3
0.034
0.091
0.287
0.051
0.149
0.311
5
0.068
0.276
0.320
6
0.09
0.292
0.328
Table Maximum Object Distance Comparison for Different Number of Cubes than Training different numbers of cubes in the N-cubes environment. We compare agents trained on the 1, 2, 2-3bbe environments with cubes of 6 different colors. Values averaged over 400 episodes with randomly initialized goal and initial configurations.
Our method's performance
Distracters: An additional scenario we consider providing the agent a goal image which contains some cuhe colors that agent is able to disregard the distracters in the goal image while successfully manipulating the other cubes their goal locations. demonstration of these capabilities are available on our website.
present in the environment. We term these cubes distracters. The
Object Properties: While we designed our algorithm to facilitate compositional generalization, interesting to study its generalization to different object properties. Dealing with novel objects would require generalization from both the DLP and the EIT. We would expect our method zero- shot generalize to novel objects in case: (1) They are visually similar to objects seen during training. (2) Their physical dynamies are similar to the objects seen during training. To test our hypothesis, we deploy our trained agent in environments including the following modifications: (I) Cuboids obtained by enlarging either the x dimension or both x and y dimensions of the cube. (II) Star shaped objects with the same effective radius of the cubes seen during training. (III) Cubes with different masses than in training.
(V) Cubes colors not seen in RL training DLP pre-training. Performance metrics for these cases presented in Table/6] Visualizations of the modified object environments and how the DLP model perceives them available in Figure)
(IV) Cubes in colors not seen in RL training.
17

---

Published as a conference paper at ICLR 2024
Change in Shape Based on study, the DLP model is able to capture the different shaped objects and their location. Observing the reconstruction of the scene, DLP models the objects using the building blocks it knows, which are cubes. Stars mapped to cubes and cuboids are mapped to a composition of multiple cubes. While this is an interesting form of generalization for image reconstruction, it is not sufficient for inferring changes in dynamics. In cases where the shape does not strongly affect dynamics such as the star and the slightly modified cuboid, the EIT agent still achieves strong performance. When this is not the case there is a significant performance drop, as expected.
Change in Color With a similar study of the DLP reconstruction we witnessed an interesting yet not surprising phenomenon: colors not seen in DLP pretraining are mapped to the closest known color in the latent space. For example brown mapped to green and orange to yellow, pink to purple. Judging by the non-negligible success rates we hypothesize that EIT generalizes to both control and matching of colors it has not seen in RL training via the DLP latent space. We believe the reason for the drop in performance originates in ambiguity caused by inconsistent mapping of colors. One failure scenario is when the goal object is mapped to a different color than the corresponding state object, due to differences originating in shading or other factors. Another failure scenario is two different objects being mapped to the same color, causing ambiguity in goal specification which the agent is not expected to generalize to. The above could provide an explanation the fact that our agent performs better with colors not seen in both DLP and RL training than with colors only seen by DLP: the performance more strongly depends on the combination of colors than on the ability of DLP to recognize each color individually.
Change in Mass Changes mass have resulted in the smallest performance drop out of all the scenarios we considered. While the change in mass has significant affect on dynamics, it does not affect dynamics related to torques (moment of inertia matrix is of the same structure) or agent-object contact. Additionally it does not affect the appearance of objects and therefore does not require generalization from the DLP. We believe our EIT policy is able to generalize well to these changes because it is Markovian and therefore reactive: it does not need to predict the exact displacement of the object in order to infer the direction of its action and can simply react to the unraveling sequence of states.
Our main conclusion from this study is that while zero-shot generalization to objects with different properties is only partial, the non-negligible success rates hint at potential few-shot generalization.
Configuration
Success Rate
Success Fraction
Max Obj Dist
Avg Obj Dist
Avg Return
Training
0.919 + 0.008
0.969 + 0.004
0.026 + 0.002
0.016 = 0.001
-0.157 + 0.007
Star
0.902 + 0.002
0.961 + 0.002
0.031 + 0.005
0.019 = 0.001
-0.167 1: 0.007
Cuboid 1.5
0.743 + 0.061
0.891 + 0.030
0.052 + 0.008
0.029 + 0.004
-0.228 1: 0.023
Cuboid
0.403 0.079
0.694 0.056
0.124 + 0.020
0.063 = 0.010
-0.408 + 0.043
Cuboid x, 1.5
0.316 + 0.077
0.635 = 0.060
0.205 + 0.037
0.097 = 0.017
-0.560 0.078
Cuboid x, = 2
0.026 + 0.008
0.274 + 0.040
0.398 + 0.009
0.217 = 0.014
- -1.089 + 0.048
Colors New to EIT
0.379 + 0.052
0.728 + 0.031
0.094 + 0.021
0.047 = 0.009
-0300 + 0.041
Colors New to EIT + DLP
0.550 + 0.100
0.810 + 0.050
0.089 = 0.009
0.042 = 0.005
-03007 + 0.026
Mass = 0.05
0.900 + 0.004
0.961 + 0.001
0.033 + 0.002
0.019 + 0.001
-0.161 4 0.005
Mass = -0.1
0.888 + 0.023
0.956 - 0.008
0.035 + 0.003
0.020 = 0.001
-0.169 + 0.007
Mass 10
0.881 - 0.014
0.945 - 0.008
0.032 = 0.002
0.021 = 0.001
-0.244
+ 0.009
Mass -20
0.751 0.019
0.876 0.010
0.061 + 0.009
0.035 + 0.004
-0.378 0.016
Table 6: Performance Metries: Zero-shot Generalization to Object Properties Methods trained
from images with GT reward and evaluated on the 3 3-Cubes environment with red, green and blue cubes. Values calculated on 400 random goals per random seed.
18

---

Published as a conference paper at ICLR 2024
Original Image
Particles
FG Reconstruction
Shape
Shape
Color
Shape
Color
Figure 9: DLP perception of environments with shape and color modifications of cube objects not
seen during training. Left to Right: Raw Image Visualization of Particle Locations Foreground
Reconstruction of the DLP Decoder. Top to Bottom: Cuboid x 2 Cuboid x, -- = 2 New Colors Star Shupe & New Colors.
C.3
PUSH-T
See Figure 10 for visualization and performance results on the Push-T task.
t=3
Gosl
Figure 10: Left Rollout of an agent trained on the Push-T task. Right Distribution of object angle difference (radians) from goal. Values of 400 episodes with randomly initialized goal and initial configurations.
19

---

Published as a conference paper at ICLR 2024
100
100
Charier
80
Tamnnn Inpit Type
Sing viee
Success
Success
40
Image Charteer Reword
Mescco
Image GT Roward
20
Ficur
0.0
0.5
1.0
1.5 Timesteps
2.0
0.0
0.5
1.0
1.5
2.0
Environment
Enviranment Timesteps
(a) Multiview
(b) Action Entity
Figure 11: Success Rate vs. Environment Timesteps Multiview and Action Entity Ablation
Values calculated based on 96 randomly sampled goals.
C.4 ABLATION STUDY
We explore how aspects found be key to the success of proposed method effect sample efficiency. Figure Taa compares our method with multi-view vs. single-view image inputs. We find that both with GT and image-based reward, multi-view inputs substantially improve sample efficiency. We believe that the connections formed between particles from different views in the Transformer blocks makes it easier for the agent learn the correlations between actions defined in 3D space and latent attributes defined in 2D pixel space. In addition, multiple viewpoints decrease the degree of partial observability. Figure TTb compares treating the action a separate input entity to the Q-function Transformer blocks vs. concatenating the action to the output of the final Transformer block, before the output MLP. We find that learning the relations between the action and the state and goal entities via the attention mechanism is crucial to the performance of our method. Without it, our experiments exhibit decreased sample efficiency in state observations and failure to learn in the given environment timestep budget with image observations.
Figure 12 presents an ablation of the contribution DLP's attributes, z = (Zps Za. Zd. 21, 24), to the agent's success the 2 2-Cubes environment. The position attribute 2p and visual features 2, contain necessary location and entityidentifying information, therefore we do not run experiments without them. The results show equivalent performance when discarding the depth za and trans- parency 2, attributes as well as without attention masking based on the transparency. The seale attribute 2s. the other hand, proves be significant for sample efficiency although does affect final performance.
100
Success Rate
75
No Transparency
No Attention Masking
No Depth
25
No Scale Ours
0
0.0
0.5
1.0
1.5
2.0
2.5
3,0
Environment Timesteps
Figure 12: Success Rate Environment Timesteps - DLP Attribute Ablation Values calculated based on 96 randomly sampled goals.
Figure 13 compares the performance of our method with DLP vs. Slot-Attention (SA, Locatello
et al. 2020) as the OCR. On the 1-Cube environment the performance is equivalent (sec [3a]. Note that the (x. y) coordinates not explicit in the SA latent representation. Nevertheless, the
20

---

Published as a conference paper at ICLR 2024
100
0
75
Success
50
Return
-50
-75
25
Slot Attention
Slot Attention
DLP
-100
DLP
0
100000 200000 300000 400000 Environment Timesteps
0.0
0.5
1.0
1.5
2.0
Environment Timesteps
(a) 1 1-Cube Success Rate
(b) 2 Cubes Return
Figure 13: OCR Performance Ablation Values calculated based on 96 randomly sampled goals.
EIT is able to infer object location from the slots. These results showcase similar capabilities to the ones presented in the Push-T task, where the EIT was able to infer orientation from the DLP latent visual attributes. In the 2-Cubes environment, despite observing a moderate increase in return during training (see Figure [3b), our method with SA was unable to solve the task, achieving approximately 0% success rates. From an investigation of these experiments and the representations produced by SA, we found that often, both cubes were assigned to a single slot. This led to the agent learning to push both cubes to the middle point between the two cubes' goals. We hypothesize that this behavior is optimal given that the agent can only infer a single location for both objects in the same slot. We slots. Our design choice of utilizing DIP, coupled with training a model with a relatively large number particles (24) and limiting the capacity of each particle's visual latent features (zf 12²), effectively prevented multiple cubes from being assigned to a single particle.
able to train SA model which consistently separated the cubes different
D IMPLEMENTATION DETAILS AND HYPER-PARAMETE
In this section, we provide extensive implementation details in addition to the open-source code that can be found in the official repository: https: //tpt//ttttt//// com/ Hat.rrti/HHERrt
D.1 ENVIRONMENT
We implement environments with IsaacsGym (Makoviychuk al. 2021). by adapting code from IsaacGymEnvs and OSCAR"
Ground-truth State Denote s; (z,, y:). 9: (x.. vi" the xy coordinates of the state and goal of entity respectively. The input to the networks in the structured methods are two sets of vectors {us) = [s; one-hot [i]])] E 22+ { = [g,, one-hot [i] E R2+N [-] denoting concatenation, where the one-hot vectors serve entity-identifynng features. In the unstructured case, the input is [s), $2. N 91, 92, 9N..
Ground-truth Reward The reward calculated from the ground-truth state the system, which we refer to as the ground-truth (GT) reward, is the mean negative L, distance between each cube and its desired goal position on the table:
1 F4 N
10
9 1/
where of and of denote the desired and achieved goal for object i respectively, the number of objects, the immediate reward at timestep & and L a normalization constant for the reward corresponding to the dimensions of the table.
htteess//gtthbb.cco///////aaaaa. htt.:omwwwieeeeeeeeeaeooooooo http::///gtthbb..cm/////bbbb/
saacGymenvs
21

---

Published as a conference paper at ICLR 2024
Image-Based Reward The reward calculated from the DLP OCR for method is the negative GDAC distance (see Eq. between state and goal sets of particles, averaged over viewpoints:
K
DistaDAc (pmmm tam fm)m=))
(3)
K
1=
2| 2/ and D2 (x. v) |, 2/ in the GDAC distance, 2$²
where D1 (Ir, v)
2f denoting DLP latent attribute 2p. 2f of particle (-) respectively. We filter out particles that do not correspond to cubes (see section B.1 for the distance calculation. When particle has no match C), a negative bonus is added to the reward to avoid "reward hacking" by pushing blocks off the table or occluding them intentionally.
(i.e. miny 2
Evaluation Metries We evaluate the performance of the different methods based on the following:
Success: T TE T[loo - 2اl12 < R) =N) all N objects are at a threshold distance from their desired goal. R denotes the success threshold distance and is slightly smaller than the effective radius of cube. I denotes the indicator function. This metric most closely captures task success, but does not capture intermediate success or timestep efficiency.
Success Fraction: t ΣN (1la - of 1/2 < R), fraction of objects that reach individual success.
Maximum Object Distance: max; {1lof ot 21 largest distance of an object from its desired goal.
1 Σ | of Il average distance of objects from their desired goal. Average Object Distance: N
1= rt. the immediate GT reward averaged across timesteps, where T is the evaluation episode length. A high average return means that the agent solved the task quickly.
Average Return: 3 "E=
D.2 REINFORCEMENT LEARNING
ines3 (Raffin al., 2021]. Specifically, we use TD3 (Fujimoto et al.) 2018) with HER (Andrychowicz et al. (017). We e-greedy and Gaussian action noise for exploration, that decays half its initial value with training progress, similar to Zhou et al. (2022). We use Adam for neural network optimization. Related
We implement RL algorithm with code adapted from
can be found in Table and Table
Learning Rate
5e-4
Batch Size
512
?
0.98
0.05
Episodes Collected per Training Loop
16
Update-to-Daan Ratio
HER Ratio
0.8
Exploration Action Noise a
0,2
Exploration
0.3
Table 7: General hyper-paramette used for RL training.
Number of Cubes
2
3
Episode Horizon
30
50
100
Replay Buffer Size
100000
100000
200000
Table 8: Environment specific hyper-parameters used for RI. training.
For the Entity Interaction Transformer (EIT) we adapted components from DDLP's Particle Interac- tion Transformer (PINT,Daanie & Tamar (2023)), which is based Transformer decoder architec- ture and utilizing the open-source minGPT (Karpathy) 202T) code base. Related hyper-parameters can be found in Table
22

---

Published as a conference paper at ICLR 2024
Attention Dimension
64
Attention Heads
8
MLP Hidden Dimension
256
ML.P Number of Layers
Table 9:
for the EIT architecture.
Attention Masking: The DLP model extracts a fixed number of particles, which often include particles that do not represent objects in the image. These particles assigned low transparency (2t) values by the DLP model as to not affect the reconstruction quality. We disregard these particles in our policy and Q-function by directly masking the attention entries related them. We found that this slightly improves sample efficiency but is not crucial to performance as the ETT is able to learn to disregard these particles by assigning them very low attention values.
Policies for the unstructured baselines have 5 layer MLPs with hidden dimension 256.
D.3 PRE-TRAINED IMAGE REPRESENTATIONS
In this section, we detail the various unsupervised pre-trained image representation methods used in this work. We begin with the non-object-centric baselines, i.e., methods that given an image R encode single-vector representation = E IR.,, where D is the latent dimension, of the entire input image. Then, we describe the object-centric representation (OCR) method that provides a structured latent representation = E R xx of given image 1, where K the number of entities in the scene, each described by latent features of dimension d.
Data: We collect 600, 000 images from viewpoints by interacting with the environment using a random policy for 300, 000 timesteps. For all methods, we use RGB images at a resolution of 128 128, i.e., I E IR 128x 128x3
Variational Autoencoder (VAE): We train a B-VAE (Higinns et al. (017)) with latent bottleneck of size D 256, i.e. each image is encoded as 2 E 225 We adopt a similar autoencoder archi- tecture as Rombach et al. (2022) based on the open-source implementation and add 2-layer MLP with 512 hidden units after the encoder and before the decoder to ensure the latent representation is of dimension 256. We use B = le 10, batch size of 16 and an initial learning rate of 2e 4 which is gradually decayed with a linear schedule. The model is trained for 40 epochs with a per- ceptual reconstruction loss and L1 pixel-wise loss, similarly to Rombach et al. (2022). and we keep the default values for the rest of the hyper-parameter
Deep Latent Particles (DLP): We train DLPv2 (Daniel & Tamar 2023) using the publicly avail- able code base as our unsupervised OCR model. We modify the DLP model to have background particle features of dimension 1, and discard the background particle for RL purposes. The background is static in our experiments and setting its latent particle to have a single feature is meant limit its capacity to capture changing parts of the scene such the objects the agent. Recall that DLP provides disentangled latent space structured as a set of foreground particles
z = (Ep: Ze, 2d. Zt, 24) 2 K-1 E (6++1) where K' is the number of particles. Figure 14 il- )/
lustrates an example of the object-centric decomposition for single image using a DLP model pre-trained on our data. We keep the default recommended paree-pprameters and report the data- specific papammerreettrrs Table TO Note that our DLP model represents an image by a total of K (6 +1)= = 24 (6 + 4) = 240 latent features.
Slot-Attention: For the OCR ablation study we train a Slot-Attentin (Locatello al. 2020) model with 10 slots, each of size D = 64, i.e., each image T is encoded as 2 RI We use the implementation from and keep most of the hyper-parameters similar, corresponding to the ones used in the original paper, and we provide the set of tlpper-paaameten our code repository. We performed multiple training runs with each set of paraamearseseess and took the run that yielded the best object separation to slots, similar to the training procedure in Wu et al. 2022).
com/ HHwusee/eeeeeeeoeee....e. ob ject-di siccttddssovvery pecovery-pytorree
htttes///tthbb..oo///////////sss -ifuniio httpss//giittbbb.com//aaaaaatt
23

---

Published as a conference paper at ICLR 2024
Original Inage
Reconstruction
Bounding Boxes
Foreground
Background
Porticles
Figure 14: Object-centric Decomposition with DLP - DLP decomposes a single image into la- tent particles, each characterized by attributes including position (keypoints the images), scale (bounding boxes), and visual appearance features around the keypoint (displayed as decoded glimpses from these features).
Batch Size
64
Posterior KP K
24
Prior KP Proposals L.
32
Reconstruction Loss
MSE
3KL
0.1
Prior Patch Size
16
Glimpse Size S
32
Feature Dim m
4
Background Feature Dim mag
1
Epochs
60
Table 10: Hyper-parameters used for the Deep Latent Particles (DLP) object-centric model.
D.4 SMORL REIMPLEMENTATION
We re-implement SMORL (Zadaianchuk al. 2020) based the official implementationf? using the same code-base as we used for the EIT for the SMORL attention architecture. SMORL specific hyper-parameters are detailed in Table []] We extend SMORL multiple views, which includes modifications to several aspects of the algorithm:
Attention Architecture: We extend SMORL's attention policy by adding a goal-conditinned and goal-unconditional attention block for the additional view. The outputs of attention layers from both views are concatenated and fed an MLP, as the single-view version. An outline of SMORL's single-view attention-basee architecture is described in Figure 15
https::/ggithbb.cmm// 1tt
24

---

Published as a conference paper at ICLR 2024
GUA
unconditooned
(BUA)
Query
[+]
Attentoon
Multi-Head
Stone Particles
Concat
MIP
GCA
Partides
Multi-Hed Attention
Action
Choose
Goal
Goal conditioned
Attaniion (GCA)
Figure 15: Outline of SMORL's Attention Architecture The policy is conditioned the goal by choosing a single goal particle and feeding it to gross-attention block between the goal particle and the state particles. In parallel, cross-attention between learned particle and the state particles is performed to extract features from the state that are not goal-dependant. The outputs of the two attention layers are then concatenated the original goal particle and fed an MLP to produce the action. For the Q-function, the input action is additionally concatenated to the output of the attention to produce the value.
Selecting Single Goal: SMORL decomposes the multi-object goal-conditioned task to single objects by selecting a single goal at a time, and rewarding the agent with respect to this sub-goal alone. Working with multiple views requires selecting goal particle corresponding object from both viewpoints, which requires explicit matching. We do this by selecting a goal particle from one viewpoint and choosing the closest matching particle from the second viewpoint based on the Lg distance in latent attribute z
the same
The image-based reward calculated from the DLP OCR for SMORL is the negative L2 distance in attribute 2p between the goal particle to the closest matching state particle based on attribute zy, averaged over viewpoints:
Reward:
K
T1
K
P
2p
my arg min
(4)
ge denoting the goal particle from view k and stt denoting particle m from view k. When there is
no match for the goal particle in viewpoint & (i.e. min, 10 -
C), the minimal reward
given for that view. Note that this reward a special case of the Chamfer reward we define this work, with the goal set consisting of single entity per view.
Attention Dimension
64
Unconditional Attention Heads
8
Goall-conditioned Attention Heads
8
MLP Hidden Dimension
256
MLP Layers
4
Scripted Meta-policy Steps
15
Table 11: SMORL hyper-parameters.
25

---

Published as a conference paper at ICLR 2024
E ATTENTION IN RL POLICIES COMPARISON TO PREVIOUS WORK
In this section, we compare our use of attention to two previous approaches, Zhou et al. (2022)
which is state-based and SMORL (Zadaianchuk al.) 2020) which is image-based.
Zhou et al. (2022) also propose a Transformer-basee policy. They define an entity as the concate- nation of each object's state, goal, and the state of the agent (and the action when we consider the input to the Q-function). This requires explicitly matching between entities in state and goal as well as identifying the agent, which trivial when working with GT state observations. When learning from OCRs of images, this is not at all trivial. Matching is not always possible due to lack of one-to-one match occlusion, which also limits the use of multiple viewpoints. We tackle this in our EIT by using cross-attention block for goal-conditioning Additionally, we learn which parti- cles correspond to the agent implicitly through the RL objective. Figure 16 describes the differences between their definition of an input entity to ours.
SMORL uses an OCR images and does not require matching between entities the single-view case. The goal-conditioned attention in their architecture (see Figure [5) matches single goal par- ticle to the relevant state particle via cross-attention. Different from us, the attention mechanism is used only extract sub-goal specific entities from the set of state entities, and does not explicitly model relationships between the different entities in the state. We do this explicitly by incorporating self-attention Transformer blocks our architecture.
Goal
Action
Irput Telen
Iegut Teken
Input 1akan
INGT
Ineal
Input Toben
X0N-11
Attention based
Iouut Tolen
Attent on based Network
Irput Taken
JMSt
12g
Irput Tolen
[Goa
Ingut Toben
[GaM
Ipptt [Goal
Figure 16: Entity Definition Comparison Left descripiinn of an input token defined by Zhou et al. (202), where each token is concatenation of the object and corresponding goal as well as global entities such as the agent and action; Right Our definition of input tokens, where each entity is treated separate token.
26

---

Published as a conference paper at ICLR 2024
F COMPOSITIONAL GENERALIZATION THEORY
We begin by defining a notion of compositionally generalizing functions F.1.1 and provide an ex- ample of such class of functions F.1.2 We additionally provide example of case where class of functions can accurately approximate another class of functions up to N entities but does not generalize increasing number of entities Following these, we prove our main theorem F22 We then present modifications to the assumptions of our main theorem to obtain result F3 that adheres Definition 3
Relating the theorems to compositionally generalizing policies, we show that an e-optimal Q- function for M objects is also approximately optimal for M objects under our structural as- sumptions. This implies that a policy that is trained to maximize this Q-function with M objects should also achieve high returns for M + objects, i.e. achieve zero-shot compositional generaliza- tion.
F.1 COMPOSITIONALDY GENERALIZING FUNCTIONS
F.1.1 DEFINITION
In the following, we define notion of compositionally generalizing functions:
Definition Denote the space variable-sized sets of entities S UNENSN where S.V {[s)),lss 5}. S being the space of a single entity. Denote SEN C S the subspace con- taining sets of up size N.
A class of functions : + RR is said to admit compositional generalization on S S if there exists N such that for any f C that satisfy (f f(z)] <E, Va E XEN we have that for any M > 0 and x E P M: [lf((z) - f(=)l (C1 + C2 M) -E, Vx E X' where C1, C2 constants that do not depend M,
f
Put simply, our definition of compositional generalization asks that adding additional M objects to the problem incurs an error that is at most linear M.
F.1.2 EXAMPLE
An example of a class of compositionally generalizing functions is the class of DeepSets-style (Za- heer et al., 2017) function approximators, which are an aggregation of functions defined on single entities:
Theorem 4. Let CDs S - R. be the class of functions of the form: Q(s) ト/4 D. (s:).
where v S R is a function that operates on single entities. Then CDS admits compositional generalization on S.
Proof. Assume Q" Q E CDS satisfy (xQ* QQ))l <E, Vs EEE
From the case of N = 1 we have:
(s) ((s)) < [|0 0(s)ll <e, S.
From the above equation we obtain VN e N:
N Σ
[+0(()))(()1 = Σ
(s:)
Σ 100 (s))-
Σ = Vs E SN.
N
Thus we have that by Definition 3 CDS admits compositional generalization on S with C1 1, C2 0.
F.1.3 NON-GENERALIZING Q EXAMPLE
Assuming Q" has self-attention structure, Theorem shows that if we obtained an E-optimal Q for 1,..., M objects where Q also has self-attention structure, then the sub-optimality w.r.t. M k objects grows at most linearly in M k. This raises the question: there Q structures that lack the
27

---

Published as a conference paper at ICLR 2024
compositional generalization quality in this case? In this section provide simple example for structure that can accurately approximate a self-attention Q" for 2 objects but does not generalize to increasing number of objects.
Consider the following Q structure:
AA E SN, Va E YN E N: (1,,,..,,,00 + Q (51, ..., sN, where ()) RR.
Qi (81, 0N Ex a(81.85.0)
8y, 0) Σj=1 a 85,, We denote N min (2, N where the indices j ,,...,N NN are ordered by the value of (sj. a) in increasing order. Note that this operation still invariant to permutations of the input state
{8i {8
For N = 2, this structure is identical to the one assumed in Theorem 2 which is that of Q" and therefore there exists a Q such that: IQ Q" = 0. In this case, both u = u* and a = a².
For N > 2 on the other hand, the approximation error cannot be bounded by E as Q does not account for the entire set of states:
N
Q ($1 SN, SN, a) Q" SN,
Σ SN, - Q (s.....sN.a))=
N
2
N
Σ j=1
a (8 ())
Σ
a
(si, S5, a) a)" (s),
N
j=1 a (8t, 83. a)
Σ a Sj.
j=1
a 83)
1 (8
a (5,,
(8).0)
Σ
a (8% Sg.
2- (8),a)
- a 88,)
N
Σラ= a* (8,
N
N
a" (8 8g.
i=1 j=1
j=3
S5) a)
1" (8).0)
20
Σイ= a" (si, a)
The approximation error not zero for all inputs ($1...., SN. a) so long as v* > 0 for some state 8; E S and action a, which is true unless Q" (s1, SN a) = 0, E Va E A.
This example illustrates how some function approximators generalization. Although the approximated model perfectly fits the training distribution, containing up to objects in this case, it does not fit test distributions with more objects.
not well-suited for compositional
Relating this Definition if we consider class of functions that contains both Q and Q* de- scribed above, the example llustrates that this class does not admit compositional generalization.
F.2 MAIN THEOREM PROOF
F.2.1
LEMMAS
We start by showing that two positive functions < f (s) , (s) normalized by sum of their values over set of inputs { of size M are 8-close to each other, the difference in their normalized value under a set size 2M is bounded by 28.
f(x)) EE [(,) k E [1, M - - [1, M]:
Lemma 5.
g(v))
< < < f(8) Vsi, 8j, VN E [1, then for
EN g(e,) )(e,)
((s)
g (s:)
$ 2I
Σ [(3))
Σ g (s,)
28

---

Published as a conference paper at ICLR 2024
Proof.
(8i)
9 (sp)
((i) Ej=1 (13) - (8:) Ej= S(85)
(+)
Σ==1 (s))
Ej=1 g (83)
(E (8))) (E)=1 15=19(1) ((3))
1f (9 (si) f ((,)]
<8
E (I)=
For k E [1, M- - 1) E [];
f f(a)
(s.)
Σ Σ (s:) g))))((())))) ( [s/)]
E
M-R f())
Σ 1==1 (1²))
M+R
E I (5,) (E 9 (8,))
Σ Ditij= If (8:) g (8)) - 9 (s,) f (s,)
M
Σ Σj=m+1 If (8:) g (
- g(;) f(;)]
<
E [8,)) [ 9 (87)
Σイ= {(,) [E==] ((,)))
ocf(sha(M)
DiAy= If (8:) (()) (s:)
M
Σ M+k
M+ If (s:) (83)
[s,]]
5 26 (0)
Eا f (8))) (
(f + 2==M++ {
M+*
(g (8:) + E/= 9 (8))
I (s))
g (si)
26
M+K (8,)
ΣM* g (8,)
In the following two lemmas we bound the difference between two weighted sums of single-inpur functions applied individually on a set of inputs {s. assuming the weights are ô-close and the function values are e-close.
Lemma 6. or, B,v,u 0 and la B1 < 8, 10
ul E, then:
Jav - Bul < a
V ++
8 +
2
Proof.
Jare - Bu| = lav cru B uu + aru| < |aee aul++aau B. ul a |0 -ul-uuueeeeeeeee la B < a8
Jane - Bu] = -vvv- 3 v- B u + 3 vl < 180 BulBBuloooo - 3 vl = 3 // uutuulue B < 86
Combining the above two inequalities we get:
lav - Bul < + 3 -+
Lemma 7. Let F (N)) = Eim1 ((s)0((s)
En f(s,) and F (s1, = =EN, (x:))
E/ (()))) and
assume
((,))
f(s))
ΣE E (+))
(s,)
6g - < <( 9" (ai) SC.
0 < f (8) f" (8) 9" (s). Then:
F (S)...., SN) - F* (s1. S
29

---

Published as a conference paper at ICLR 2024
Proof.
N (()))
(si) g" (8i)
| (1)),.,8N): - F ($1..., sN)[
$
i=1 {
Σ fr (13)
N
f(;)
(si)
9 (5t)
g" (8)) = =+))
j=1 (s))
Σ/=1 f (s))
Using Lemma 6 we get:
N [(a.]
+ (e)
2jst (18z)
Lis (1))
g + g" (si)
(+)
2
Σ f(ss)
ΣN (x)
2i f(a5)
L
((())
9 (5:) + gg" (s:)
8 5
2
2
S Σ
C-
F.2.2 PROOF
We now prove Theorem
Proof. Using S r S 1 and the discounted reward definition we get that VN:
Q (81, 8N, a) S 1-7
By definition of the structure of the Q-function and setting N we get that Vs S:
1
(s, - (s, v (s + 0 < (8, S 1-7
Again using the definition of the structure of the Q-function we get that VN:
N
Q (81, (81, SN,,)
1
Σ
a (8i, 83, 83.0) (8) s
E/= (8t, 88, a)
/=1
1
(88, a)
1==1 (8i, 8j.
1-7 --7
Repeating the above for the optimal Q-function we obtain VN:
0 us, < 0 0" (8) <
(5)
1-
1
1
0 s Q (s1. a) v 0 Q: (51 ..... (s1...., SN S
(6)
1-7
Setting N = 1 and using the soopimalliyy assumption:
(s. a) Q (s.a) lo (s o" (s,a)) <E, Vs $
(7)
We restate the theorem assumption that the attention weights are d-close for up to M objects:
a 83.))
a" Sy.a)
áō VjEL...,,,, N,
M.
(8)
Σm a (3,,
Σ a ())
30

---

Published as a conference paper at ICLR 2024
Using equation and the fact that
E RR,, from Lemma 5 we obtain that bk E [1 - 1]:
Σ 1=1 a (5i, Sy. a)
ΣØ j==1 a 8g.0)
N
(si; 8j,
a (5 Sj. a)
E Li=1 a (8,, a)
N
11 a* (8%, 81, a)
21=1
N+* (sis 01
E at (sis $1,
N
a (3).
(5,)
26 = 2No 7
Ei=1
N+ a (8), st.a)
Ei=1
D a (st, 81,0)
j=1
ΣE= a (s),
Σz=1 a" (5 83.0)
$ 2N6, Vi E [1,N]
(9)
N+ 2i=1 a (8%, 81, a)
N+k 1=l a (8), (1,88, a)
Using equation equation 7 and equation 8 from Lemma
we have:
N @ (81, ... a) Q (61,,.., SN. Set 1-7
o [I,
(10)
Note that we can decompose Q; in the following manner:
M+*
Qi (S1, 0M
la (0s, 85, a) v [s) =
Mー 25=1 88
8j, a)
M
M+k
1
(a (8i, 88, a) v (8), + W+ (82, 83, a)
fax (81, 85, v [s)
j-11 a (si,
8j, a)
j=1
j=M+1
2==1 (8), 8.: a)
M
M++ M+1 a (8t, 85.00 0)
Qi
+
Qi (8i, 1M,, 8
j=1 a (8i, ()
M++
j=1 a 88;
M+E
a (1, a)
1 as
Σ Ene a (8..
Using this decomposition for Qi and Q;:
2 ($1, m. SMKK: s/s)) Qi (51, mn (s BM+4) s
Σ j=1 a 83. a)
Em (si, 83. a)
< (#)
2, (s1, sM
Qi (s), SM, +
Σ M+k a (s,,
Σ) a 83.0)
f=1
Σ j M41 (si, 83, a)
M+k
> M4= a" (8t, 8).
###)
Q SMH1 0)
Qi (s,, SMM4s..... a) +
Σ M+k (3))
Ejet at (si,
M+*
a (Si, a)
a" S1,
##
U (8i, a)
v" a))
M+h a (s),sj, j=1
$j, a)
Σ وド=
M+E a (s;, a)
Using equation and equation 10 with Lemma we obtain the following bound for the first term:
M j=1 (Si, 8).
M De= a" (Si,
(#)
Q (M,
Q:
a) s
Σ == a (81,8,,0) 8) MM a(sussji) + Mh a" (s11850))
5= a (8r,
E/a a (s.is)aa))
Ej a(snosa)
M m1)
Qs (S1...., SM, "et...,,M,,a, + Q7 (s1, aM, 2M8 <
8 (
M
2M 5
<3
8 +
7
1 7
31

---

Published as a conference paper at ICLR 2024
Similarly for the other two terms obtain:
k
2k
(# #) E +
1
1-1
2
# E
8
1 1-9
Putting the three terms together we have:
3E 3(M+h) +
@ (81, SM+K, a) Qi (81, SM+R, al)
8
1-7
The same result is obtained for M < i S M k by similarly repeating derivations above, starting from the decomposition of Q;
Using this final result we obtain our desired upper bound:
@ (51 ... SAS+k,a) a) Q" .... SM+1:0) (M
$ 1 M+A
Q (31 mm SM+k. a) Q7 (MH s
M +
3 (M + +2 vk [L, -1]
33 +
1-7
F.3 RELAXING THE ASSUMPTION ON THE ATTENTION WEIGHTS
Theorem 2 does not exactly fit Definition because it makes an assumption the difference be- tween the optimal and approximated attention weights (equation 8. this section we present mod- ifications to the assumptions of the theorem that allow to alleviate this assumption, and bound the difference as function of the Q-value approximation error & alone.
The first additional assumption make the state space S is that individual object states distinguishable, and separated by some constant C:: Assumption 8. S
) 8; E
Vay ||s| sklli 2 C> 0).
The second assumption we add is that the attention-value function v²(-) is object-distinguishing, i.e. A ||s sy1h1 2 C > 0 I0" (si,a) v² (sj, a) > X 0:
Assumption 9. VS E SN Va E A, VN E we have: Q' ($1, a) N ΣN Q (Si, SN, aN, a), where
Q (s1, 0 = Σm (s..s,aa Σj= a (8i, v" (8j, o" ()) (R) o² E R and satis- fies 10t (st, - (s) 2 A>0
We state compositional generalization result under these assumptions as follows:
Theorem 10. Let Assumptions 8 and 9 hold. Let Q be an upproximation of Q" with an identicul structure. Assume that Vs E S.,, Va E A, VN E M1 we have
Q (81..,, 8N, Q" ($1,... SN a) <. Then Vs SMIK Va A, Vk [1, 1]:
(ss SMrk,, Q" (si, SM+K. S
12(M + k)
A(1 x(1-7) A(1 - 7) +3 +3)E, E,
8
where X > O is constant independent o
This theorem thus states that the class of functions described in Assumption 9 admits compositional generalization on (defined in Assumption &). as defined by Definition 3
We now prove Theorem 10
We start by showing that for the self attention class of functions, if two functions are s-similar for any set of objects that are X-different, must mean that the attention weights are also similar.
32

---

Published as a conference paper at ICLR 2024
Lemma 11. Consider the functions v²= : 1S+R and o² : S - RR. Assume that for any
N l...., M, and for any {s,}i=1 E SN it holds that:
lu(s)) - u(se)) 2 vjskeL...,, N, A=0,
(11)
and
Ei a(si)o(s))
Σ a (s()))((8))
Se, 8>0,
(12)
Ei1 a(si)
Σi (si)
Then
((si)
(s())) ΣØ= Σ/=aa (8j)
S 4 ,e,,...,, VN E M.
N E}=1((8)) as5))
N
Proof. From the case of N = of equation
12 we have that for every E S:
(o(s) (s) S
(13)
Consider the case of N 2. Let ā Eis a(s,)*
a(s:)
and similarly a = -E= (x,)
(=1)
We have:
(ām((1)) + (1 (2) (ā'v(s1) + ā² a')x(s)))
=|āv((11 (1 a)v(s2) (ā'v² (s1) + -1-1
a² v0 (52))
u(s))) + (1- (1)0
u(2)))
S {sāv($1) + +1 ())(s2) (ā'v" + -1--a") j²) (52)) + + (1-1111)))) ā') $ = 25
(14)
where the first incquality is from equation 1 and the second inequality is from equation 12
Manipulating the terms above we obtain:
((ā a" (u(s1) v(s2))l < 2s.
(15)
Now, from equation 11 and equation 15 we get that for any $1, 82 S2:
(ss))
(s:)
2: 48
i= 1=122
Σ asj))
E ((
A X'
The claim holds for N = 2.
Let us now consider the case of N = 3. Assume without loss of generality that v(81) < v(82) < a(a,) vss)). Denote % شزز assy)* i vs()) and 02.3
724(22) 7se(sa)
7277
7a
We rewrite the weighted sum using this notation:
E a(s))u(s) m1v(11) + (22(32) + 730(s))
Σn ((s)))
= + (72 + ra) (220(32)
+ 73
7101 + (1 m) 72.3
72 73 72 73
Then we can rewrite equation 12 for this case in the following manner:
h1 01 + 71) 02,3 (r + -
Since 02,3 is a weighted average of (12) and vs3)), we can say that |0] 02.3l 2 X. Additionally,
from the case of N = 2 of equation 12 we have that |02.3 05 < s. Using the above with similar derivation of equation 14 we obtain for 51
a(s1)
(31)) a* (8)))
S X
3 Σ a(s;)
33

---

Published as a conference paper at ICLR 2024
and a($2) + a(sa)
+ a aa (5z)
2x
S
(16)
a(,)
Σ 3 (r()))
Symmetrically we can obtain the following bound for 83:
(1(3))
a= (83)
28
(17)
Σ j=1 a(3))
Σ}=1a a- (83)
and we use equation 16 and equation 17 to obtain a bound for sg:
a(82)
x (82)
E3= (s3)
Σ}=1 a (8))
((82) + +a(8)) - a((3)
0 (82) + (83()))) - a (8(3))
Σi=1a(())
Σj=1 a Σ)=1a (83)
a(32) a(3))
(5()))
a(33)
a (33)
+
5
Σ)=1 (83)
E a()))
3 a(8))
EE= a" (8}))
$ x + 25_ 48
28
X X
The claim holds for N 3.
Now consider the case of 3 N S M. Assume without loss of generality that v(s1) < vS22) <
< v(SN). Denote the normalized attention weights 7
o(as) and the partial weighted
Σソ= als))
sum De..... = =Ei= 3,0(3 Given i E , - 1, we can regroup the values such that we have E5=
1 U1....5-1 < 0 < Gi+1....N N and corresponding weights Σ}1 7/a %o Σ Yi. Notice that by definition, 6 = Vi.
The new set of values ( éz, 0s+1,...NN) are X-far from each other since ô1.....--1 and 1:+1....,N are weighted averages of values, {UI,... 4-13 and (U:+1...., uN}, }. that are X-far from 0 vi.
From equation 12 we have that each is e-close to its 8* counterpart. To see why this is true, notice that 7/2 is normal lzed or and therefore:
a(*,) a(em)
((s))
1 a(s))u(ss)
Ei a(s))v(s))
a(,,)
j=k Ej= a(s;)
Ej=x a(s,)
j=
k= 2jzk j=nk mm=i a(sm)
So have:
a a(ss)0(0)
fk.... Gf....l C1
SE
E) (x
We therefore return to the case of N = 3 for values { 0z, Citt,..N and weights
(E 7t. N. Σj } and have that by 75 a(es) applies for every E ......N 1.
(x:)
< 1s This
a(s))
Σ (z,)
The bound on I'1 7 and InN ?5 is obtained from the case of = 2 and i = N - 1 since the
attention weights in these cases can be divided (mi 72, Σn }53 and (E Yj -N-1. 7N3 respectively.
N-2
And thus, we obtain the desired bound for all N e 1,...,M M.
34

---

Published as a conference paper at ICLR 2024
follows by substituting equation g with the following bound the nor- malized attention weights:
The proof of Theorem 10
From Assumption 9
|o(s), a) - v(s a)
2 vj
(18)
Using Assumption the e-optimality assumption and equation 18 we can bound the difference
between the sets of approximated and optimal normalized attention weights with Lemma [1 and obtain:
a 85,)
a' S5.
4 X S
VjEI, N, EN..,... M. (19)
L== a 81,0)
En a² (84, 81,0)
We substitute the bound of with 4e in the derivation and obtain the desired bound: X
[12(M + +k)
② ($1,.., m. SM+k. a) Q- ($1, aM s
(s),..,
X(1 A(1 71 +
+
35


Reviews:
[{"review_id": "YvN0Xb1MMU", "reviewer": "Reviewer_sQ6n", "summary": "This manuscript introduces an innovative approach that seamlessly integrates an object-centric model with a transformer to master structured representations crucial for goal-conditioned reinforcement learning (RL), particularly in scenarios entailing multiple objects or entities. The employed object-centric model, denoted as DLP, equips the framework with the capability to capture a structured portrayal of the environments. Concurrently, the transformer component adeptly models the dynamics of the entities and their intricate physical interactions. The clarity of the conceptual foundation is commendable, and the results showcased, particularly in the challenging realm of image-based control, are robust and hold promise. Furthermore, the paper hints at potential advancements in the field of compositional generalization. Given these strengths, I am inclined to recommend this paper for acceptance, acknowledging its significant contributions and merits. However, there are some unclear points in the current version and it would be better if the authors could provide clarification on them.", "questions": "I list the weaknesses and questions together in the above section.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- **[General idea]** Overall, the concept presented in the paper is elegantly simple and straightforward—a notable strength, as this simplicity bodes well for better understanding and potential scalability of the framework. This is of particular importance, despite the approach essentially being a synthesis of OCR and transformer-based MBRL.\n\n- **[Presentation]** The clarity and coherence of the presentation, spanning both the main paper and the appendix, are commendable, facilitating easy comprehension for the reader. Nevertheless, I have enumerated several recommendations in the subsequent sections to further enhance the manuscript.\n\n- **[Experiments]** The experiments conducted using IsaacGym validate the method's efficacy, and the exploration of compositional generalization yields valuable insights. However, I have outlined several suggestions in the sections that follow, aimed at verifying some claims made in the algorithm's design.", "weaknesses": "I list the weaknesses and questions together here.\n\n**[About the matching]** \n\nI concur with the authors regarding the permutation invariant block in the EIT, acknowledging its potential to obviate the need for matching post-OCR. However, the rationale behind the decision to forego a straightforward matching step subsequent to OCR is not entirely clear to me. Is this choice motivated by a desire for increased flexibility, or are there other factors at play? From my perspective, matching algorithms can serve as modular, plug-and-play components, exemplified by their seamless integration in slot attention mechanisms as outlined in [1]. I recommend a more thorough elucidation of this particular point in the rebuttal, as it would greatly enhance the clarity and comprehensiveness of the explanation.\n\n**[About the evaluation]**  \n\nIn order to rigorously assess the contribution of each individual component within the algorithm’s design, I recommend broadening the scope of the ablation studies conducted. Specific areas to consider include: (1) experimenting with alternative OCR methodologies in lieu of DLP, to evaluate the framework’s adaptability and performance consistency across varying OCR techniques; (2) a detailed evaluation of the impact that each component recognized by DLP has on the ultimate policy learning. This is particularly pertinent for elements that do not share a direct correlation with dynamics and rewards, such as background features. \n\n **[About the compositional generalization]**  \n\n-  Can the method generalize to the case where the novel objects (e.g., different shape but similar to the ones seen in the training, e.g., cuboid versus cube) exist during the inference phase? \n\n- Does the model possess capabilities for both extrapolation and interpolation with respect to the quantity of objects involved? To illustrate, consider a scenario wherein the model is trained on sets of 2, 4, 6, and 8 objects, and subsequently tested on sets of 3, 5, 7 (interpolation) as well as 1, 9, 10 (extrapolation). While I acknowledge the presence of some relevant results in Figure 5, a more systematic and thorough analysis of the model’s extrapolation and interpolation capabilities would be beneficial. This approach would align with the high-level conceptualization of generalization discussed in [2].\n\n **[About the interaction]**  \n\n-  I would appreciate additional clarification from the authors regarding the nature of entity interactions within the model. From my perspective, these interactions can be broadly classified into two categories: (1) interactions that influence dynamics without impacting the reward, and (2) interactions that affect both dynamics and reward. While I understand that the transformer is capable of capturing both types of interactions, a more explicit discussion on how it accomplishes this, and the implications of these interactions on the model’s performance, would be highly beneficial and contribute to a more thorough understanding of the model’s capabilities.\n\n- I am interested in understanding how the density and frequency of interactions influence the performance of policy learning within the model. Could the authors possibly quantify and assess the model’s precision in predicting interactions across varying levels of interaction density and frequency? One potential metric for this evaluation could be the accuracy of the predicted entity state in comparison to the ground truth state, especially if direct interaction capture proves challenging within the simulator. I hypothesize that a reduction in workspace size, given a constant number of objects, is likely to increase interaction occurrences. Focusing on this aspect would provide valuable insights into the model’s robustness and adaptability under different operational conditions.\n\n **[About the presentation]**\n\nMinor: I would recommend transferring the contents of either Appendix A or E to the main paper. This adjustment not only enhances the overall presentation but also efficiently utilizes the remaining available space (currently less than 9 pages).\n\n\n\n\n\n*References*\n\n[1] Locatello, Francesco, et al. \"Object-centric learning with slot attention.\" Advances in Neural Information Processing Systems 33 (2020): 11525-11538.\n\n[2] Balestriero, Randall, Jerome Pesenti, and Yann LeCun. \"Learning in high dimension always amounts to extrapolation.\" arXiv preprint arXiv:2110.09485 (2021).", "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": "Xw1rMoC9Ck", "reviewer": "Reviewer_SBTK", "summary": "The paper presents an object-centric RL model that can learn to manipulate many objects and shows generalization capabilities. The main contribution is the combination of Deep Latent Particles (DLP) as entity-centric perception pipeline and a transformer for policy and Q function. By defining a reward based on feature closeness and geometric distance there is no matching between goal-image and current image required. The only caveat is that the objects need to be filtered (the robot needs to be removed).  \nThe proposed method is pretty simple in comparison to prior work that considers entity-entity interactions.", "questions": "- how important is it that the robot is mostly white on a white background? What happens if a larger part of the robot is seen in the images? I suggest discussing this in the limitations. Also, the need to filter non-object entities. Other works would also move the robot to a particular position in the scene if part of the goal. \n- what happens if the number of latent particles is higher than the number of entities?\n- How do you compare to the above-mentioned baselines?\n\n--- Post rebuttal update. My concerns were addressed. I changed my score from 5 to 8.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- scalability to many objects\n- a relatively simple method\n- no explicit matching is required\n- multiview\n- supplementary contains important comparisons w.r.t. the reward etc.", "weaknesses": "- number of objects known. \n- missing related work and baselines:\n   - SRICS [1] is like SMOURL but dealing with object-object interactions \n   - DAFT-RL[3]: also tackles the interaction problem and baselines therein\n     DRAFT-RL is fairly recent, but it contains, IMO, relevant related work and further baselines, such as:\n     NCS [3], STOVE [4] etc.\n- supervision/filtering of entities such that only objects go into chamfer reward computation is hidden in the appendix \n- only empirical results on one type of environment: I am wondering how well it would generalize to more cluttered scenes, e.g. to a kitchen environment\n \nDetails:\n- Fig 5: too small font in the right subplot\n- Appendix A: Chamfer rewards:\n  The definition of $X_j$ and $Y_i$ after Eqn (1): what is the $i$ in the definition of $X_j$? Do I understand correctly, that it is all $x$ that have $y_j$ as their closest entity in $Y$? \n  Also afterward, when you write how to obtain standard Chamfer, the $sum_j$ is somehow missing for the second fraction. \n- I think some more information about the Generalized Density Aware Chamfer reward should go into the main text, and also that non-object particles are removed.\n- A paper that also addresses many-object manipulation with an object-centric representation is [5] (not from images)\n\n- citations/references are often published at conferences but listed as arXiv papers\n\n[1] https://proceedings.mlr.press/v164/zadaianchuk22a.html\n[2] https://arxiv.org/abs/2307.09205.pdf\n[3] https://openreview.net/forum?id=fGG6vHp3W9W\n[4] https://openreview.net/forum?id=B1e-kxSKDH\n[5] https://openreview.net/forum?id=NnuYZ1el24C", "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": "uIUsMlkULF", "reviewer": "Reviewer_PeH5", "summary": "Authors solve table-top goal conditioned tasks form pixels using particles encoding and a transformer based RL. This is an interesting improvement over previous SOTA works and its major weakness are: 1) entities are fixed cubes-with-specific-color, thus there is no possibility to generalize to other objects with different properties, 2) interaction between the objects has not properly been demonstrated and 3) related work could be improved.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "-\tThe paper is clear, well written and the topic is very interesting for the community\n-\tThe viewpoint is a nice work around to solve depth ambiguities \n-\tThe conditional goal transformed is sound and nicely implemented for an actor-critic RL.\n-\tInputs are just pixels, thus making the problem very complex.\n-\tClever use of the Chamfer Distance with the particles.\n-\tExamples animations provided that shows the system working.", "weaknesses": "**Related work**\n\nThe literature review is too focused on RL and could be improved.\n\nExample of missing SOTA object-centric perception:\nM. Traub et al. Learning what and where: Disentangling location and identity tracking without supervision,”  International Conference on Learning Representations, 2023.\n\n“learning to manipulate” The citations related to manipulation are only for image segmentation, there is a scarce but very good literature on object-centric manipulation. The majority with full observability but some from pixels.\n\n- Works based on Interaction Networks, Propagation Networks and Graph networks. E.g., A. Sanchez-Gonzalez, N. Heess, J. T. Springenberg, J. Merel, M. Riedmiller, R. Hadsell, and P. Battaglia. Graph networks as learnable physics engines for inference and control.\n\n- Examples from pixels: \nvan Bergen & Lanillos (2022). Object-based active inference. Workshop on Causal Representation Learning @ UAI 2022.\nDriess et al. \"Learning multi-object dynamics with compositional neural radiance fields.\" Conference on Robot Learning. PMLR, 2023.\n\n- Finally, regarding the use of particles for robotic control I really think that this work is seminal: \nLevine, S. et al. (2016). End-to-end training of deep visuomotor policies. The Journal of Machine Learning Research, 17(1), 1334-1373.\n\nAs an aside comment, we can find similar table-top behaviours as the one presented here using LLMs, e.g., “Palm-e: An embodied multimodal language model.”\n\n**Methods:**\n\nThis sentence requires elaboration: “Obviously, we cannot expect compositional generalization on multi-object tasks which are different for every N.”\n\nAssumption 1. Probably you are using a standard notation but please explain what is \\alpha* and v*.\n\nCould you explain the consequence of Theorem 2.\n\nWhy did you use off-policy algorithm TD3?\n\n¿Why do you need RL to train the DLP? It was mentioned that this module is pretrained, so no goal would be needed. Otherwise, you constraint the training for the defined goals that are set by the designer.\n\nGoal definition – Using the encoder. This is a common technique but prevents for proper generalization. How you would encode in this architecture non-predefined goals?, like move red objects to the left.\n\nParticles and only cubes. Using particles is very interesting, but evaluation with non-cube objects is not tested. This means that it could be that the experiments are assuming that the objects are point-mass entities. This would prevent generalization. In particular, the definition of cube-red as a single entity seems very restricted so you cannot perform behavioural operations on other shapes with different colours or other properties. \n\nAlso this rises the problem of permutation invariant, maintaining the identity of an object may is important in tasks that object permanence is needed for instance in dynamic-sequential tasks.\n\n**Experiments**\n\n¿Why adjacent goals require interactions? This can be solved reactively.\n\nI find very interesting the Ordered-Push. Should be the EIT trained for each task or it is trained on all tasks and the executed?\n\nI understand that you relegated the Chamfer distance to the Appendix, but it could be great that at least a written explanation is placed (or the equation) to understand how the rewards works.\n\nUsing this distance (and the L2)  as rewards why is RL needed, would it be enough to use a KL as objective function? Or are there other rewards used?\n\nWhat is state input? Full observability?\n\nThe agent is learning arm-object interaction thanks to the RL approach but it is not clear that the system is learning objects interaction.\nCompositional generalization. While I agree that training on N objects and then executing the task with less and more objects shows generalization capabilities. This does not necessarily endorses composition. \n\nCould you explain how the system changes when including more objects at the level of the  DLP and the EIT?\n\nBaselines: The text says: “We use DLP as the pre-trained OCR for this method for a fair comparison”, but then SMORL is only compared in the results showed with “state” access. Does this mean that this is without using pixels as input.\n\nIt is interesting that using RL also unstructured approach cannot handle the complexity. We obtained similar results using an ELBO loss. However, this makes the comparison too naïve. As the comparison of your algorithm is against full observable (state) and unstructured.\n\n**Minor comments**\n\n- Please check open quotes, in latex you can use ``word”\n- Self attention -> Self-attention", "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": "03BJl4zj2A", "reviewer": "Reviewer_aSoT", "summary": "This paper presents an object- or entity-centric RL algorithm for learning goal-conditioned manipulation. As object-centric representations, the authors use the Deep Latent Particles (DLP) method. The novelty is in the policy and Q-network, for which they propose an Entity interaction transformer (EIT), which is a transformer-based architecture to process the structured per-object latent representations. They test the method in an object manipulation task, with a robot manipulator and 2 static viewpoints provided as observations. They adopt a goal-conditioned RL setup, where the goal state is provided as target images, and introduce a Chamfer reward term to train the policy and Q function. The experiments show that their method can match the performance of another structured latent state method (SMORL), and outperform it when using image goals. Moreover, they demonstrate compositional generalization, where an agent trained on i.e. 3 colored cubes can generalize to a task with N colored cubes.", "questions": "- An important rationale for object-centric representations for RL is that once you learn a policy on one object, you can apply it to other objects (i.e. explaining the sample efficiency gap with SMORL). Why did the authors choose to discard this feature in their architecture, and would there be options to combine the strengths of both?\n\nP.S: Fig 6 caption has a typo \"mathcing\"", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper provides various novel contributions such as the transformer architecture for the Q and policy network, the Champfer reward to train policies conditioned on goal images, and demonstrates compositional generalization.\n\n- The experimental results show ablations for the various components, such as the Champfer reward, using object-centric structured latent state spaces and using multiple views.", "weaknesses": "- The method seems very tied to the experimental setup of having a robot manipulator that needs to push objects to a particular location. Some of the proposed novelties such as the Champfer reward don't seem very applicable beyond this use case.\n\n- The experiments are limited to a single environment of colored cubes. It would be interesting to see whether the approach can scale to various objects (for example YCB objects), and more cluttered scenes.\n\n- As hinted by the authors SMORL is more sample efficient, as it learns to manipulate a single object and can then generalize to the others. This seems to be an essential feature / reason to go to object-centric approaches.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)