Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The presented work addresses the task of robotic manipulation tailored to environments containing many objects. Task planning in such environments entails large number of possible combinations of actions given the number of objects. To address this challenge, the authors propose to use an attention module in combination with a graph neural network to predict the effects of the manipulation action, limiting the objects that need to be included into the execution planning. Experiments are conducted in simulated environment using a robotic arm interacting with two sets of shapes. The model performs well when for both test scenarios, one where the model operates on a task that is was trained on, and two on a task that was not presented during training showing the generalization capabilities of the presented approach.
main_review: The authors present an interesting approach with several novel ideas. Using an attention mechanism to address the complexity of environments containing larger numbers of objects to focus only those objects that are going to be affected by an action is interesting and novel. Based on the ablation study, however, that this mechanism does seem overly important. Would that change if the number of objects in the environment increases dramatically? The approach performs well in the scenario where the task was shown during training. More importantly, this is also the case for the scenario were the task that was not shown during training and shows the generalization abilities of the approach. The numbers for the two metrics (RMSE and Hits @ 1) are reasonable. The baselines, however, don't seem to be sufficient since (Qi et al. 2021) is either not failing to provide results or scores sufficiently worse compared to the presented approach, and the second baseline is a modified version of the presented approach. Veerapaneni et al. (2019) and (Kipf et al. 2020) use a similar test environment and (Huang et al., 2020) compare their results to (Kipf et al. 2020). The authors neither use those methods in their  evaluation nor do the discuss their reasons for excluding those. It would be great if the authors could provide more details on this part.
The submission is well written, the authors are able to communicate their ideas well, but placing the related work section at the end of the submission is confusing. 
summary_of_the_review: The presented approach introduces interesting and novel ideas and shows that the model is able to generalize to tasks that it was not trained on. Only using one exterior baseline method that does not perform well seems not sufficient given that three related publications are evaluated using a similar setup. While the experimental section needs to be improved, I would rate the positive aspects the approach, generalization abilities and the novelty higher and would, therefore, argue for the submission to be accepted.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This work proposes a pipeline to learn factored world models to predict the robot actions’ effects. This work proposes to leverage a graph neural network to extract state information and adopt the contrastive loss to train the encoder and latent transition model. This paper applies the proposed approach to pick-and-place tasks with simple shapes, and extensive experiments indicate its effectiveness.
main_review: Strengths:

s1)  The paper proposes to leverage a graph neural network to extract the state feature and learn a latent dynamic model. The contrastive loss is applied to update the neural network weights.

s2) The paper is cleary-written and easy to follow. 

Weaknesses:

w1) In Sec 2, "We will assume a setting in which the state of the world is represented as an image, which has been pre-processed into a factorized state s = $ <s_1 , s_2 , ..., s_K >$ in which each si is an image centered on the ith object". The assumption may not hold in many applications where some objects are fully occluded and can not be perceived.

w2) In the experiments, the evaluation setting is too simple. It would strengthen the paper if more complicated manipulation tasks (including more complex shapes and manipulation tasks pushing, pulling, throwing) are included.
summary_of_the_review: This work proposes a pipeline to learn factored world models to predict the robot actions' effects leveraging the graph neural network to encode the latent space and the contrastive loss to update the neural network's weight. However, the technical novelty is limited, and the current setting is too simple. It's unclear whether the proposed pipeline could work with complicated manipulation tasks.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: In this paper, the authors expand on previous works on contrastive learning with factored world models (specifically the C-SWM model by Kipf et al. - 2020) and apply their model for the task of predicting the dynamics of object manipulation with a continuous action space. Such previous works have shown that by leveraging residual graph neural networks (GNNs) to contextualize the learned latent states with possible actions, the model learns representations that are equivariant to permutation of objects in the scene, thus allowing scaling to a high number of objects regardless of the combinatorial explosion of the state space. Unlike previous works, where new objects are discovered by the model from raw pixels, the authors here use the simulation environment to segment objects appearing in the scene and factor them as part of the state representation. Another difference is that instead of factorizing the action space, they use two action primitives for picking and placing objects from given continuous coordinates and propose an action attention module that predicts how each action will affect each object in the next state. They claim that by departing from the standard action-object dependence assumption made by previous works, the current model can effectively model the effect of future actions for structured tasks unseen during training. They finally propose to stack multiple GNN layers instead of a single layer as mostly practiced by earlier works. The authors evaluate their methodology on two pick-and-place tasks regarding building structures of abstract object shapes or cubes. The objective is to correctly predict the effect of certain action sequences in other blocks in order to build a certain structure (e.g. pile). They sample noise from the correct action sequence in order to generate negative examples for contrastive learning. They train in some sequence configurations and test for zero-shot performance in unseen tasks (new structures of objects). They report results for both mean square error of the block
position as well action ranking (how well the final state of the action sequence matches the final latent state). 
main_review: Strengths:
 - The authors demonstrate that factored world models can indeed propagate multi-object structured relation modeling in robotic manipulation scenarios while enabling zero-shotting in unseen object/action configurations. 
 - The addition of the attention module as well as stacking multiple GNN layers evidently increases  (although marginally and not always) the performance of the model across both tasks and evaluation metrics
 - The proposed factorization of the state representation is relatively straightforward to implement in the context  of robotic manipulation applications, as actuators are most commonly trained in simulation environments that  come packed with object segmentation tools (e.g. segmentation masks in Pybullet) 

