From: NeurIPS.cc/2022/Conference/Program_Chairs
title: Paper Decision
decision: Accept

From: NeurIPS.cc/2022/Conference/Paper10705/Area_Chair_V5ys
metareview: Reviewer zPKR summarizes the paper well: "This paper an agent that learns a world model through exploratory play, which it then deploys to solve downstream tasks via a planner. The world model is an ensemble of GNNs, where nodes represent objects, which introduces an object-centric inductive bias. Effective playful exploration is achieved by using disagreement among the ensemble as a proxy for uncertainty, which allows the agent to select actions whose outcomes it cannot yet reliably predict. The authors show that this method works better than baselines, in the sense that it quickly starts engaging in “interesting” interactions (with objects)."

All reviewers unanimously vote to accept the paper. The limitations of the work are well summarized in the comment: "The most significant limitation of this paper is the limited application in the real-world setting as the proprioceptive states instead of image observations used in the paper acts as an important constraint. Besides, the heaviness in using an ensemble of GNNs may be a limitation." 

The benefits of the method outweigh the limitations. I hope authors work on extending their work to operate on sensory observations in the future. 
recommendation: Accept
confidence: Certain
award: No

From: NeurIPS.cc/2022/Conference/Paper10705/Reviewer_Fu6z
title: Thanks for the clarification.
comment: Thanks for the clarification.

This was a nice paper and my score remains very positive.

From: NeurIPS.cc/2022/Conference/Paper10705/Reviewer_zPKR
title: Thanks for response
comment: Many thanks to the authors for their response, and especially for clarifying the white noise issue. I remain impressed by the paper, and I'm now more confident in that assessment, so I have increased my confidence accordingly

From: NeurIPS.cc/2022/Conference/Paper10705/Reviewer_JJXz
title: Thanks for authors' response
comment: The rebuttal from authors have addressed my concerns in:
(1) how to train the proposed method using image observations
(2) why this method is designed to use ensemble of GNNs and the corresponding heaviness.
Therefore, I have re-rated the evaluation to 6 (weak accept).

From: NeurIPS.cc/2022/Conference/Paper10705/Authors
title: General Response
comment: We would like to thank all reviewers for their valuable feedback. Here we summarize the major changes we made to the manuscript based on the feedback we received from the reviewers:

- We added section D in the supplementary showing the multi-step prediction performance of GNN vs. MLP ensembles as well as the asymptotic behavior comparison. **(R1)**

- We added new citations to related work. **(R2)**

- We added Section E to the supplementary to compare our method to a new baseline that uses model-based control (with GNN and MLP) together with Random Network Distillation. **(R3)**

