Paper: Entity-Centric Reinforcement Learning for Object Manipulation from Pixels (ICLR-2024)

Reviewer comment (Novelty):
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.

Author response:
*Focus on Pushing Tasks* -  We have added experiments with a more complex task of orienting a T-shaped object. We believe that various other manipulation tasks could easily be handled by our policy, given a proper reward function during training.

*Chamfer Reward* - Designing a reward function from images is a difficult task, and while we make some progress on it using the Chamfer reward (which can also be extended to the orientation task), we acknowledge that it is not yet solved.

Reviewer comment (Experiments):
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.

Author response:
*Objects in the Environment* - The experiments we have added in the new revision may shed light on some of your concerns regarding the types of objects our method can deal with. The push-T task demonstrates manipulation of an object that is not a cube with more complex dynamics (see Appendix B.4 and “Push-T” tab on our website). In addition, we present zero-shot generalization results on variations of object properties, including cuboids of different dimensions (results are presented in Table 6 and discussed in Appendix B.2  in the new revision). While zero-shot generalization to objects with different properties is partial on some scenarios, the success rates are not negligible and therefore hint at potential for few-shot generalization. These results also imply that our method can handle training with different types of objects to begin with.

*Cluttered Scenes* - The Group-Push experiment shows the ability of our agent to deal (zero-shot) with a large amount of objects densely occupying the table (see our website for demonstrations).

Reviewer comment (Experiments):
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?

Author response:
*Best of Both Approaches* - We designed our method to learn simultaneous multi-object manipulation because it is applicable to a larger group of tasks than learning on a single object is, as we demonstrate in our experiments (see Figure 4d, 4f and Table 1). We present types of sub-goals that can’t be achieved without taking all sub-goals into account (e.g. Ordered-Push).

Compared to SMORL, there is indeed a tradeoff between sample efficiency and performance that originates from simplification of the learning problem. In some cases, which we claim to be closer to real-world scenarios, oversimplification significantly hurts performance.

Regarding combining the strengths of both, this could be an interesting direction for future work. One way to do this would be to create a curriculum of an increasing number of objects similar to Li et al. [1]

[1] Richard Li, Allan Jabri, Trevor Darrell, and Pulkit Agrawal. Towards practical multi-object manipulation using relational reinforcement learning. In 2020 ieee international conference on robotics and automation (icra), pp. 4051–4058. IEEE, 2020.

Reviewer comment (Presentation):
The literature review is too focused on RL and could be improved.

Author response:
Thank you for pointing out the various papers to us, widening the scope of our related work to both model-based planning and active inference methods. We have added them to the related work section of the new revision.

Reviewer comment (Novelty):
An example of missing SOTA object-centric perception work is: M. Traub et al. "Learning what and where: Disentangling location and identity tracking without supervision,” International Conference on Learning Representations, 2023.

Author response:
Thank you for pointing out the various papers to us, widening the scope of our related work to both model-based planning and active inference methods. We have added them to the related work section of the new revision.

Reviewer comment (Novelty):
As an aside, similar tabletop behaviours as the one presented here can be obtained using LLMs, e.g., “Palm-e: An embodied multimodal language model.”

Author response:
Regarding Palm-e, it is an LLM-based planning algorithm that assumes access to a low-level manipulation policy. Here we use RL to learn a single policy that solves the entire task, both single-object manipulation and planning to achieve the overall goal while accounting for object-object interaction.

Reviewer comment (Theory):
The sentence “Obviously, we cannot expect compositional generalization on multi-object tasks which are different for every N” requires elaboration.

Author response:
*”This sentence requires elaboration: “Obviously, we cannot expect compositional generalization on multi-object tasks which are different for every N.””*

We expect compositional generalization to an increasing number of objects where the objects are similar to the ones seen during training and the basic task remains the same with the increased number of objects, such as pushing each object to a desired goal location.

Reviewer comment (Theory):
Assumption 1 needs clarification — please explain what is \alpha* and v*.

