A structured understanding of our world in terms of objects, relations, and hierarchies is an important component of human cognition. Learning such a structured world model from raw sensory data remains a challenge. As a step towards this goal, we introduce Contrastively-trained Structured World Models (C-SWMs). CSWMs utilize a contrastive approach for representation learning in environments with compositional structure. We structure each state embedding as a set of object representations and their relations, modeled by a graph neural network. This allows objects to be discovered from raw pixel observations without direct supervision as part of the learning process. We evaluate C-SWMs on compositional environments involving multiple interacting objects that can be manipulated independently by an agent, simple Atari games, and a multi-object physics simulation. Our experiments demonstrate that C-SWMs can overcome limitations of models based on pixel reconstruction and outperform typical representatives of this model class in highly structured environments, while learning interpretable object-based representations.
This paper tackles the problem of learning an encoder and transition model of an environment, such that the representation learnt uses an object-centric representation which could favor compositionality and generalisation. This is trained using a contrastive max-margin loss, instead of a generative loss as previously explored. They do not consider RL or follow-up tasks leveraging these representations and transition models yet.
They perform an extensive assessment of their model, with many ablations, on 2 gridworld environments, one physical domain, and on Atari.
The paper is very well motivated, easy to follow, and most of its assumptions and decisions are sensible and well supported. They also provide interesting assessments and insights into the evaluation scheme of such transition models, which would be of interest to many practitioners of this field.
Apart from some issues presented below, I feel that this work is of good quality and would recommend it for acceptance.
1.	The model is introduced in a very clear way, and most decisions seem particularly fair. I found the presentation of the contrastive loss with margin to be clear, and the GraphNet is also well supported (although see question below).  However, two choices are surprising to me and would deserve some clarification and more space in the main text, instead of the Appendix:
	a.	Why does the object extractor only output a scalar mask? This was not extremely clear from reading the main text (and confused me when I first saw Figure 1 and 3a), but as explained in the Appendix, the CNN is forced to output a sigmoid logit between [0, 1] per object channel. 
	This seems overly constraining to me, as this restricts the network to only output “1 bit” of information per “object”.
	However, maybe being able to represent other factors of these objects might be necessary to make better predictions?  
	This also requires the user to select the number of output channels precisely, or the model might fail. This is visible in the Atari results, where the “objectness” is much less clear.
	 Did you try allowing the encoder to output more features per objects? 
	Obviously this would be more complicated and would place you closer to a setting similar to MONet (Burgess et al. 2019) or IODINE (Greff et al. 2019), but this might help a lot.
	b.	It was hard to find the dimensionality D of the abstract representation $z_t$. It is only reported in the Appendix, and is set to $D=2$ for the 2D gridworld tasks and $D=4$ for Atari and the physics environments.  These are quite small, and the fact that they exactly coincide with your assumed sufficient statistics is a bit unfortunate.  
	What happens if D is larger? Could you find the optimal D by some means?
2.	The GraphNet makes sense to me, but I wondered why you did not provide $a_t^j$ to $e_t^{(i, j)}$ as well? I could imagine situations where one would need the action to know if an interaction between two slots is required.
3.	Similarly, the fact that the action was directly partitioned per object (except in Atari where it was replicated), seemed slightly odd. Would it still work if it was not directly pre-aligned for the network? I.e. provide $a_t$ as conditioning for the global() module of the GraphNet, and let the network learn which nodes/edges it actually affects.
4.	In your multi-object contrastive loss, how is the mapping between slot k in $z_t$ and $\tilde{z}_t$ performed? Do you assume that a given object (say the red cube) is placed in the same $k$ slot across different scenes/timesteps? This may actually be harder to enforce by the network than expected (e.g. with MONet, there is no such “slot stability”, see [1] for a discussion).
5.	It was unclear to me if the “grid” shown in Figure 3 (b) and 5 is “real”? I.e. are you exactly plotting your $z_t$ embeddings, and they happen to lie precisely along this grid? If yes, I feel this is a slightly stronger result as you currently present, given this means that the latent space has mirrored the transition dynamics in a rather impressive fashion.
6.	Related to that point, I found Figure 3 b) to be slightly too hard to understand and parse. The mapping of the colours of the arrows is not provided, and the correspondence between “what 3D object is actually moving where” and “which of the coloured circles correspond to which other cubes in the image” is hard to do (especially given the arbitrary rotation).  Could you add arrows/annotations to make this clearer?  Alternatively, presenting this as a sequence might help: e.g. show the sequence of real 3D images, along with the trajectory it traces on the 2D grid.
7.	Figure 4 a) was also hard to interpret. Seeing these learnt filters did not tell much, and I felt that you were trying too hard to impose meaning on these, or at least it wasn’t clear to me what to take of them directly. I would have left this in the Appendix. Figure 4 b) on the other hand was great, and I would put more emphasis on it.
8.	There are no details on how the actual test data used to generate Table 1 was created, and what “unseen environment instances” would correspond to. It would be good to add this to the Appendix, and point forward to it at the end of the first paragraph of Section 4.6, as if you are claiming that combinatorial generalization is being tested this should be made explicit. I found Table 1 to be great, complete, and easy to parse.
9.	It would be quite interesting to discuss how your work relates to [1], as the principles and goals are quite similar.  On a similar note, if you wanted to extend your 2D shape environment from a gridworld to a continuous one with more factors of variations, their Spriteworld environment [2] might be a good candidate.
References:
[1] Nicholas Watters, Loic Matthey, Matko Bosnjak, Christopher P. Burgess, Alexander Lerchner, “COBRA: Data-Efficient Model-Based RL through Unsupervised Object Discovery and Curiosity-Driven Exploration”, 2019, https://arxiv.org/abs/1905.09275
[2] Nicholas Watters, Loic Matthey, Sebastian Borgeaud, Rishabh Kabra, Alexander Lerchner, “Spriteworld: A Flexible, Configurable Reinforcement Learning Environment”, https://github.com/deepmind/spriteworld/