Entity-Centric Reinforcement Learning for Object Manipulation from Pixels

Thank you for acknowledging the novelty and contribution of our work.

Experimental Setup

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.

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).

Sample Efficiency

Comparison to SMORL - We would like to emphasize that SMORL’s policy does not generalize to multiple cubes - it is trained to perform single-cube manipulation and does so also during inference. A scripted meta-policy is used to cycle between the objects in the environment and solve the sub-goals one by one during inference, disregarding the other objects while doing so. While SMORL is more sample efficient in the sense that it only requires learning single-object manipulation, it is fundamentally limited in performing complex multi-object tasks that require modeling interactions, as we demonstrate in our experiments (see Figure 4d, 4f and Table 1).

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.

Thank you for acknowledging the complexity of the tasks we deal with and the strengths of different aspects of our method. We also thank you for your questions and interest in our work.

Related Work

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.

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.

Method

”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.

”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).

”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.

”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.

”¿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.

”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.

Related Work - Continued

STOVE

STOVE is an object-centric video prediction model which also demonstrates world-modeling capabilities for solving a sequential decision making task via planning (Monte Carlo Tree Search). They extend their video prediction model to dealing with control by conditioning the predictions on actions and predicting not only the next frame but also the reward to enable reward-maximizing planning. In their evaluation, they consider visually simple 2D environments and discrete action spaces. It is not clear how this method would scale to complex 3D environments and continuous action spaces such as the ones we consider in our experimental setup. Additionally, expanding their model to multiview inputs is highly non-trivial.

DAFT-RL

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.

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.

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.

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.

Thank you for acknowledging the strengths of our work and simplicity of our method. We also thank you for your questions and referring us to the various methods related to our work.

Related Work and Baselines

We have added the various methods you referred us to in the related work section. Our method is model-free and arguably much simpler than the methods you mentioned. Additionally, we introduce multiview inputs which are seamlessly integrated in a single model, and show zero-shot generalization to a substantially larger number of objects than previous work.

In the following, we address each method in comparison to ours in detail.

SRICS

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.

Linear Relation Networks (LRN)

This method uses state inputs. It is similar to SMORL in the sense that it only considers relationships between objects and goals, therefore not accounting for object-object interaction. Due to the above, we believe SMORL is a better fit as a baseline which covers the essentials of the LRN method and improves over them (image-based, goal-conditioned via images).

CEE-US

Thank you for pointing out this paper to us. Although it is not an RL method and uses state observations, it does deal with exploration and compositional generalization in the multi-object manipulation setting. We have added it to the related work section.

NCS

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.

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.

Thank you for the insightful and constructive feedback, we appreciate your acknowledgement of our work’s contribution.

Matching

The tasks we consider in our experiments require matching between entities in different views of the state as well as matching between state and goal entities. There are several reasons for our choice to take a learning approach for the matching rather than explicitly incorporating modular matching components in the algorithm. The overall rationale behind the decision is achieving flexibility, simplicity and robustness, as follows:

• When working from image observations, a match for each entity does not always exist, for example in the case of occlusion in one of the viewpoints (happens frequently due to the robot arm). This case requires special handling in an explicit matching algorithm and there are not always a set of clear rules as to how to handle these discrepancies. This also requires incorporating additional hyper-parameters such as a threshold for what is considered a match in order to prevent matching with the next closest match even if it is not the same object (e.g matching a red cube to a purple cube by mistake). This increases complexity and hurts flexibility, as the hyper-parameters are not input dependent and may increase the chance of mismatches.

• DLP often assigns multiple particles to a single object which creates a situation where there is not a one-to-one match. This makes it difficult to aggregate matching particles to a single entity of fixed size. Our learning approach deals with this automatically via the attention mechanism.

• Generality: we designed our algorithm to have specific structure but try to keep it as general as possible. One direction for future work includes goal specification via natural language. In this case, the state tokens would be particles while the goal tokens would consist of natural language tokens. This is another case where the matching is not straightforward and a one to one match does not necessarily exist.

In addition, explicit matching is not always a simple task. Matching entities in consecutive observations within a single-view object-centric framework poses significant challenges, as demonstrated in prior works (https://arxiv.org/abs/2107.09240). While incorporating matching algorithms, such as the Hungarian algorithm, is feasible, it introduces computational complexities that hinder scalability to datasets involving more than 3 objects. In our multi-view setting, this challenge intensifies as we must match entities across views, contending with potential occlusions and the absence of a one-to-one correspondence. Furthermore, the extensive literature on multi-view image matching (https://openaccess.thecvf.com/content_ICCV_2017/papers/Maset_Practical_and_Efficient_ICCV_2017_paper.pdf, https://arxiv.org/pdf/2205.01694.pdf) highlights the inherent difficulty of this problem, independent of the additional complexities associated with training a generative object-centric model.

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.