Author response:
*”Assumption 1. Probably you are using a standard notation but please explain what is \alpha* and v*.”*

We reference the supplementary background for the attention mechanism (Appendix E) where $\alpha$ represents the attention weights (produced from keys and queries) and $v$ represents the values. * denotes the optimal versions of those functions (per our assumption that the optimal Q function has an attention structure).

Reviewer comment (Theory):
The consequence of Theorem 2 should be explained.

Author response:
*”Could you explain the consequence of Theorem 2.”*

The theorem bounds the approximation error of a learned Q-function vs. the optimal Q-function, when both have a self-attention structure. The Q-function is learned on environments with $M$ objects or less, and is epsilon-close to the optimal Q-function in this case (less than $M$ objects). The optimal Q-function is optimal for all $N$. We show that with no additional training, the approximation error increases linearly in the number of objects. This is a notion of compositional generalization defined through the Q-function, and we show that a Q-function with a self-attention structure obtains it.

What this implies is that the Q-function used to optimize a policy on $M$ objects is not very far from the optimal Q-function for $M+k$ objects and thus the same policy should obtain similar returns even when deployed on environments with an increasing number of objects.

Reviewer comment (Experiments):
The rationale for using the off-policy algorithm TD3 should be clarified.

Author response:
*”Why did you use off-policy algorithm TD3?”*

TD3 is a SOTA off-policy model free deep RL algorithm based on DDPG which has been used in similar tasks in previous work, including the original HER paper. It is well suited for goal-conditioned RL because it allows dynamically relabeling goals in the replay buffer transitions sampled during training, thereby improving sample efficiency.

Reviewer comment (Theory):
Why RL is needed to train the DLP should be explained, since it was mentioned that this module is pretrained and no goal would be needed. Otherwise, training would be constrained to the defined goals set by the designer.

Author response:
*”¿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.”*

We do not need RL to train the DLP nor do we claim this. The DLP is pretrained on image data collected by a random policy. This is not related to the goals. Goals in our experimental setup are defined as images of the scene with a specific object configuration.

Reviewer comment (Theory):
The goal definition using the encoder — while common — prevents proper generalization. The reviewer asks how non-predefined goals, such as "move red objects to the left," would be encoded in this architecture.

Author response:
*”Goal 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.”*

What are you referring to by predefined goals? The goals in our setting are not predefined, they are provided as an image of the desired configuration by the user. This image is converted to the latent particle representation using the pre-trained DLP encoder, which is able to encode images it has not seen during training but that come from the same distribution. Moreover, it is able to generalize to images with a larger/fewer number of objects than it has seen during training. These are demonstrated by the success of our method as well as in the generalization capabilities in the Cube-Sort scenario.

If by “move red objects to the left” you are referring to goal specification through natural language, this is very possible and integrates well with our proposed method, assuming access to a reward function. The goal entities will consist of the natural language tokens specifying the goal and their relation to the DLP-extracted image entities will be learned via the cross-attention block in the EIT. Reward design might be a challenge in this case but this is not exclusive to our method. We leave such an interesting investigation to future work.

Reviewer comment (Experiments):
The evaluation only uses particles and cube objects. While using particles is interesting, evaluation with non-cube objects is not tested. This could imply that the experiments assume that objects are point-mass entities, which would prevent generalization. Furthermore, the definition of cube-red as a single entity seems very restricted, preventing behavioural operations on other shapes with different colours or properties.

Author response:
*”Particles 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.”*

We have studied the generalization capabilities of our method to object properties not seen during training, please refer to Appendix B.2 of the new revision.

We additionally provide experimental results on non-cube objects in our revision. The new task involves pushing a T-shaped block to a goal orientation. This task demonstrates that the DLP latent representation is expressive enough to infer implicit object properties (orientation from visual latent attributes) as well as the ability of the EIT to capture these properties and learn how to manipulate them to achieve desired goals.

Reviewer comment (Theory):
Why adjacent goals require interactions is unclear, as this can be solved reactively.

Author response:
*”¿Why adjacent goals require interactions? This can be solved reactively.”*