Weaknesses:
 - Even though this work seems to be an extension of the factorized world model proposed by Kipf et al. (2020) there are no direct comparisons given with their model in the experiments section.
 - The authors claim they wish to explore the potential of factored world models in generalizing to unseen tasks, but in the concerned environments, new "tasks" simply correspond to novel sequences of picking and placing actions (i.e. novel configurations), something previous works that impose permutation equivariance already demonstrate that they can achieve. This work's novelty focuses on the application of an attention mechanism and stacking multiple GNN layers. Both additions provide minor performance boosts (and not robustly across all experiments, as seen in rows 3, 4, and 5 of the bottom part of Table 1). 
 - A few grammatical mistakes and forgotten to actually cite works in the first two paragraphs of Section 2.
- Inconsistent presentation of figures and tables throughout the text.
summary_of_the_review: This work presents an adaptation of an existing factored world model specifically tailored for robotic manipulation, using motion primitives (picking/placing objects) as the action space and pre-segmenting objects in order to factorize them into the state representation. The paper's novelty comprises the addition of an attention operation to contextualize future actions with latent object states. Even though experimental results suggest positively towards the potential of this method in task generalization, the lack of extensive comparison with previous factored
(structured) world models do not convince me 100% of the merits of the proposed methodology, especially when considering the capability of the standard C-SWM model to explore new object categories which are lost in the proposed model in favour of a "black-box" object detection system that factorizes the scene as input to the model.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: The paper proposes Object-Factored world models, a world model learned using contrastive loss, that tends to generalize over the number of objects in the scene. It also alleviates the assumption that previous world models hold about the object-action association and the author(s) propose an "action-attention module" that estimates the probability of an action affecting a particular object (using self-attention). Their experiments mainly focus on robotic manipulation tasks and they show generalization on unseen-test tasks with little drop in performance as compared to the training set.

main_review: Strengths:
1) The paper is well written and is easy to understand in a go. 
2) The action-attention module seems to be a novel idea especially if I think about a multi-agent scenario where it'd be useful to have such a module to direct the policy (or an accurate world model learning) and be sample-efficient (as opposed to just concatenating the action's latent to every object's latent state).

3) I appreciate the author(s) adding Figure 5 as it provides a  good visualization of the action-attention module.


Weaknesses:


4) Action Attention Module: Although I like the idea of having such a module, I am not convinced by the experiments performed to evaluate this. The task of pick and place is simple enough to implicitly figure out which object does an action correspond to -- this can be seen from Table 1, FWM and FWM - No attention. There is barely any improvement in the Hits@1 (and even the block's position if you consider the std. deviation along with the mean). This implies that not having action attention suffices for the network to figure out which object is the action intended on.

This is not to say that the action attention is useless. It is to say that for this specific task of pick and place, it is not very useful.

5) Claim about combinatorial generalization: The combinatorial generalization is due to the use of GNNs which is a core module of C-SWM [1]. Hence I do not agree with the claim that the zero-shot generalization shown in the paper is a contribution by FWM rather the author(s) just extends the claim made by C-SWM by showcasing it in a different environment (the robotic manipulation environment).

6) Zero-Shot generalization: Based on my points (5) and (6), it is important to make the distinction that the action attention does not contribute to the generalization (at least from the results shown on pick and place task in Table 1 on both Cubes and Shapes environment), rather it is the GNN that contributes to the zero-shot generalization. These two are often worded one after the other which makes the reader think that it is the action-attention that is major contributing factor to the generalization. Kindly make this clear in the text.

7) "No RGB" Experiment: There is no discussion about this experiment but in the environments considered by the author(s) all the objects are in red and hence RGB doesn't really play a role in the task. This is also shown clearly in Table 1, where the performance is almost similar (a slight increment/decrease on the RMSE/Hits@1) to FWM.

8) Real-world experiments? - The first line of the Experiments section says "Our empirical evaluation focuses on pick-and-place robotic manipulation tasks performed in simulation as well as on a real-world robotic arm", but I do not see any real-world robotic arm experiments. Could the author(s) clarify this?

9) Why only Robotic manipulation? - The action-attention module and (other components from C-SWM) are very generic, so why did the author(s) restrict themselves to only robotic manipulation environments? What about Atari games or other object-based RL environments? It would be interesting to see if action-attention plays an important role there (especially in getting to know the agent in Atari game a.k.a the most controllable agent).

10) In Fig 5, are the objects colored only for visualization purposes or is this the RGB input?

[Minor thing which I have not considered for rating the paper]

11) There are some places where the citations have not been mentioned (rather the literal word "citations" appear twice in the text). For example: pg 2(last line) and pg 3 (first paragraph, 5th line). This is being disregarded by me for evaluating the paper since it seems to be a genuine mistake caused by overlooking. I'd request to add the citations in the future version of the manuscript.

---
References

[1] Contrastive Learning of Structured World Models, Thomas Kipf et. al, ICLR 2020
summary_of_the_review: Overall I feel that the paper lacks a concrete quantitative evaluation of the action attention module (which is the core contribution of the paper). The experiments don't prove that this module is helpful for robotic manipulation tasks (though it can be useful for some other tasks which the author(s) have to investigate further). Hence, I would like to vote for a weak rejection of the paper at this stage. I will be making a final decision after engaging with the author(s) post rebuttal.
review_title: null

Meta-review: This paper presents a GNN-based attention mechanism and tests it on a robotic stacking task.

While all the reviewers agree that this work is novel and interesting, they also are unanimous (even after the rebuttal) in pointing to the insufficient experimental evaluation of the proposed method.

I encourage the authors to incorporate the feedback of all the reviewers.
Decision: Reject