- We updated our supplementary website (https://cee-us.github.io) to contain the application of CEE-US to the Robodesk environment as proof of concept for applicability to more complex environments. **(R3)**

We also corrected minor typos in the main text and the supplementary.

From: NeurIPS.cc/2022/Conference/Paper10705/Authors
title: Response to Reviewer 3 [2/2]
comment: > Also, I wonder if it is possible to use some designed intrinsic rewards instead of the ensemble of M GNNs to train a world model with good data-efficiency?

The reason for us using the ensemble is to approximate the epistemic uncertainty of the model itself, since we want to take actions which lead to states maximally informative about the dynamics model. Note that ensemble disagreement provides a computationally tractable approximation of the predicted information gain [2]. On the other hand, e.g. a count-based intrinsic reward signal would be decoupled from the true performance of the model itself. We tested a new baseline of our method, putting together Random Network Distillation together with model-based planning. RND can be seen as an extension of count-based metrics to continuous domains. We couple RND both with a GNN model and with an MLP and test both variants. As shown in the table below, CEE-US achieves superior performance. Here we evaluate the model checkpointed after 100 iterations of free-play, showcasing the sample-efficiency of our method in terms of downstream task performance. We put the full temporal evolution of success rates as well as the interaction metrics in the Supplementary Section E.

|      Task@100      |   CEE-US    |   GNN+RND   |  MLP+iCEM   |   MLP+RND   |
|--------------------|-------------|-------------|-------------|-------------|
| Stacking 2 Obj     | 0.58 ± 0.07 | 0.01 ± 0.01 | -           | -           |
| Pick & Place 4 Obj | 0.92 ± 0.03 | 0.67 ± 0.04 | 0.48 ± 0.07 | 0.15 ± 0.01 |
| Flipping 4 Obj     | 0.56 ± 0.03 | 0.38 ± 0.09 | 0.22 ± 0.02 | 0.11 ± 0.01 |

> Could the author show the heaviness (e.g., the training time of the ensemble of M GNNs compared the MLP baseline) of the proposed method? 

Ensembles are parallelizable across the GPU and do not lead to a large slow-down. We are reporting the inference times of a single GNN and the ensemble on a NVIDIA GeForce RTX 3060 GPU on a batch of 125 samples:

| GNN Ensemble | Single GNN | MLP Ensemble | Single MLP |
|--------------|------------|--------------|------------|
|      0.00205 |   0.00115  |      0.00046 |    0.00020 |

The time comparison to the MLP Ensemble is not surprising, as internally the GNN has 3 MLPs: the node MLP, edge MLP and the global node MLP. We want to highlight that we didn’t optimize our code for computational efficiency and we still believe there is room for improvement on these times. 

**More information on object-centric unsupervised representation learning:**

As an example of an unsupervised object-centric representation learning method, we consider SCALOR [1]. This vision module could potentially be pre-trained on random data or end-to-end trained during free play. Object-centric representation learning methods such as SCALOR, represent an image / observation / scene as a set of latent feature vectors, where each vector corresponds to an object in the image / observation / scene. These vectors are also learned in an unsupervised fashion, such that there is no hand-crafted image processing that is required to capture the object-centric representations. In particular, for SCALOR, each feature vector is additionally structured in a location part and an object appearance part: $z=[z_{\text{where}}, z_{\text{what}}]$. The feature vectors $z$ would be input to the GNN. This is in fact very similar to the representation that is now fed into the network. Because similar objects would likely have similar $z_\text{what}$, due to the reconstruction loss, the learned object representation might allow generalizing better to new, but similar, objects.

[1] Jiang, Jindong, et al. "Scalor: Generative world models with scalable object representations." arXiv preprint arXiv:1910.02384 (2019).

[2] Pathak, Deepak, Dhiraj Gandhi, and Abhinav Gupta. "Self-supervised exploration via disagreement." International conference on machine learning. PMLR, 2019.

From: NeurIPS.cc/2022/Conference/Paper10705/Authors
title: Response to Reviewer 3 [1/2]
comment: Thank you for your time reviewing our work and your valuable feedback. We have improved our paper based on your concerns, as addressed in the following:

> Could the author show in detail how to train the proposed method using image observations instead of proprioceptive states as the input, and further transfer the proposed method in real-world applications?

We acknowledge that our method in its current version is only applied to proprioceptive inputs. However, as discussed in the Discussion part of our paper, our method can be extended in future work to deal with images by combining CEE-US with object-centric unsupervised representation learning methods, which exactly produce the state representations we use, namely the location and visual properties of each object/entity separately. One such example is the method SCALOR [1]. Below we provide more information about these methods. 
Combining SCALOR with our pipeline in its current form requires some non-trivial additional technical effort, such that we defer this to future work. However, we believe that it is well possible to make CEE-US deal with image input in this way.

In relation to applicability to images, you raised another concern: 

> objects with diverse geometries (e.g., real-world cabinets) do not have simple proprioceptive states, and therefore need image observations as input, otherwise the application range will be limited

In order to test if our method can deal with diverse geometries of objects with only proprioceptive state information, we now ran experiments on the Robodesk environment (https://github.com/google-research/robodesk). This environment has complex objects/entities such as a drawer, a sliding cabinet, buttons and other blocks. In our experiments, we observe the robot arm interacting with the entities during free play, e.g. opening drawer and cabinet, pushing blocks and pushing buttons. The corresponding videos can be found on our supplementary website: https://cee-us.github.io.

For Robodesk, we encode each entity’s state purely as proprioceptive information of position, quaternion, linear and angular velocities. Note that the entities have even different joint types, where the drawer has a slide joint in only $y$ direction and the sliding cabinet a slide joint only in $x$ direction, and the blocks and the ball being free joints. The different entity types are encoded as static object features and are categorical variables with one-hot encoding. An easy way of obtaining this information in real-world setups (where we do not have access to the ground-truth states) is by using motion capturing for object tracking as it is commonly done in robotics. 

The learned GNN ensemble can then be used in the extrinsic phase to solve downstream tasks zero-shot. We test opening the drawer, sliding the cabinet, pushing buttons, and moving blocks yielding the following success rates: 

| Open Drawer | Open Slide Cabinet | Push Green Button | Push Flat Block Off Table |
|-------------|--------------------|-------------------|---------------------------|
| 0.97 ± 0.02 | 0.873 ± 0.066      | 0.870 ± 0.071     | 0.58 ± 0.09               |

We agree that image features are an important part of perception and aid learning correct dynamics. However, our results in the Robodesk environment show that CEE-US with proprioceptive inputs can still be successfully applied to a more realistic environment.

From: NeurIPS.cc/2022/Conference/Paper10705/Authors
title: Response to Reviewer 2
comment: We thank the reviewer for their review and for pointing us in the direction of the two mentioned papers. They are indeed closely related to our work, and we updated our manuscript to include these citations as following:

We added Kim et. al, 2020 [1] to the list of prediction error based methods (l. 247) as states are retrospectively labeled as interesting based on the $\gamma$-Progress curiosity which is defined as the predicticion error between the current world model and an exponential average of old world models. 

We will add the following discussion to the final version of the paper (due to space constraints we could not include it now):

_Haber et al. [2] implements directed information-seeking behavior by using a loss model to predict the prediction error at future time steps. However, the agent might get stuck on states that are notoriously difficult to predict due to e.g. stochasticity._

Regarding the reviewer’s question on the white noise problem: this is indeed an important problem in novelty-based curiosity methods. However, deterministic ensembles inherently sidestep this problem as also put forward in Pathak et. al, 2019 [3]. Since the intrinsic reward in our case is computed as the disagreement of the ensemble members, the intrinsic reward goes to zero, when all members agree such that the variance of the predictions is zero. In the case of noise, with enough samples, each ensemble member should converge to predict the mean of the noise. This in turn means that the variance goes to zero and so the intrinsic reward also goes to zero. As a result, the agent shouldn’t get stuck in these situations. However, a more elegant way of dealing with inherent stochasticity would be to train stochastic ensembles outputting the parameters e.g. of a Gaussian distribution. This would help us isolate the aleatoric uncertainty from the epistemic uncertainty. This could potentially reduce the time spent in a noisy part of the state space, and learning the inherent system noise can also be exploited for risk-averseness later on in planning, as in [4]. 

Since in our work we focused on deterministic environments, we didn’t include aleatoric uncertainty modeling. However, we see this as an exciting future direction.

[1] Kuno Kim, Megumi Sano, Julian De Freitas, Nick Haber and Daniel Yamins. “Active World Model Learning with Progress Curiosity.” International Conference on Machine Learning. ICML, 2020.

[2] Nick Haber, Damian Mrowca, Li Fei-Fei and Daniel L. K. Yamins. “Emergence of Structured Behaviors from Curiosity-Based Intrinsic Motivation.” Annual Meeting of the Cognitive Science Society. CogSci, 2018.

[3] Pathak, Deepak, Dhiraj Gandhi, and Abhinav Gupta. "Self-supervised exploration via disagreement." International conference on machine learning. PMLR, 2019.

[4] Vlastelica, Marin, et al. "Risk-averse zero-order trajectory optimization." 5th Annual Conference on Robot Learning. 2021.

From: NeurIPS.cc/2022/Conference/Paper10705/Authors
title: Response to Reviewer 1
comment: We thank the reviewer for their review and for sharing our excitement on this work’s potential impact at the intersection of robotics and psychology. 

> During evaluation, you give the task reward function to a planner in order to select the appropriate action. 

The world model only learns the dynamics of the environment, in a completely task-agnostic fashion. We take the stance that a downstream task is given with a reward function that the agent can evaluate on real and imagined states. (Typically, either a distance to a desired goal, or a sparse reward). In future work, such a reward function could potentially be synthesized from a language instruction or similar. We give the planner access to this reward function, such that the cost of the state trajectories that were rolled out in the imagination of the learned GNN models can be evaluated. We would like to highlight that this is in line with other works in model predictive control, in which the planner has access to the cost/reward function. However, the planner doesn’t have any additional knowledge about the ground truth states in the environment, and only computes the reward/cost given the state trajectories produced in imagination of the model. If the model is wrong, then an actually high-reward action sequence can be mislabeled as low-reward. Therefore, accurate dynamics prediction of the learned world model is key here.

> Is is possible to run MLP + iCEM but instead of concatenating the object-representations, you simply sum them, or do something that exploits structure but is not a GNN? This would help isolate the benefits of using a GNN.

As we need to do dynamics predictions per node/object, we cannot put the aggregated information through a flat layer without a subsequent per object output head. It is therefore unclear to us, how we could put the structural bias only partially in. 
The important part for sample-efficiency is sharing the parameters of the node and edge update functions across the nodes in the GNN, as the experience gathered on one object/entity can generalize over the other objects. It would be indeed very interesting to investigate different variations of GNNs in future work. 

> It seems that an MLP leads to the same relative amount of exploration with different objects as the GNN but takes more time to do so? Why do you think it leads to similar results? 
> Figure 3: It seems structured world models lead to object interactions more quickly? With infinite time, how do you think of GNN & MLP would have the same asymptotic behavior?

Structured world models indeed lead to object interactions more quickly due to the reasons discussed in the second paragraph to the previous question, resulting in high gains in sample-efficiency. In the case of the Playground environment, the asymptotic behavior of both methods are on par in terms of interactions and success rates. However, in terms of multi-step prediction accuracy of the dynamics, the GNNs are superior to MLPs. We have now improved the manuscript and added a new section to the supplementary (see section D) to illustrate this. We have also trained the MLP+iCEM variant for an additional 100 iterations in the more complex Construction environment and added this new figure (Fig. S6) to the supplementary. We see that the asymptotic behavior of both methods in the Construction environment show a discrepancy, especially for more complex behaviors such as two object interactions and flipping. And we also show in the newly added figure S7, that the asymptotic downstream task performance generalization of the MLP variant, even after an additional 100 steps of training is inferior to CEE-US.
For Playground, the interaction metrics being similar is also due to the fact that the search space in the Playground environment is constrained by the walls, and as the objects have only 3 degrees of freedom (x-y translation and rotation around z axis), the complexity of the behavior that can be shown in free play is limited. 

> What are the numbers in Table 3? Success rate?

These are the success rates for offline RL. We have now updated the caption in the main paper for clarity.
 
> Plan2Explore seems like a good baseline since it seems similar in principal but doesn't offer combinatorial optimization. Why is this not a baseline? Could you expand on the differences to this baseline?

We will add the following discussion to the final version of the paper (due to space constraints we could not include it now):

_Plan2Explore works with the latent dynamics prediction model, that has a recurrent structure, and has been applied to domains with image observations. Since we are considering a fully-observable setting with proprioceptive states, we didn’t apply this image-based method as a baseline._

From: NeurIPS.cc/2022/Conference/Paper10705/Reviewer_Fu6z
rating: 8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.
confidence: 4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
summary: The authors propose to leverage an ensemble of GNN-based forward-models to estimate an agent's uncertainty with interacting with some part of an object-centric state-space. This uncertainty is used as an intrinsic reward signal during task-agnostic environment exploration with a model (dubbed "Free Play"). Given a fully observable MDP, the authors show that this kind of "Free Play" enables qualitatively object-centric exploration along with the ability to perform zero-shot generalization to reward functions defined over object attributes.
strengths_and_weaknesses: Strengths:
* Very clear method, results, and introduction
* High potential for impact in psychology and robotics
* While the results are somewhat expected, no one has combined disagreement with GNN based methods and the results are convincing


Weaknesses:
* The method has some hand-crafting in parsing out the agent from the objects, though this seems like a fine first-step.

questions: Experiments:
- Is is possible to run MLP + iCEM but instead of concatenating the object-representations, you simply sum them, or do something that exploits structure but is not a GNN? This would help isolate the benefits of using a GNN.

Questions:
- During evaluation, you give the task reward function to a planner in order to select the appropriate action. Your model however, only learns to predict the next state, given an action. It doesn't learn to predict the environment reward---in fact, the agent never sees environment reward. How does the planner know how to leverage the task reward to construct a plan towards that reward? Does the planner have knowledge about all ground-truth states, along with the reward R(s,a) for all ground-truth states?
- It seems that an MLP leads to the same relative amount of exploration with different objects as the GNN but takes more time to do so? Why do you think it leads to similar results? 
- Figure 3: It seems structured world models lead to object interactions *more quickly*? With infinite time, how do you think of GNN & MLP would have the same asymptotic behavior?
- What are the numbers in Table 3? Success rate?

Baselines:
- Plan2Explore seems like a good baseline since it seems similar in principal but doesn't offer combinatorial optimization. Why is this not a baseline? Could you expand on the differences to this baseline?


limitations: Limitations are adequately addressed.
ethics_flag: No
soundness: 3 good
presentation: 4 excellent
contribution: 3 good

From: NeurIPS.cc/2022/Conference/Paper10705/Reviewer_zPKR
rating: 7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
confidence: 4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
summary: This paper an agent that learns a world model through exploratory play, which it then deploys to solve downstream tasks via a planner. The world model is an ensemble of GNNs, where nodes represent objects, which introduces an object-centric inductive bias. Effective playful exploration is achieved by using disagreement among the ensemble as a proxy for uncertainty, which allows the agent to select actions whose outcomes it cannot yet reliably predict. The authors show that this method works better than baselines, in the sense that it quickly starts engaging in “interesting” interactions (with objects). They do this in two simulated environments with dynamics - a planar world with a variety objects that the agent can push around, and a robot-arm environment where the agent can stack, flip, and throw objects. The authors show that the resulting agent achieves zero-shot generalisation on downstream tasks, to numbers of objects unseen at training time.
strengths_and_weaknesses: The paper makes progress on an important problem in the RL field, namely learning with intrinsic motivation. As the authors argue, a learned world model that is reward-neutral is a key ingredient for better generalisation and transfer learning. The benefits of learning an object-centric representation in this context are well supported by the authors’ work. The experimental methodology seems sound, and the results are persuasive. The paper is well written and clear and, as far as I can determine, technically sound. Overall I think this is a very good paper that is well above the acceptance threshold for NeurIPS.

I am not sufficiently familiar with the literature on intrinsic motivation to assess the related work section, or to judge the originality of the work with great confidence (hence my “low confidence” rating). However, the following paper got my attention a few years back, which seems very related, and isn’t cited:

Nick Haber, Damian Mrowca, Li Fei-Fei, Daniel L. K. Yamins. Emergence of Structured Behaviors from Curiosity-Based Intrinsic Motivation. In Proceedings of the 40th Annual Meeting of the Cognitive Science Society, CogSci 2018. (https://arxiv.org/abs/1802.07461)

I notice that a subset of the authors of that paper have a more recent ICML paper that also looks relevant, and isn't in the bibliography: https://arxiv.org/abs/2007.07853


questions: See below
limitations: I believe the ensemble disagreement method that the authors use is susceptible to the white noise problem, wherein the agent is motivated to fixate on dynamics that are random and therefore cannot be learned. Perhaps the authors could comment on this.
ethics_flag: No
soundness: 4 excellent
presentation: 4 excellent
contribution: 3 good

From: NeurIPS.cc/2022/Conference/Paper10705/Reviewer_JJXz
rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
confidence: 4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
summary: This paper proposes a curiosity-driven method for exploration via structured world models. It uses GNN ensemble disagreement for computing intrinsic motivation signals in RL, which achieves sample-efficient and interaction-rich exploration in multi-object manipulation environments. Further, the trained world models combined with planning-based methods can achieve zero-shot generalization to many object manipulation tasks. Experiments shows the data-efficiency, the zero-shot generalization ability and manipulation accuracy of the proposed method.


strengths_and_weaknesses: Strengths:
1.	This paper achieves quite a high sample-efficiency compared with other methods using the curiosity-driven method.
2.	This paper uses GNN with high combinatorial generalization abilities, which makes the method generalizes to environment with novel number of objects.
3.	The world models trained using the proposed method can easily generalize to many downstream manipulation tasks without specific trainings.


Weakness:
1.	The world model takes proprioceptive states of objects, instead of the image observation. This may be a hinder to the transfer to the real-world applications, in that: (1) we can not easily get the accurate proprioceptive states of objects; (2) objects with diverse geometries (e.g., real-world cabinets) do not have simple proprioceptive states, and therefore need image observations as input, otherwise the application range will be limited; (3) training using images has a much harder setting than proprioceptive states, which may potentially require the adjustment of the method proposed in this paper. Therefore, in my opinion, it is really important to train the proposed method based on image observations, instead of proprioceptive states.
2.	Though resulting in good data-efficiency, using an ensemble of M GNNs makes it heavy to train the model. Training with designed intrinsic rewards (e.g., the number of agent-object / object-object interactions, the state prediction accuracy, etc) may make the training simpler with only one GNN, while keeping the curiosity-driven good data-efficiency.

questions: 1.	Could the author show in detail how to train the proposed method using image observations instead of proprioceptive states as the input, and further transfer the proposed method in real-world applications?
2.	Could the author show the heaviness (e.g., the training time of the ensemble of M GNNs compared the MLP baseline) of the proposed method? Also, I wonder if it is possible to use some designed intrinsic rewards instead of the ensemble of M GNNs to train a world model with good data-efficiency?

limitations: The most significant limitation of this paper is the limited application in the real-world setting as the proprioceptive states instead of image observations used in the paper acts as an important constraint.
Besides, the heaviness in using an ensemble of GNNs may be a limitation. And this limitation may be even serious when the setting becomes more complex, which requires the ensemble to contain more GNNs with more model parameters.

ethics_flag: No
ethics_review_area: I don’t know
soundness: 3 good
presentation: 3 good
contribution: 2 fair