Adjacent-Goals can sometimes be solved reactively but there is a non-negligible probability that the goal won’t be achieved if not taking object-object interaction into account. This is supported by our empirical results, and is demonstrated in the supplementary material and our website (please see “Entity-Entity Interaction” tab on our website). In the video demonstration, we see that a policy that does not take object-object interaction into account consistently pushes objects away from their goal while attempting to bring another object to its desired subgoal. These cases happen often enough that SMORL experiences a performance drop of ~20% in success rate compared to non-adjacent goals and compared to our method on Adjacent-Goals (see Table 1).

Reviewer comment (Experiments):
For the Ordered-Push task, clarification is needed on whether the EIT should be trained for each task individually or trained on all tasks and then executed.

Author response:
*”I find very interesting the Ordered-Push. Should be the EIT trained for each task or it is trained on all tasks and the executed?”*

The EIT is trained for each task separately except for Adjacent-Goals. The reason for this is that the physical constraints of Ordered-Push and Small-Table are different and need to be learned (i.e., cubes will fall off the smaller table or cubes will run into a wall).

Reviewer comment (Presentation):
Although the Chamfer distance is detailed in the Appendix, it would be helpful to have at least a written explanation or the equation in the main text to understand how the rewards work.

Author response:
*”I 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.”*

Thank you for pointing this out, we have added an explanation about the Chamfer reward in the main text.

Reviewer comment (Experiments):
When using Chamfer distance (and L2) as rewards, it should be explained why RL is needed — would it be sufficient to use a KL objective function, or are there other rewards used?

Author response:
*”Using 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?”*

Can you please elaborate what you mean by using the KL Divergence as the objective function? Objective for policy learning? KL between what two probability distributions?
RL is needed because while the reward might be differentiable with respect to the state, the dynamics of the environment are not. RL allows maximizing long horizon discounted rewards and model-free RL allows doing so without explicitly modeling the dynamics of the environment.

Reviewer comment (Writing):
Clarification is needed on the definition of state input — is it full observability?

Author response:
*”What is state input? Full observability?”*

State input is defined precisely in Appendix C.1 and is the ground truth $(x,y)$ coordinates of each object. For the factored representation, we add a one-hot vector as entity identifying features.

Reviewer comment (Evaluation):
Training on N objects and then executing the task with fewer or more objects shows generalization capabilities, but this does not necessarily demonstrate compositionality.

Author response:
*”Compositional 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.”*

We refer to compositional generalization in our context as generalizing to a different number of objects (Section 4.4). While we acknowledge that this is not the only type of compositional generalization, this term is used in previous literature and includes our use of the term for variations in number of entities (“A survey on compositional generalization in applications”, Lin et al.).

Reviewer comment (Experiments):
Regarding baselines, the text states: “We use DLP as the pre-trained OCR for this method for a fair comparison,” but SMORL is only compared in the results provided with “state” access. Does this mean that this is without using pixels as input?

Author response:
*”Baselines: 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.”*

In the image-based SMORL we use DLP as the OCR. The comparison with SMORL in the image-based case is in the Chamfer Reward section of the Appendix A.2. In the graphs of the main text we compare both image-based and state-based methods with GT reward. Image-based SMORL is designed in a way that does not allow using it with a reward that is not image-based, therefore does not appear in those graphs.

Reviewer comment (Experiments):
The point that even RL-based unstructured approaches cannot handle the complexity is interesting, but comparing against a naive ELBO loss baseline makes the comparison too naive, especially as your algorithm is compared against full observability (state) and unstructured approaches.

Author response:
*”It 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.”*

Can the reviewer please elaborate on this point? Which comparison do you believe is too naive and why?

If you are referring to the comparison between the state-based unstructured approach to our image-based method, we don’t believe this comparison to be naive. While it might be clear that structured approaches should perform better here, unstructured data from ground truth state observations still explicitly holds all the relevant attributes (and only those) to perform the task. What is mostly missing is the permutation invariance as an inductive bias. The fact that it learns slower than our method which is image-based, emphasizes the importance of this structure for object-centric tasks.

Reviewer comment (Theory):
number of objects known.

Author response:
**Number of Objects**

We wish to emphasize that the **number of objects is not known** and is not used explicitly in any part of the algorithm. The number of objects should generally affect the choice of the DLP model’s “number of particles” hyper-parameter. In our experimental setup, we set “number of particles”=24, which is much larger than the amount of objects/entities we consider in any of our environments. We found that choosing a relatively high upper bound is beneficial for RL training as it reduces the chance of objects in the scene not being represented by at least a single particle, and allows representing large entities (such as the robot arm) with multiple particles. In addition, this allows for zero-shot generalization to a significantly larger number of objects, as we demonstrate in the Group-Push experiment.

Reviewer comment (Novelty):
missing related work and baselines: SRICS [1] is like SMOURL but dealing with object-object interactions.

Author response:
[SRICS](https://arxiv.org/abs/2109.04150)

We find this work an interesting extension of SMORL (by the same authors) for *state inputs* which takes into account object-object interaction. 

*Method Complexity*: compared to our method from state inputs, it requires training a separate network to consider relationships and interactions. This network is essentially a world-model that is used during both RL training to define rewards and during inference to decide on the ordering of sub-goals to solve using the RL policy. Our method incorporates all of the above aspects in a single transformer-based model, making it more robust, easy to train and requiring less hyper-parameters.

*Assumptions*: An assumption of SRICS is that when trying to achieve a given sub-goal, the agent should minimize its effect on other sub-goals. This assumption is explicitly integrated in the method through the “selectivity reward” used to train the SAC agent. This can be a restrictive assumption that could negatively impact performance. Consider for example the task of moving objects located at one side of the table to the other. The SRICS agent will move the objects one by one, although the optimal behavior with respect to timestep efficiency (and a negative distance reward) would be to move multiple objects simultaneously. Our method does not make such assumptions.

*Inputs*: While we find this work related to ours, it is not a relevant baseline because it is not designed for image inputs. The authors mention the extension of this method to image-based inputs as future work, implying that this would not be a trivial extension.

Reviewer comment (Novelty):
missing related work and baselines: DAFT-RL[3] also tackles the interaction problem and baselines therein.

Author response:
[DAFT-RL](https://arxiv.org/abs/2307.09205)

This paper proposes an object-centric model-based approach, where the state is not only factored into entities but also to their individual attributes such as physical and visual properties. It does so by learning multiple separate networks which include 3 graphs. One is in charge of modeling the relationships between object attributes, the agent’s action and the reward. A second graph is in charge of modeling object-object interaction and the third is in charge of agent-object interaction and forward dynamics prediction.

*Method Complexity*: The proposed method consists of several components and training stages. The stages are roughly separated into learning each of the graphs detailed above for each class of objects which is assumed to be known. Each network is learned by maximizing a distinct objective using different data. These do not include the policy network which is learned subsequently using the various model components. Compared to their method, our transformer-based model is trained to maximize a single objective and learns relationships between entity attributes as well as entity-entity interaction in an implicit manner in order to perform multi-object tasks.

*Assumptions*: DAFT-RL makes several assumptions that explicitly affect their algorithm design choices. They assume a known set of object classes and that their attributes are both accessible from images and are disentangled in the latent image representation. The Push-T task we consider in our experiments is an example where this assumption does not hold, as both color and orientation are implicit in the latent visual features. In addition, the mass or friction coefficient of the objects for example, can’t be inferred from images alone. They also assume each action directly affects a single object at a time. We do not make these assumptions. The attention mechanism in the EIT along with representing the action as an entity in the Q-function network enables modeling agent interaction with multiple objects simultaneously.

Reviewer comment (Experiments):
missing related work and baselines: DRAFT-RL contains relevant related work and further baselines, such as NCS [3], STOVE [4] etc.

Author response:
[NCS](https://arxiv.org/abs/2303.11373)

This method presents a structured approach for solving object rearrangement tasks. They use a modification of SLATE as an object-centric representation of images. They then use it to construct a factored transition graph between clustered states where the nodes are states and the edges are actions. They use this graph to plan and execute actions on a high-level object rearrangement environment.

*Method Complexity*: Generally, NCS is a complex method containing multiple manual (not learned) task specific substeps in each stage of the algorithm, increasing the number of hyper-parameters and making it less generally applicable. These include clustering states using K-Means, isolating which object moved in each transition based on the difference in state attributes, explicitly matching state and goal objects using the Hungarian algorithm (assumes a 1-to-1 match exists) and at each timestep, choosing the object with the largest distance from the goal as the next candidate to move, not taking into account goal constraints such as the presence of other objects in the goal position.

*Assumptions*: This work makes several simplifying assumptions. They do not consider entity-entity interaction, do not handle occlusion (no agent in image), do not deal with low level control (assume high level actions). Our method is designed to address challenges that arise in the real world where these assumptions do not hold. We consider entity-entity interaction via the attention mechanism and handle occlusions by integrating multiple viewpoints in our EIT architecture. Additionally, our method is aimed at learning both low level and high level control, which prevents us from using NCS as a baseline in our experimental setup.

Reviewer comment (Reproducibility):
supervision/filtering of entities such that only objects go into chamfer reward computation is hidden in the appendix.

Author response:
**Filtering of Entities in the Chamfer Reward**

The filtering of entities for the reward is an implementation choice that is not directly connected to our method but to the specific tasks we consider in our experiments.

Even when solving from GT simulation state, the agent is “filtered” from the reward in the sense that the reward is only defined on the objects of interest. This is the standard task and standard reward for block pushing in previous works. Additionally, previous work such as SMORL consider performance metrics that exclude the agent. While SMORL aims to remain completely self-supervised, including the arm in the reward when it is not the goal of the task could potentially negatively affect performance.

While disregarding the agent in the reward when working with state observations is trivial, it is not trivial when working with images. Without an object-centric representation of images, this separation is not feasible because the agent and objects are not disentangled in the representation. Using an OCR not only enables this, but makes it straightforward with minimal supervision. The supervision we detail in the Appendix A.1 is very light and easy to acquire both in simulation and the real world.

To your suggestion, we will refer to this choice in the main text.

Reviewer comment (Experiments):
only empirical results on one type of environment; uncertainty about how well it would generalize to more cluttered scenes, e.g., to a kitchen environment.

Author response:
**Environments**

We have demonstrated that our method solves core challenges in multi-object manipulation such as occlusions, interactions between objects, a continuous action space, and compositional generalization. Our new results also demonstrate success on more geometrically complex objects (Push-T). These challenges, which were largely ignored in previous studies, are important for real world scenes such as kitchens. There may be additional challenges, but we believe our results show a clear progress towards handling such scenes.

Reviewer comment (Presentation):
Fig 5: too small font in the right subplot.

Author response:
Thank you for your response. To your suggestion, we address the supervision required for the Chamfer reward in our setting in the main text and in more detail in Appendix A.1. We have increased the font size in Figure 5, but we can further increase it if the reviewer still finds it necessary. We will add the Push-T results and discussion to the main text in the next revision.

Reviewer comment (Writing):
Appendix A: Chamfer rewards — The definition of $X_j$ and $Y_i$ after Eqn (1): what is the $i$ in the definition of $X_j$? Question about whether it is all $x$ that have $y_j$ as their closest entity in $Y$.

Author response:
*Chamfer Reward Definition Clarification* - You understood the definition correctly. $x_i$ refers to $x \in X$ that have $y_j$ as the closest entity in $Y$. In the standard Chamfer distance, we have a single sum in each term which is over $x_i$ or $y_j$. In the DAC distance, the additional sum comes from the fact that we first take an average over all $x$ that are mapped to the same $y$, and then average over the different groups that are defined by this mapping. We have added details about the GDAC reward to the main text in the new revision.

Reviewer comment (Writing):
Appendix A: Chamfer rewards — When describing how to obtain standard Chamfer, the $\sum_j$ is missing for the second fraction.

Author response:
*Chamfer Reward Definition Clarification* - You understood the definition correctly. $x_i$ refers to $x \in X$ that have $y_j$ as the closest entity in $Y$. In the standard Chamfer distance, we have a single sum in each term which is over $x_i$ or $y_j$. In the DAC distance, the additional sum comes from the fact that we first take an average over all $x$ that are mapped to the same $y$, and then average over the different groups that are defined by this mapping. We have added details about the GDAC reward to the main text in the new revision.

Reviewer comment (Reproducibility):
More information about the Generalized Density Aware Chamfer reward should go into the main text, including the fact that non-object particles are removed.

Author response:
*Chamfer Reward Definition Clarification* - You understood the definition correctly. $x_i$ refers to $x \in X$ that have $y_j$ as the closest entity in $Y$. In the standard Chamfer distance, we have a single sum in each term which is over $x_i$ or $y_j$. In the DAC distance, the additional sum comes from the fact that we first take an average over all $x$ that are mapped to the same $y$, and then average over the different groups that are defined by this mapping. We have added details about the GDAC reward to the main text in the new revision.

Reviewer comment (Experiments):
Question — 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? Suggestion to discuss this in the limitations, along with the need to filter non-object entities. Other works might move the robot to a particular position in the scene if part of the goal.

Author response:
*“How important is it that the robot is mostly white on a white background?”*

The fact that the robot is mostly white on a white background could potentially make it harder for the DLP model to distinguish them from each other when modeling the scene. In preliminary experiments during the development of our method, the DLP modeled the robot in the background. After decreasing the capacity of the latent background particle by setting its dimension to $1$, this no longer happened and the DLP successfully distinguished the robot from the 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.”*

There are many cases in our setting where a large part of the robot is seen in the images. This can be seen on our website or in the supplementary material where the particle locations are visualized on the video - many of the particles are on the agent. We do not see this as a limitation of our method as the EIT learns which entities are relevant for solving the task.

Filtering non-object entities in the EIT via masking in the attention mechanism slightly improves sample efficiency and requires no supervision over the input, as we use the latent transparency attribute of the DLP to do so (i.e., it is learned unsupervised as part of DLP pretraining). This is not a limitation of our method but rather a feature that is enabled via the use of attention and DLP. Additionally, it is not required for the agent to learn. When conducting experiments during the development of our method we saw that the agent achieves the same performance in slightly longer training times without the attention masking.

*”Other works would also move the robot to a particular position in the scene if part of the goal.”*

If the robot position is part of the goal, there would be no need to filter the robot particles in the Chamfer reward calculation. We believe an object pushing task should not include the robot position as part of the goal hence the need for filtering.

Previous work such as RIG added the robot position as part of the goal - we see this as a limitation of their method as they can not separate the goal to individual entities due to the single-vector representation of the scene. SMORL considered the robot position as part of the goal during both training and inference, although the definition of their task is object manipulation, as highlighted by their performance metric which only considers the object distances from their desired goals. We believe that SMORL chose to include the robot as part of the goal in order to remain completely self-supervised. We see this as a limitation as well, as it may negatively affect performance. Such negative effects include counter productively moving the agent to its goal before moving the rest of the objects to their desired goal.

Reviewer comment (Experiments):
Question — what happens if the number of latent particles is higher than the number of entities?

Author response:
*“What happens if the number of latent particles is higher than the number of entities?”*

This is exactly the case in our setting. The number of latent particles is 24, much larger than the number of entities in the scene which is 4 (3 objects and the agent). This design choice is motivated by the fact that we prefer redundancy of particles in order to decrease the chance of entities being overlooked by the DLP model. This is emphasized in the video on our website (“In the Eyes of the Agent”) where multiple particles represent the agent and each object.

In the case it is the other way around (more entities than latent particles), this would be a problem, but also a bad choice from the user’s side. This is not a limitation of the method as the number of particles is a controllable hyper-parameter.

Reviewer comment (Experiments):
Question — How do you compare to the above-mentioned baselines?

Author response:
**Related Work Conclusion**

We compare our method to what we see as the most relevant SOTA baselines: goal-conditioned model-free algorithms that are image-based. There are clear advantages to model-free over model-based methods (as is true for the other direction) and we see our work as a step forward in goal-conditioned object-centric model-free RL.

Reviewer comment (Experiments):
In 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.

Author response:
**Evaluation**

(1) We agree this would be an interesting ablation. We intend to experiment with the slot-based Slot-Attention ([1] Locatello et al.), which was shown to perform well in the investigation performed in OCRL ([2] Yoon et al.). 
This ablation will require some time to find the right hyper-parameters for SLATE and integrate it with our code. We have started implementing this and will conclude the experiments after the rebuttal period.

We pre-trained Slot-Attention (SA) using the same data as DLP with 10 slots (the maximum number feasible for 128x128 resolution given our computational resources). Notably, the training process was considerably slower than training DLP on similar hardware. Visual results can be found in the following image:  https://i.postimg.cc/PtWm4r1q/slots-25.png. It's important to highlight that, in the SA representation for our data, the robot and the table are often contained in the same slot, and occasionally, multiple cubes are grouped in a single slot instead of being individually represented. We are currently training our model using object-centric representations from the pre-trained Slot-Attention model. However, we anticipate that the imperfect SA decomposition may impact final performance, and we will provide detailed results in the upcoming version.

Slot-Attention image columns: (1) original image, (2) reconstruction, (3) mask, (4) refined mask, (5) ->(14) slots.



(2) We would like to emphasize that we discard the background features when extracting the latent representation from images, as it is constant in all of our experiments. We can add experiments that ablate some DLP components. Note that our new results on Push-T show that the policy makes use of the "feature" component in DLP, as it is the only feature that (implicitly) contains orientation information.

Reviewer comment (Experiments):
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?

Author response:
*Novel Objects* - This is an interesting question. Our method in its current version is not specifically designed to deal with novel objects zero-shot. Dealing with novel objects would require generalization from both the DLP and the EIT.
We would expect our method to zero-shot generalize to novel objects in case:
- They are visually similar to objects seen during training.
- Their physical dynamics are similar to the objects seen during training.

To test our hypothesis, we deploy our trained agent in environments including modifications to shape, color and mass. Results are presented in Table 6 and discussed in Appendix B.2  in the new revision. Our main conclusion from this study is that while zero-shot generalization to objects with different properties is only partial, the success rates are not negligible and therefore hint at potential for few-shot generalization.

Reviewer comment (Experiments):
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].

Author response:
*Extrapolation / Interpolation* - Due to long training times, we trained our image based method on at most 3 cubes. To test the interpolation capabilities of our method we trained an agent on a mixture of episodes from the 1-Cube and 3-Cubes environments and tested the performance on the 2-Cubes environment. The interpolation results for 2 cubes were similar to the results for 2 cubes with an agent trained only on 3 cubes.
For further insight on the extrapolation capabilities, we present a comparison of the performance of our agent with respect to the number of objects it was trained on for 1, 2 and 3 objects. Results are presented in Table 5 and discussed in Appendix B.2 in the new revision.

Reviewer comment (Theory):
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.

Author response:
**Interaction**

We do not understand the questions regarding interaction. Since our method is model-free, the EIT is not trained to predict the next state but only the action (conditioned on the current observation and goal).
Therefore, we do not understand how to classify the interactions to dynamics/reward, nor how to report the precision in predicting interactions. We can visualize the attention pattern in the transformer, but it's not clear to us how to use this signal to answer the questions the reviewer raised. If we misunderstood the reviewer's question, we would be happy for a clarification.

Reviewer comment (Presentation):
Minor: 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).

Author response:
**Presentation**

Thank you for the suggestion, we have moved additional details about the Chamfer reward to the main text.