Contrastive Learning of Structured World Models
Authors: Thomas Kipf
Elise van der Pol
Max Welling
Venue:  (2020)
Keywords: state representation learning
graph neural networks
model-based reinforcement learning
relational learning
object discovery

Abstract: 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). C-SWMs 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.

Published as a conference paper at ICLR 2020
CONTRASTIVE LEARNING OF STRUCTURED WORLD MODELS
Thomas Kipf
Elise van der Pol University of Amsterdam UvA-Bosee Delta Lab e. e. vanderpo @uva.nl
Max Welling
University of Amsterdam t.n.k s.n.kipf@uva..o
University of Amsterdam CIFAR m. wew inguuva......... .nl
ABSTRACT
A structured understanding of our world in terms of objects, relations, and hierar- chies an important component of human cognition. Learning such structured world model from raw sensory data remains challenge. As a step towards this goal, we introduce Contrastively-train trained Structured World Models (C-SWMs), C- SWMs utilize contrastive approach for representation learning in environments with compositional structure. We structure each state embedding as set of ob- ject representations and their relations, modeled by graph neural network. This allows objects to be discovered from raw pixel observations without direct super- vision part of the learning process. We evaluate C-SWMs on compositional environments involving multiple interacting objects that can be manipulated inde- pendently by an agent, simple Atari games, and multi-object physics simulation. Our experiments demonstrate that C-SWMs can overcome limitations of models based pixel reconstruction and outperform typical representatives of this model class in highly structured environments, while learning interpretable object-based representations.
INTRODUCTION
Compositional reasoning in terms of objects, relations, and actions is central ability in human cognition (Spelke & Kinzler, 2007). This ability serves as a core motivation behind range of recent works that aim at enriching machine learning models with the ability to disentangle scenes into objects, their properties, and relations between them (Chang et al., 2016; Battaglia et al., 2016; Watters et al., 2017; van Steenkiste et al., 2018; Kipf et al., 2018; Sun et al., 2018; 2019b; Xu et al., 2019). These structured neural models greatly facilitate predicting physical dynamics and the consequences of actions, and provide strong inductive bias for generalization novel environment situations, allowing models to answer counterfactual questions such as "What would happen if I pushed this block instead of pulling it?".
Arriving at a structured description of the world in terms of objects and relations in the first place, however, is a challenging problem. While most methods in this area require some form of human annotation for the extraction of objects or relations, several recent works study the problem of object discovery from visual data in completely unsupervised self-supervised manner (Eslami et al., 2016; Greff et al., 2017; Nash et al., 2017; van Steenkiste et al., 2018; Kosiorek et al., 2018; Jan- ner et al., 2019; Xu et al., 2019; Burgess al., 2019; Greff al., 2019; Engelcke al., 2019). These methods follow generative approach, i.e., they learn to discover object-based representa- tions by performing visual predictions or reconstruction and by optimizing an objective pixel space. Placing loss in pixel space requires carefully trading off structural constraints latent variables vs. accuracy of pixel-based reconstruction. Typical failure modes include ignoring visu- ally small, but relevant features for predicting the future, such bullet in an Atari game (Kaiser et al., 2019), or wasting model capacity on visually rich, but otherwise potentially irrelevant features, such static backgrounds.
To avoid such failure modes, we propose to adopt a discriminative approach using contrastive learn- ing, which scores real against fake experiences in the form of state-action-state triples from an

---

Published as a conference paper at ICLR 2020
experience buffer (Lin, 1992), in similar fashion as typical graph embedding approaches true facts the form of untity-relation-ent triples against corrupted triples or fake facts.
We introduce Contrastively -eriinnt--teedd Structured World Models (C-SWMs), class of models for learning abstract state representations from observations an environment. C-SWMs learn a set for each object in partieular observation. Environment transitions
of abstract state variables, are modeled using a graph neural network (Scarselli et al., 2009; Li et al., 2015; Kipf & Welling, 2016; Gilmer al., 2017; Battaglia et al., 2018) that operates on latent abstract representations.
This paper further introduces a novel object-level contrastive loss for unsupervised learning of object-based representations. We arrive this formulation by adapting methods for learning trans- lational graph embeddings (Bordes et al., 2013; Wang et al., 2014) to our use case. By establishing connection between contrastive learning of state abstractions (François-Lavet et al., 2018) and relational graph embeddings (Nickel et al., 201 2016 we hope to provide inspiration and guidance for future model improvements in both fields.
al., 2018; Thomas
In set of experiments, where we use novel ranking- based evaluation strategy, we demonstrate that C-SWMs learn interpretable object-level state abstractions, accurately learn to predict state transi- tions many steps into the future, demonstrate combinatorial generalization to novel environment configurations and learn identify objects from scenes without supervision.
2 STRUCTURED WORLD MODELS
Our goal is to learn object-oriented abstraction of particular observation environment state. In addition, we would like to learn an codiin-condiiooned transition model of the environment that takes object representations and their relations and interactions into account.
We start by introducing the general framework for contrastive learning of state abstractions and transition models without object factorization in Sections 2-122.2, and the following describe variant that utilizes object-factorized state representations, which we term Structured World Model.
2.1 STATE ABSTRACTION
We consider an off-policy setting, where operate solely on buffer of offline experience, e.g.. obtained from an exploration policy. Formally, this experience buffer B = {(+t, 0t: 81+11 })= con- tains T' tuples of states s, S, actions a, E A, and follow-up states 8t-1 S, which reached after taking action ag We do not consider rewards as part our framework for simplicity.
Our goal is to learn abstract latent representations 2; E Z of environment states st S that discard any information which is not necessary to predict the abstract representation of the follow- up state 2++1 Z after taking action dg. Formally, have encoder E S Z which maps
observed states to abstract state representations and a transition model T x A > Z operating solely abstract state representations.
2.2 CONTRASTIVE LEARNING
Our starting point is the graph embedding method TransE (Bordes et al., 2013): TransE embeds facts from knowledge base I. = {(e,, rt, ot) which consists of entity (ex, (t,,Ti,0) 0t), where e, the subject entity (analogous to the source state 8, in our case), r, is the relation (analogous to the action in our experience buffer), and is the object entity (analogous to the target state $1+1).
triples
TransE defines the energy of triple (e,, 0t) as H d(F(e;) + G(re), F(o))), where F (and G)) are embedding functions that map discrete entities (and relations) to R where the dimen- sionality of the embedding space, and d(-,-) denotes the squared Euclidean distance. Training is carried out with an energy-based hinge loss (LeCun et al., 2006), with negative samples obtained by replacing the entities in a fact with random entities from the knowledge base.
We can port TransE to our setting with only minor modifications. As the effect of an action is in general not independent of the source state, replace G(r,) with T( ax), i.e., with the transition
2

---

Published as a conference paper at ICLR 2020
function, conditioned on both the action and the (embedded) source state via 24 energy of a sate-action-stal triple then can be defined as follows: H d(2) + ( az), 2(+1).
E(st). The overall
This additive form of the transition model provides strong inductive bias for modeling effects of actions the environment as translations the abstract state space. Alternatively, one could model effects as linear transformations rotations in the abstract state space, which motivates the of graph embedding method such as RESCAL (Nickel et al., 2011), CompleX (Trouillon et al., 2016), or HolE (Nickel et al., 2016b).
With the aforementioned modifications, we arrive at the following energy-based hinge loss:
L d(z + T(z,, or), 2++1) + max(0, 7 d(2,, 2++1))
(I)
defined for a single (st, Gr: St+1) with a corrupted abstract state 2, E(;)). 5, is sampled random from the experience buffer. The margin is a hyperparameter for which we found y 1 to be good choice. Unlike Bordes al. (2013), we place the hinge only on the negative term instead of on the full loss and we do not constrain the norm of the abstract states 2g, which we found to work better in our context (see Appendix A.3). The overall loss is to be understood as an expectation of the above over samples from the experience buffer B.
2.3 OBJECT-ORIENTED STATE FACTORIZATION
Our goal is to take into account the compositional nature of visual scenes, and hence we would like learn a relational and object-orient orreetee model of the environment that operates on a factored abstract state space Z 2 X Zk. where K is the number of available object slots. We further assume an object-factorized action space AA Ai .... x AK. This factorization ensures that each object is independently represented and it allows for efficient sharing of model parameters across objects in the transition model. This serves as a strong inductive bias for better generalization to novel scenes and facilitates learning and object discovery. The overall C-SWM model architecture using
representations is shown Figure
GNN
CNN
MLP
Object extractor
m
Object encoder
Transition model
2/ A
Contrastive loss
CNN-based object extractor, 2) an MLP-based object encoder, 3) a GNN-based relational transition model, and 4) an fbctrtefactorized contrastive loss. Colored blocks denote abstract states for particular object.
Figure 1: The C-SWM model is composed of the following components:
Encoder and Object Extractor split the encoder into two separate modules: 1) a CNN-based object extractor Fext: and 2) MI.P-based object encoder Func The object extractor module is CNN operating directly on image- based observations from the environment with K feature maps in its last layer. Each feature map m [EEn(())) can be interpreted as object mask corresponding to one particular object slot, where [...]x denotes selection of the k-th feature map. For simplicity, only assign single feature map per object slot which sufficed for the experiments considered in this work (see Appendix A.4). To allow for encoding of more complex object features (other than, e.g., position/velocity), the object extractor can be adapted produce multiple feature maps per object slot. After the object extractor module, we flatten each feature map m. (object mask) and feed it into the object encoder E The object encoder shares weights across objects and returns an abstract state representation: z = (m)f) with 2f Zz. We set Z = RD in the following, where is hyperparameter.
Relational Transition Model We implement the transition model as graph neural network (Scarselli et al., 2009; Li et al., 2015; Kipf & Welling, 2016; Battaglia et al., 2016; Gilmer et al., 2017; Battaglia al., 2018), which allows model pairwise interactions between object states
3

---

Published as a conference paper at ICLR 2020
while being invariant to the order in which objects represented. After the encoder stage, we have an abstract state description of E Z and an action of E A₁ for every object in the scene. We repre- sent actions one-hot hot vectors (or vector of zeros if no action is applied note that other choices are possible, e.g., for continuous action spaces. The transition function then takes as input the tuple of object representations 2 (z/.-..,+3 2{)) and actions d, (a) -.... ) at a particular time step:
particular object), but
A = T(2,, a.) = GNN( +f))))))))))))
(2)
T(2; at) is implemented as graph neural network (GNN) that takes 2f as input node features. The model predicts updates A., = (Az, (Azi....... Azk) The object representations for the next time step + ). The GNN consists of node update functions
are obtained via 2 = (z) + A2) .........
fudee and edge update functions Sedge with shared parameters across all nodes and edges. These functions are implemented as MI.Ps and we choose the following form of message passing updates:
(R.5) fedge 2/11)
et
A = foode (oo al. Σnzj 6t ()
(4)
(i.) an intermediate representation of the edge interaction between nodes and j. This corresponds to single round of node-to-edge and edge-to-node message passing. Alternatively, one could apply multiple rounds of message passing, but we did not find this to be necessary for the experiments considered in this work. Note that this update rule corresponds to message passing on a fully-connected scene graph, which is O(²)). This can be reduced to linear complexity by reducing connectivity to nearest neighbors in the abstract state space, which we leave for future work. We denote the output of the transition function for the k-tt object as A= * = T (2z, az) the following.
where et
Multi-object Contrastive Loss We only need to change the energy function to take the factoriza- tion of the abstract state space into account, which yields the following energy H for positive triples and H for negative samples:
K H= H
K
= K
Σ k=1
d(zt + (zn A= K
Σ d( kk
& 2441),
(5)
where 2/ is the k-th object representation of the negative state sample ă contrastive loss for a single sttte-cction-ee sample from the experience buffer then takes the form:
E(5t). The overall
L = H + max H).
3 RELATED WORK
For coverage of related work in the area of object discovery with autoencoder-based models, we refer the reader to the Introduction section. We further discuss related work on relational graph embeddings in Section 2.2.
Structured Models of Environments Recent work modeling structured environments such as interacting multi-object or multi-agent systems has made great strides in improving predictive accuracy by explicitly taking into account the structured nature of such systems (Sukhbaatar et al., 2016; Chang et al., 2016; Battaglia et al., 2016; Watters et al., 2017; Hoshen, 2017; Wang et al., 2018; van Steenkiste et al., 2018: Kipf al., 2018; Sanchez-Gonzalez et al., 2018: Xu et al., 2019). These methods generally make use of some form of graph neural network, where node update functions model the dynamics of individual objects, parts or agents and edge update functions model their interactions and relations. Several recent works succeed in learning such structured models directly from pixels (Watters et al., 2017; van Steenkiste et al., 2018; Xu et al., 2019; Watters et al., 2019), but in contrast to our work rely on pixel-based loss functions. The latest example in this line of research is the COBRA model (Watters et al., 2019), which learns an action-conditioned transition policy on object representations obtained from unsupervised object discovery model (Burgess et al., 2019). Unlike C-SWM, COBRA does not model interactions between object slots and relies on pixel-based loss for training. Our object encoder, however, is more limited and utilizing an iterative object encoding process such as in MONet (Burgess et al., 2019) would be interesting for future work.
4

---

Published as a conference paper at ICLR 2020
Contrastive Learning Contrastive learning methods are widely used in the field of graph rep- resentation learning (Bordes et al., 2013; Perozzi et al., 2014; Grover Leskovec, 2016; Bordes et al., 2013; Schlichtkrull et al., 2018; Veličković et al., 2018). and for learning word representations (Mnih Teh, 2012; Mikolov et al., 2013). The main idea is to construct pairs of related data exam- ples (positive examples, e-g., connected by edge in graph co-occuring words in sentence) and pairs of unrelated or corrupted data examples (negative examples), and use loss function that scores positive and negative pairs different way, Most energy-based losses (LeCun al., 2006) are suitable for this task. Recent works (Oord et al., 2018; Hjelm et al., 2018; Hénaff et al., 2019; Sun al., 2019a; Anand al., 2019) connect objectives of this kind to the principle learning representations by maximizing mutual information between data and learned representations, and successfully apply these methods to image, speech, and video data.
State Representation Learning State representation learning in environments similar to ours is often approached by models based on autoencoders (Corneil et al., 2018; Watter et al., 2015; Ha & Schmidhuber, 2018; Hafner et al., 2019; Laversanne-Fing et al., 2018) via adversarial learning (Kurutach et al., 2018; Wang et al., 2019). Some recent methods learn state representations without requiring decoder back into pixel space. Examples include the selectivity objective in Thomas al. (2018), the contrastive objective in François- Lavet et al. (2018), the mutual information objective in Anand et al. (2019), the distribution matching objective in Gelada al. (2019) using causality- based losses and physical priors in latent space (Jonschkowski & Brock, 2015; Ehrhardt et al., 2018) 8). Most notably, Ehrhardt et al. (2018) propose a method learn an object detection module and physics module jointly from raw video data without pixel-based losses. This approach, however, can only track a single object at a time and requires careful balancing of multiple loss functions.
EXPERIMENTS
Our goal of this experimental section to verify whether C-SWMs can 1) learn to discover object representations from environment interactions without supervision. 2) learn accurate transition model in latent space, and 3) generalize to novel, unseen scenes. Our implementation is available under https:tpp..... /tt:://ggtttie thu.. coo/tkiipi
4.1 ENVIRONMENTS
We evaluate C-SWMs on two novel grid world environments (2D shapes and 3D blocks) involving multiple interacting objects that can be manipulated independently by agent, two Atari 2600 games (Atari Pong and Space Invaders), and a multi-object physics simulation (3-body physics). See Figure 2 for example observations.
For all environments, we use a random policy to collect experience for both training and evaluation. Observations are provided 50 50 x 3 color images for the grid world environments and as 50 x 50 x 6 tensors (two concatenated consecutive frames) for the Atari and 3-body physics environments. Additional details on environments and dataset creation can be found Appendix
Move left
Move right
(c) Atari Pong
(d) Space Invaders
(e) 3-Body Physics
2D Shapes
(b) 3D Blocks
Figure 2: Example observations from block pushing environments (a-b), Atari 2600 games (c- d) and a 3-body gravitational physics simulation (e). In the grid worlds (a-b), each block can be independently moved into the four cardinal directions unless the target position is occupied by another block or outside of the scene. Best viewed color.

---

Published as a conference paper at ICLR 2020
4.2 EVALUATION METRICS
In order to evaluate model performance directly in latent space, we make use of ranking metrics, which are commonly used for the evaluation of link prediction models, as in, e.g., Bordes et al. (2013). This allows us to assess the quality of learned representations directly without relying on auxiliary metrics such as pixel-based reconstruction losses, or performance in downstream tasks such planning.
Given an observation encoded by the model and an action, we use the model to predict the rep- resentation of the next state, reached after taking the action in the environment. This predicted state representation then compared to the encoded true observation after taking the action in the environment and set of reference states (observations encoded by the model) obtained from the experience buffer. We measure and report both Hits at Rank 1 (H@) 1) and Mean Reciprocal Rank (MRR). Additional details these evaluation metrics can be found in Appendix C.
4.3 BASELINES
Autoencoder-based World Models The predominant method for state representation learning is based on autoencoders, and often on the VAE (Kingma & Welling, 2013; Rezende et al., 2014) model in particular. This World Model baseline inspired by Ha Schmidhuber (2018) and uses either deterministic autoencoder (AE) or VAE to learn state representations. Finally, MLP is used to predict the next state after taking an action.
Physics As Inverse Graphics (PAIG) This model by Jaques et al. (2019) is based on an encoder- decoder architecture and trained with pixel-based reconstruction losses, but uses differentiable physics engine in the latent space that operates on explicit position and velocity representations for each object. Thus, this model is only applicable to the 3-body physics environment.
4.4 TRAINING AND EVALUATION SETTING
We train C-SWMs on an experience buffer obtained by running a random policy on the respective environment. We choose 1000 episodes with 100 environment steps each for the grid world envi- ronments, 1000 episodes with 10 steps each for the Atari environments and 5000 episodes with 10 steps each for the 3-body physics environment.
For evaluation, we populate separate experience buffer with 10 environment steps per episode and total of 10.000 episodes for the grid world environments, 100 episodes for the Atari environments and 1000 episodes for the physics environment. For the Atari environments, we minimize train/test overlap by 'warm-starting' experience collection in these environments with random actions before we start populating the experience buffer (see Appendix B). and we ensure that not a single full test set episode coincides exactly with an episode from the training set. The state spaces of the grid world environments are large (approx. 6.4M unique states) and hence train/test coincidence of full 10-step episode is unlikely. Overlap is similarly unlikely for the physics environment which has continuous state space. Hence, performing well on these tasks will require some form of generalization new environment configurations an unseen sequence of states and actions.
All models are trained for 100 epochs (200 for Atari games) using the Adam (Kingma & Ba, 2014) optimizer with a learning rate of 5 10 and batch size of 1024 (512 for baselines with decoders due to higher memory demands, and 100 for PAIG as suggested by the authors). Model architecture details are provided in Appendix D.
4.5 QUALITATIVE RESULTS
We present qualitative results for the grid world environments in Figure 3 and for the 3-body physics environment in Figure 4. All results are obtained hold-out test data. Further qualitative results (inel. on Atari games) can be found in Appendix A.
In the grid world environments, ters for a particular scene, without direct supervision. Further, each object is represented by two coordinates which correspond (up random linear transformation) the true object position in
observe that C-SWM reliably discovers opject-specific fil-

---

Published as a conference paper at ICLR 2020
-2
-2.5
0.0
-2.5
0.0
(b) Learned abstract state transition graph of yel- cube (left) and the green square (right), while keeping all other ohject positions fixed at test time.
(a) Discovered object masks in scene from 3D cuhes (top) and 2DD shapes (hottom) environments.
Figure 3: Discovered object masks (left) and direct visualization the 2D abstract state spaces and transition graphs for single object (right) in the block pushing environments. Nodes denote state embeddings obtained from test set experience buffer with random actions and edges are predicted transitions. The learned abstract state graph clearly captures the underlying grid structure of the environment both terms of object-specific latent states and in terms of predicted transitions, but randomly rotated and/or mirrored. The model further correctly captures that certain actions do not have an effect if a neighboring position is blocked by another object (shown as colored spheres), though the transition model does not have access to visual inputs.
(a) Observations from 3-body gravitational physics simulation (hottom) and learned abstract state transition graph for a single object slot (top).
-2.5
0.0
2.5
(h) Abstract state transition graph from
50 test episodes for single object slot.
Figure 4: Qualitative results for 3-body physics environment for a single representative test set episode (left) and for dataset of 50 test episodes (right). The model learns smoothly embed object trajectories, with the circular motion represented the latent space (projected from four to two dimensions via PCA). In the abstract a particular episode, green links correspond to ground truth transitions and violet links correspond transitions predicted by the model. One trajectory (in the center) strongly deviates from typical trajectories seen during training, and the model struggles to predict the correct transition.
transition graph, orange nodes denote starting states for
the scene. Although we choose a two-dimensional latent representation per object for easier vi- sualization, find that results remain unchanged if we increase the dimensionality of the latent representation. The edges in this learned abstract transition graph correspond to the effect of a par- ticular action applied the object. The structure of the learned latent representation accurately captures the underlying grid structure of the environment. We further find that the transition model, which only has access latent representations, correctly captures whether an action has an effect or not, e.g., neighboring position is blocked by another object.
Similarly, we find that the model can learn object-specific encoders in the 3-body physics envi- ronment and can learn object-specific latent representations that track location and velocity of particular object, while learning an accurate latent transition model that generalizes well to unseen environment instances.
4.6 QUANTITATIVE RESULTS
We set up quantitative experiments for evaluating the quality of both object discovery and the quality of the learned transition model. We compare against autoencoder baselines and model variants that do not represent the environment in object-factorized manner, do not GNN, do not make use of contrastive learning. Performing well under this evaluation setting requires some degree of (combinatorial) generalization unseen environment instances.

---

Published as a conference paper at ICLR 2020
We report ranking scores (in %) in latent space, after encoding source and target observations, and taking steps in the latent space using the learned model. Reported results are mean and standard error of scores over 4 runs hold-out environment instances. Results summarized in Table 1.
Table 1: Ranking results for sult prediction in latent space. Highest (mean) scores in bold.
Step
5 Steps
10 Steps
Model
H@1
MRR
H@1
MRR
H@)
MRR
C-SWM
100:00
100
100
100
99.9
100
2D SHAPES
latent GNN
99,9=00
100:00
97.4=a0
98,4
89.7=03
93.1
factored states
54.5
65.0
34.41166
47.4
24.1
37.0
contrastive loss
49.9=a9
55.2
6,5
9.3
1.4ga.
2.6
World Model (AE)
98.7=as
99.2
36. =811
44. =881
6,5
10,5
World Model (VAE)
94.2
96.4
14.1
21.4
1.4
3.5
3D BLOCKS
C-SWM
99.9.mm
100g
100g00
99.9.n0
99.9.no
latent GNN
99.9 -9.9
99.9
96.3 H044
97.7+a0
86.0#1a
90,2+115
factored states
74.2=93
82,5
48.7+122
62.6+13.0
65.8=140
49.6+110
contrastive loss
48.9
52.5
12.2
16.3
3.1
5.3
World Model (AE)
93.5
95.6206
26.7 197
35.6108
4.0
7.6
World Model (VAE)
90.9=a7
94.2#;06
31.3
41.8
7.2
12.9
C-SWM (K 5)
20.5
41.8
9,5
22.2.33
5.3
15,8
ATARI
PONG
C-SWM (K-
34.8
54.3
12.8
28.1
9.5
21.1
C-SWM (K = 1)
36.5
56.2262
18.3
35.7
11.5
26.0
World Model (AE)
23.8 =3.
44.7
1.7 +0.5
8,0
1,2
5,3
World Model (VAE)
1.0
5. +o.1
1.0
5,2+00
1.0
5,2+00
SPACE INVADERS
C-SWM (K- 5)
48.5
66.1
16.8 #27
35.7
11.8#30
26.0
C-SWM (K- 3)
46.2+13.0
10.8
28.5
6.0+04
20.9
C-SWM (K- = 1)
31.5
48.61118
10.0
23.9
6,0
19.8
World Model (AE)
40.2
59.6
5.2
14.1
3,8
10.4
World Model (VAE)
1.0g00
5.3
0.8
5.2
1.0
5.2+00
3-BODY PHYSICS
C-SWM
100:00
100
97.2.09
98.5
75.5
85.2
World Model (AF)
100:00
100..0
97.7
98.8
67.9
78.4
World Model (VAE)
100:00
100.00
83.1
90.3
23.6242
37.5
Physics WM (PAIG)
89,2-355
90.7
57.7+120
63.1
25.1 +13.0
33.1
We find that baselines that make use of reconstruction losses in pixel space (incl. the C-SWM model variant without contrastive loss) typically generalize less well to unseen scenes and learn latent space configuration that makes it difficult for the transition model to learn the correct transition function. See Appendix A.1 for a visual analysis of such latent state transition graphs. This effect appears to be even stronger when using VAE-based World Model, where the prior puts further constraints on the latent representations. C-SWM recovers this structure well, see Figure 3.
On the grid-world environments (2D shapes and 3D blocks), C-SWM models latent transitions al- most perfectly, which requires taking interactions between latent representations of objects into account. Removing the interaction component, i.e., replacing the latent GNN with an object-wise MLP, makes the model insensitive to pairwise interactions and hence the ability to predict future states deteriorates. Similarly, we remove the state factorization, the model has difficulties general- izing to unseen environment configurations. We further explore variants of the grid-world environ- ments in Appendix A.S, where 1) certain actions have no effect, and 2) one object moves randomly and independent of agent actions, to test the robustness of our approach.
For the Atari 2600 experiments, we find that results can have high variance, and that the task is more difficult, as both the World Model baseline and C-SWM struggle to make perfect long- term predictions. While for Space Invaders, large number of object slots (K = 5) appears to be beneficial, C-SWM achieves best results with only single object slot in Atari Pong. This suggests that one should determine the optimal value of K based on performance on a validation set if it is not known a- a-priori Using an iterative object encoding mechanism, such as in MONet (Burgess et al., 2019), would enable the model to assign 'empty" slots which could improve robustness w.r.. the choice of K, which we leave for future work.
8

---

Published as a conference paper at ICLR 2020
We find that both C-SWMs and the sutoenccoder-aseee World Model baseline excel at short-term predictions in the 3-body physics environment, with C-SWM having slight edge in the 10 step pre- diction setting. Under evaluation setting, the PAIG baseline (Jaques al., 2019) underperforms using the hyperparameter setting recommended by the authors. Note that we do not tune hyperpa- rameters of C-SWM separately for this task and use the same settings as in other environments.
4.7 LIMITATIONS
Instance Disambiguation In our experiments, we chose a simple feed-forward CNN architecture for the object extractor module. This type of architecture cannot disambiguate multiple instances of the same object present in one scene and relies on distinct visual features or labels (e.g., the green square) for object extraction. To better handle scenes which contain potentially multiple copies of the same object (c.g., in the Atari Space Invaders game), one would require some form of iterative disambiguation procedure break symmetries and dynamically bind individual objects to slots or object files (Kahneman & Treisman, 1984; Kahneman et al., 1992), such as in the style of dynamic routing (Sabour et al., 2017), iterative inference (Greff et al., 2019; Engelcke et al., 2019) or sequential masking (Burgess et al., 2019; Kipf et al., 2019).
Stochasticity Markov Assumption Our formulation of C-SWMs does not take into account stochasticity in environment transitions or observations, and hence is limited to fully deterministic worlds. A probabilistic extension of C-SWMs an interesting avenue for future work. For sim- plicity, we make the Markov assumption: state and action contain all the information necessary to predict the next state. This allows us to look at single state-action-state triples in isolation. go beyond this limitation, would require some form of memory mechanism, such RNN as part of the model architecture, which we leave for future work.
CONCLUSIONS
Structured world models offer compelling advantages over pure connectionist methods, by enabling stronger inductive biases for generalization, without necessarily constraining the generality of the model: for example, the contrastively trained model the 3-body physics environment is free to store identical representations in each object slot and ignore pairwise interactions, i.e., an unstruc- tured world model still exists a special case. Experimentally. use of this additional structure, likely because it allows for transition model of significantly lower complexity, and learn object-oriented models that generalize better unseen situations.
find that C-SWMs make effective
We are excited about the prospect of using C-SWMs for model-based planning and reinforcement learning in future work, where object-oriented representations will likely allow for more accurate counterfactual reasoning about effects of actions and novel interactions in the environment. We further hope to inspire future work to think beyond autoencoder based approaches for object-based, structured representation learning, and to address some of the limitations outlined in this paper.
ACKNOWLEDGEMENTS
We would like to thank Marco Federici and Adam Kosiorek for helpful discussions. We would further like to thank the anonymous reviewers for valuable feedback. T.K. acknowledges funding by SAP SE.
REFERENCES
Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Marc-Alexandre Côté, and R Devon Hjelm. Unsupervised state representation learning in Atari. arXiv preprint S.KV:1906.082222 2019.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.066 2016.
Peter Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, et al. Interaction networks for learning about objects, relations and physics. In NIPS, 2016.

---

Published as a conference paper at ICLR 2020
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzaler Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. networks. Relational inductive biases, deep learning, and networs.. arXiv preprint arXiv:18060011 01261, 2018.
Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning envi- ronment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 2013.
Antoine Bordes, Nicolas Usumier, Alberto Garcia-Duran, Jason Weston, and Oksana Yakhnenko. Translating embeddings for modeling multi-relational data. NIPS, 2013.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint arXiv:1606.01544 2016.
Christopher P Burgess, Loic Matthey, Nicholas Watters, Rishabh Kabra, Trina Higgins, Matt Botvinick, and Alexander Lerchner. Monet: Unsupervised scene decomposition and represen- tation. arXiv preprint
2019.
Michael B Chang, Tomer Ullman, Antonio Torralba, and Joshua B Tenenbaum. A compositional object-based approach learning physical dynamics. arXiv preprint arXiv: 2016.
Dane Corneil, Wulfram Gerstner, and Johanni Brea. Efficient model-based deep reinforcement learning with variational state tabulation. arXiv preprint arXiv: 04325, 2018.
Sebastien Ehrhardt, Aron Monszpart, Niloy Mitra, and Andrea Vedaldi. Unsupervised intuitive physics from visual observations. In Asian Conference on Computer Vision. Springer, 2018.
Martin Engelcke, Adam R Kosiorek, Oiwi Parker Jones, and Ingmar Posner. Genesis: Gener- ative scene inference and sampling with object-centric latent representations. arXiv preprint arXiv:1907.1305,, 2019.
SM Ali Eslami, Nicolas Heess, Theophane Weber, Yuval Tassa, David Szepesvari, Koray Kavukcuoglu, and Geoffrey E Hinton. Attend, infer, repeat: Fast scene understanding with gen- erative models. NIPS, 2016.
Vincent François-Lavet, Yoshua Bengio, Doina Precup, and Joelle Pineau. Combined reinforcement learning via abstract representations. arXiv preprint
2018.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare G Bellemare. Deep- MDP: Learning continuous latent space models for representation learning. arXiv preprint 19X6::9906.02/11 2019.
Justin Gilmer, Samuel Schoenholz, Patrick Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In ICML, 2017.
Klaus Greff, Sjoerd van Steenkiste, and Jürgen Schmidhuber, Neural expectation maximization. In NIPS. 2017.
Klaus Greff, Raphaël Lopez Kaufmann, Rishab Kabra, Nick Watters, Chris Burgess, Daniel Zoran, Loic Matthey, Matthew Botvinick, and Alexander Lerchner. Multi-object representation learning with iterative variational inference. arXiv preprint /903v 004500 2019.
Aditya Grover and Jure Leskovec. node2vee: Scalable feature learning for networks. In KDD, 2016.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv:1803.10122. 2018.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In ICML, 2019.
Olivier J Hénaff, Ali Razavi, Carl Doersch, SM Eslami, and Aaron van den Oord. Dfta-efficient image recognition with contrastive predictive coding. arXiv preprint arXiv: 1905.0922, 2019.
10

---

Published as a conference paper at ICLR 2020
R Devon Hjelm, Alex Fedorov, Samuel Yoshua Bengio. Learning deep representations by mutual information estimation and maximiza- tion. arXiv preprint (XXX::8808.066777 2018.
Karan Grewal, Adam Trischler, and
Yedid Hoshen. Vain: Attentional multi-agent predictive modeling. In NIPS, 2017.
John D Hunter. 2007.
Matplotlib: A 2d graphics environment. Computing In Science Engineering,
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint 100. 2015.
Michael Janner, Sergey Levine, William Freeman, Joshua Tenenbaum, Chelsea Finn, and Jiajun Wu. Reasoning about physical interactions with object-oriented prediction and planning. In ICLR, 2019.
Miguel Jaques, Michael Burke, and Timothy Hospedales. bessio--s-inverse-graahiic Joint unsu- pervised learning of objects and physics from video. arXiv preprint 19Xiv:190..1119,, 2019.
Rico Jonschkowski and Oliver Brock. Learning state representations with robotic priors. Au- tonomous Robots, 2015.
Daniel Kahneman and Anne Treisman. Changing views of Attention and Automaticity. Academic Press, Inc.. San Diego, CA, 1984.
Daniel Kahneman, Anne Treisman, and Brian J Gibbs. The reviewing of object files: Object-specific integration of information. Cognitive psychology, 1992.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, al. Model-based reinforcement learning for atari. arXiv preprint 19Xxvv:9033.00 2019.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1rXiv:1412.6988 2014.
Diederik Kingma and Max Welling. arXiv:1312.6114, 2013.
Auto-encoding variational bayes. arXiv preprint
Thomas Kipf, Ethan Fetaya, Kuan-Chiel Wang. Max Welling, and Richard Zemel. Neural relational inference for interacting systems. In ICML, 2018.
Thomas Kipf, Yujia Li, Hanjun Dai, Vinicius Zambaldi, Alvaro Sanchez- -Gobalzzz, Edward Grefen- stette, Pushmeet Kohli, and Peter Battaglia. Compile: Compositional imitation learning and execution. In /CML, 2019.
Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional net- works. arXiv preprint arXiv: 1609002099 2016.
Adam Kosiorek, Hyunjik Kim, Yee Whye Teh, and Ingmar Posner. Sequential attend, infer, repeat: Generative modelling of moving objects. In NeurlPS, 2018.
Thanard Kurutach, Aviv Tamar, Ge Yang, Stuart J Russell, and Pieter Abbeel. Learning plannable representations with causal infogan. In NeurIPS, 2018.
Adrien Laversanne-Fines Alexandre Péré, and Pierre Yves Oudeyer. Curiosity driven exploration of learned disentangled goal spaces. arXiv preprint arXiv: 5221, 2018.
Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Huang. tutorial on energy-based learning. Predicting structured data, 2006.
Yujia Li. Daniel Tarlow, Mare Brockschmidt, and Richard Zemel. Gated graph sequence neural networks. arXiy preprint 1rX5v.157..05493, 2015.
Long-Ji Lin. Self-improving reactive agents based on reinforcement learning, planning and teaching. Machine learning, 8(3-4):293-321, 1992.

---

Published as a conference paper at ICLR 2020
Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation of word represen- tations in vector space. arXiv preprint arXiv:1301.378. 2013.
Andriy Mnih and Yee Whye Teh. A fast and simple algorithm for training neural probabilistic language models. arXiv preprint arXiv:
2012.
Charlie Nash, Ali Eslami, Chris Burgess, Irina Higgins, Daniel Zoran, Theophane Weber, and Peter Battaglia. The multi-entity variational autoencoder. In NIPS Workshops, 2017.
Maximilian Nickel, Volker Tresp. and Hans-Peter Kriegel. A three-way model for collective learning on multi-relational data. In ICMI., 2011.
Maximilian Nickel, Kevin Murphy, Volker Tresp, and Evgeniy Gabrilovich. A review of relational machine learning for knowledge graphs. Proceedings of the IEEE, 2016a.
Maximilian Nickel, Lorenzo Rosasco, and Tomaso Poggio. Holographic embeddings of knowledge graphs. In AAAI, 2016b.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. arXiv preprint 107.v:/807.03745. 2018.
Bryan Perozzi, Rami AI-Rfou, and Steven Skiena. Deepwalk: Online learning of social representa- tions. In 2014.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference deep generative models. arXiv preprint 4rXhv:1401.4082, 2014.
Sara Sabour, Nicholas Frosst, and Geoffrey E Hinton. Dynamic routing between capsules. NIPS, 2017.
Alvaro Sanchez-Gonzalez Nicolas Heess, Jost Tobias Springenberg, Josh Merel, Martin Riedmiller, Raia Hadsell, and Peter Battaglia. Graph networks as learnable physics engines for inference and control. In ICML, 2018.
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEEE Transactions on Neural Networks, 2009.
Michael Schlichtkrull, Thomas N Kipf, Peter Bloem, Rianne Van Den Berg, Ivan Titov, and Max Welling, Modeling relational data with graph convolutional networks. In ESWC, 2018.
Elizabeth S Spelke and Katherine Kinzler. Core knowledge. Developmental science, 2007.
Sainbayar Sukhbaatar, Rob Fergus, et al. Learning multiagent communication with backpropaga- tion. In NIP,, 2016.
Chen Sun, Abhinav Shrivastava, Carl Vondrick, Kevin Murphy, Rahul Sukthankar, and Cordelia Schmid, Actor-centric relation network. In ECCV, 2018
Chen Sun, Fabien Baradel, Kevin Murphy, and Cordelia Schmid. Contrastive bidirectional trans- former for temporal representation learning. arXiv preprint arXiv: 190,:0940605,33,, 2019a.
Chen Sun, Abhinav Shrivastava, Carl Vondrick, Rahul Sukthankar, Kevin Murphy, and Cordelia Schmid. Relational action forecasting, arXiv preprint arXiv:1904.04233 2019b.
Valentin Thomas, Emmanuel Bengio, William Fedus, Jules Pondard. Philippe Beaudoin, Hugo Larochelle, Joelle Pineau, Doina Precup, and Yoshua Bengio. Disentangling the independently controllable factors of variation by interacting with the world, arXiv preprint arXiv: arXiv:/802. 100.94480,,09444 2018.
Théo Trouillon, Johannes Welbl, Sebastian Riedel, fiii Gaussier, and Guillaume Bouchard. Com- plex embeddings for simple link prediction. In JCML, 2016.
Sjoerd van Steenkiste, Michael Chang. Klaus Greff, and Jürgen Schmidhuber. Relational neural ex- pectation maximization: Unsupervised discovery of objects and their interactions. arXiv preprint arXiv: 1802. 10353, 2018.
12

---

Published as a conference paper at ICLR 2020
Petar Veličkovié, William Fedus, William L Hamilton, Pietro Liö, Yoshua Bengio, and R Devon Hjelm. Deep graph infomax. arXiv preprint 1rXiv:1809.10341. 2018.
Angelina Wang, Thanard Kurutach, Kara Liu, Pieter Abbeel, and Aviv Tamar. Learning robotic manipulation through visual planning and acting. arXiv preprint (e05::990004411 2019.
Tingwu Wang, Renjie Liao, Jimmy Ba, and Sanja Fidler. Nervenet: Learning structured policy with graph neural networks. In ICLR, 2018.
Zhen Wang. Jianwen Zhang, Jianlin Feng, and Zheng Chen. Knowledge graph embedding by trans- lating on hyperplanes. In AAAI, 2014.
Manuel Watter, Jost Springenberg. Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. In NIPS, 2015.
Nicholas Watters, Daniel Zoran, Theophane Weber, Peter Battaglia, Razvan Pascanu, and Andrea Tacchetti. Visual interaction networks: Learning physics simulator from video. In NIPS, 2017.
Nicholas Watters, Loic Matthey, Matko Bosnjak, Christopher P Burgess, and Alexander Lerchner. Cobra: Data-efficient model-based rl through unsupervised object discovery and curiosity-driven exploration. arXiv preprint 19Xiv:1905.09275 2019.
Bing Xu, Naiyan Wang, Tianqi Chen, and Mu Li. Empirical evaluation of rectified activations in convolutional network. arXiv preprint arXiv:1505.00853 2015.
Zhenjia Xu, Zhijian Liu, Chen Sun, Kevin Murphy, William T Freeman, Joshua B Tenenbaum, and Jiajun Wu. Unsupervised discovery of parts, structure, and dynamics. arXiv preprint arXiv:/903.05/:00 1366 2019.
13

---

Published as a conference paper at ICLR 2020
A
ADDITIONAL RESULTS AND DISCUSSION
A.1 OBJECT-SPECIEG REPRESENTATIONS
We visualize abstract state transition graphs separated by object slot for the 3D cubes enviromment Figure 5. Discovered object representations very similar. We apply the same visualization technique to the model variant without contrastive loss. which is instead trained with a decoder model and a loss in pixel space. See Figure for this baseline and note that the regular structure in the latent space is lost, which makes it difficult for the transition model to learn transitions which generalize to unseen environment instances.
2D shapes dataset (not shown) are qualitatively
Qualitative results for the 3-body physics dataset are summarized in Figures and 8 for two different random seeds.
0-
-2
-2
-2
-2
(a) Object slot
(b) Object slot 2. (c) Object slot 3. (d) Object slot 4.
(e) Object slot 5.
Figure 5: Abstract state transition graphs per object slot for trained C-SWM model the 3D cubes environment (with all objects allowed to be moved, i.c., none are fixed in place). Edge color denotes action type. The abstract state graph is nearly identical for each object, which illustrates that the model successfully represents objects in the same manner despite their visual differences.
-5
10
(a) Object slot
(b) Object slot 2.
(c) Object slot 3. (d) Object slot 4.
(e) Object slot 5.
Figure 6: Abstract state transition graphs per object slot for a trained SWM model without trastive loss, using instead loss in pixel space, on the 3D cubes environment. Edge color denotes action type.
2.5
0.0
2.5
-2.5
0.0
-2.5
0.0
2.5
(a) Discovered object-specific filters.
(b) Object slot 1.
(c) Object slot 2.
(d) Object slot 3.
Figure 7: Object filters (left) and abstract state transition graphs per object slot (right) for a trained C-SWM model on unseen test instances of the 3-body physies environment (seed) 1).
For the Atari 2600 environments, we generally found latent object representations to be less in- terpretable. We attribute this to the fact that a) objects have different roles and are general not exchangeable (in contrast the block pushing grid world environments and the 3-body physics environment), b) actions affect only one object directly, but many other objects indirectly in two consecutive frames, and c) due multiple objects in scene sharing the same visual features. See Figure 9 for an example of learned representations in Atari Pong and Figure 10 for an example in Space Invaders.
14

---

Published as a conference paper at ICLR 2020
(a) Discovered object-specific filters.
(b) Object slot 1.
(c) Object slot
(d) Object slot
Figure 8: Object filters (left) and abstract state transition graphs per object slot (right) for a trained C-SWM model on unseen test instances of the 3-body physies environment (seed 2).
2.5
(a) Discovered object-speeific filters.
(b) Object slot 1.
(c) Object slot
(d) Object slot
Figure 9: Object filters (left) and abstract state transition graphs per object slot (right) for trained C-SWM model with K object slots on unseen test instances of the Atari Pong environment.
1.0
0.0
0.5
0.03
(a) Discovered spjecc-specific filters.
(b) Object slot 1.
(c) Object slot
(d) Object slot
Figure 10: Object filters (left) and abstract state transition graphs per object slot (right) for a trained C-SWM model with K object slots on unseen test instances of the Space Invaders environment.
7
10
Ground
C-SWM
World Model (AE)
Figure 11: Qualitative model comparison pixel space on a hold-out test instance of the 2D shapes environment. We train separate decoder model for 100 epochs on both the C-SWM and the World Model baseline using all training environment instances to obtain pixel-based reconstructions for multiple prediction steps into the future.
15

---

Published as a conference paper at ICLR 2020
A.2 MODEL COMPARISON IN PIXEL SPACE
To supplement our model comparison in latent space using ranking metrics, we here show direct comparison in pixel space at the example of the 2D shapes environment. This requires training separate decoder model for C-SWM. For faireess of this comparison, we use the same protocol to train a separate decoder for the World Model (AE) baseline (discarding the one obtained from the original end-to-end auto-encoding training procedure). This decoder has the same architecture as in the other baseline models and trained for 100 epochs.
For a qualitative comparison, see Figure 11. The C-SWM model, as expected from our ranking anal- ysis in latent space, performs almost perfectly at this task. Although
Pixel-based MSE
10
World Model (AE) baseline makes clear mistakes which compound over time, it nonetheless often gets several object positions cor-
103
102
rect after many time steps. The ranking loss in latent space captures this behaviour well, and, for exam- ple, assigns almost perfect score for 1-step pre- diction to the World Model (AE) baseline. The typi- cally used mean-squared error (MSE) in pixel space (see Figure 12), however, differs by several orders of magnitude between the two models and does not capture any of the nuanced differences in qualitative predictive behavior. We hence strongly encourage researchers in this area to consider ranking-based evaluations directly in latent space in addition to, or as replacement for comparisons in pixel space.
C-SWM
10a
World Model (AE)
2
3
5
8
10
Prediction steps
Figure 12: Quantitative model comparison in pixel space on hold-out test set of the 2D shapes environment. The plot shows mean squared reconstruction error (MSE) in pixel space for multiple transition model predic- tion steps into the future (lower is better), averaged over 4 runs. Shaded area denotes standard error.
A.3 HINGE LOSS
In Table 2, we summarize results a comparison between our multi-object contrastive loss from Eq. 6, denoted by C-SWM (original), and full-hinge (or triplet) loss that places the hinge both the positive and the negative energy term, i.e., L max(0, H Hf). The latter loss is similar to the loss employed in Bordes et al. (2013), but without imposing a norm constraint on the embeddings 24. Reported results are mean and standard error over 4 runs on hold-out environment instances.
Table 2: Comparison of hinge loss variants on 2D shapes environment. Numbers denote ranking scores (in %) for multi-step prediction latent space. Highest (mean) scores bold.
Step
5 Steps
10 Steps
Model
H@1
MRR
H@1
MRR
H@ MRR
C-SWM (original)
100
100
100
100
99.9
100
with full hinge (7 -1)
95.9+11
97.8
40.4
54.4
12.6
21.8
with full hinge (7 = 5)
97.7+
98.8 1.0.5
52.0
65.7 1:66
15.8
26.8 #3.5
with full hinge (7 10)
97.0...4
98.5 4:02
49.5
63.3
18.7
30.3 214
We find that placing the hinge only on the negative energy term significantly improves results, likely because the full hinge loss can be minimized by simply increasing the norms of the embeddings 21, as outlined in Bordes al. (2013). This could be addressed by introducing additional constraints, such as by forcing the embeddings to lie on the unit-hypersphere, al., 2013), which we not consider this work most of the environment dynamics in our experiments obey flat Euclidean geometry (e.g.. translation on grid).
done in TransE (Bordes
A.4 MULTIPLE FEATURE MAPS PER OBJECT SLOTT
We experimented with assigning multiple feature maps per object slot to investigate whether this can improve performance the Space Invaders environment (as example of one of the more complex
16

---

Published as a conference paper at ICLR 2020
environments considered in this work). All other model settings are left unchanged. Results summarized in Table 3. The original model variant with one feature map per object slot (using a total of K 5 slots) is denoted by C-SWM (original). Reported results mean and standard error over runs on hold-out environment instances.
Table Comparison of model variants with multiple feature maps per object on the Space Invaders environment. Numbers denote ranking scores (in %) for multi-step prediction in latent space.
Step
5 Steps
10 Steps
Model
H@L 1
MRR
H@!
MRR
H@)
MRR
C-SWM (original)
48.5
66. 166
16.8
35.7
11.8=30
26.0
with feature maps
48.2
65.7
12.8
27.7=32
8.2
21.2
with feature maps
52.5..34
71.022.
18.2
36.3
11.8
25.0
with 10 feature maps
50.010s
68.6
15.8
32.4123
9.0
22.7
We find that there is no clear advantage in using multiple feature maps per object slot for the Space Invaders environment, but this might nonetheless prove useful for environments of more complex nature.
A.5 GRID-WORLD ENVIRONMENT VARIANTS
To further evaluate the robustness of our approach, we experimented with two variants of the 2D shapes grid-world environment with five objects. the first variant, termed no-op action, we include an additional 'no-op' action for each object that has effect. This mirrors the setting present in many Atari 2600 benchmark environments that have an explicit no-op action. In the second variant, termed random object, we assign special property to out of the five objects: random action is applicd to this object at every turn, i.c., the agent only has control over the other four objects and in every environment step, total of two objects receive action (one by the agent, which only affects one out the first four objects, and one random action on the last object). This is to test how the model behaves in the presence of other agents that act randomly. As our model is fully deterministic, it is to be expected that adding a source of randomness in the environment will pose a challenge.
We report results in Table 4. Reported numbers are mean and standard environment instances. We find that adding a no-op action has little effect on our model, but adding randomly moving object reduces predictive performance. Performance in the latter case could potentially be improved by explicitly modeling stochasticity in the environment.
runs on hold-out
Table 4: Comparison of C-SWM on variants of the 2D shapes environment. Numbers denote ranking
scores (in %) for multi-step prediction in latent space.
Step
5 Steps
10 Steps
Model
H@1
MRR
H@1
MRR
H@1
MRR
C-SWM (original env.)
100
100=00
100
100
99.9
100.00
+ no-op action
100 a0.
100-oa
100
100
99.9
99.9
random object
98.4
99.2
95.9
97.6
90.4
93.7
A.6 STABILITY
The discovered object representations and identifications can vary between runs. While we found that this process very stable for the simple grid world environments where actions only affect single object, we found that results can be initialization-deenndea actions can have effects across all objects, and the 3-body physics simulation (see Figures and 8), which does not have any actions. In some cases, the discovered representation be less suitable for forward prediction in time or for generalization to novel scenes, which explains the variance in some of our results these datasets.
the Atari environments, where
17

---

Published as a conference paper at ICLR 2020
A.7 TRAINING TIME
We found that the overall training time of the C-SWM model was comparable to that of the World Model baseline. Both C-SWM and the World Model baseline trained for approx. hour wall-clock time the 2D shapes dataset, approx. hours on the 3D cubes dataset, and approx. 30min the 3-body physics environment using a single NVIDIA GTX10801 GPU. The Atari Pong and Space Invaders models trained for typically less than 20 minutes. A notable exception is the PAIG baseline model (Jaques et al., 2019) which trained for approx. 6 hours on a NVIDIA TitanX Pascal GPU using the recommended settings by the authors of the paper.
B DATASETS
B. GRID WORLDS
To generate an experience buffer for training, we initialize the environment with random object placements and uniformly sample an object and an object-specific action at every time step,
We provide state observations as 50 x 50 x 3 tensors with RGB color channels, normalized to [o, 1]. Actions are provided as a 4-dim one-hot vector (if an action is applied) or a vector of zeros per object slot in the environment. The action one-hot vector encodes the directional movement action applied to particular object, or is represented as vector of zeros if action is applied to particular object. Note that only a single object receives an action per time step. For the Atari environments, provide a copy of the one-hot encoded action vector to every object slot, and for the 3-body physics environment, which has no actions, we do not provide an action vector.
2D Shapes This environment is 5 5 grid world with 5 different objects placed at random positions. Each location can only be occupied by at maximum one object. Each object is represented by unique shape/color combination, occupying 10 x 10 pixels the overall 50 50 pixel grid. At each time step, one object can be selected and moved by one position along the four cardinal directions. See Figure 2a for an example. The action has effect if the target location in a particular direction is occupied by another object or outside of the 5 x grid. Thus, a learned transition model needs to take pairwise interactions between object properties (i.e., their locations) into account, would otherwise be unable to predict the effect of an action correctly.
To investigate to what degree our model is robust to partial occlusion and perspective changes, we implement a simple block pushing environment using Matplotlib (Hunter, 2007) as rendering tool. The underlying environment dynamics are the same as in the 2D Shapes dataset, and we only change the rendering component to make for visually more challenging task that involves a different perspective and partial occlusions. See Figure 2b for an example.
3D Blocks
B.2 ATARI 2600 GAMES
Atari Pong We make use of the Arcade Learning Environment (Bellemare et al., 2013) create a small environment based on the Atari 2600 game Pong which is restricted to the first interaction between the ball and the player-controlled paddle, i.e.. we discard the first couple of frames from the experience buffer where the opponent behaves completely independent of any player action. Specifi- cally, we discard the first 58 (random) environment interactions. We use the PONGIDETERMINISTICC v4 variant of the environment in OpenAl Gym (Brockman et al., 2016). We use a random policy and populate an experience buffer with 10 environment interactions per episode, i.c., T = 10. An observation consists of two consecutive frames, cropped (to exclude the score) and resized to 50 x 50 pixels each. See Figure 2c for an example.
Space Invaders This environment is based on the Atari 2600 game Space Invaders, using the OPACEINADDESDEETEEMISTTCCC..
variant of the environment in OpenAl Gym (Brockman et al., 2016), and processed / restricted in similar manner as the Pong environment. We dis- card the first 50 (random) environment interactions for each episode and only begin populating the experience buffer thereafter. See Figure 2d for example observation.
18

---

Published as a conference paper at ICLR 2020
B.3 3-BODY PHYSICS
The 3-body physics simulation environment is an interacting system that evolves according to clas- sical gravitational dynamics. Different from the other environments considered here, there are no actions. This environment is adapted from Jaques al. (2019) using their publicly available imple- mentation', where we set the step size (dt) to 2.0 and the initial maximum and velocities to 0.5. We concatenate two consecutive frames of x 50 pixels each provide the model with (implicit) velocity information. See Figure 4 for example observation sequence.
C EVALUATION METRICS
C.I HITS AT RANK K (H@K)
This score for a particular example if the predicted state representation is in the k-nearest neigh- bor set around the encoded true observation, where define the neighborhood of a node to include the node itself. Otherwise this score is O. In other words, this score measures whether the rank of the predicted state representation is smaller than equal to k, when ranking all reference state representations by distance to the true state representation. We report the average this score over particular evaluation dataset.
C.2 MEAN RECIPROCAL RANK (MRR)
This score is defined as the average inverse rank, i.e., MRR = 1 Σm rank,,' where rank, the rank of the n-th sample.
D ARCHITECTURE AND HYPERPARAMETERS
D.1 OBJECT EXTRACTOR
For the 3D cubes environment, the object extractor is 4-layer CNN with 3 x 3 filters, zero-padding, and 16 feature maps per layer, with the exception of the last layer, which has ie., one per object slot. After each layer, apply BatchNorm (Ioffe & Szegedy, 2015) and ReLU(z)= = max(0, x) activation function. For the 2D shapes environment, we choose a simpler CNN architecture with only single convolutional layer with 10 x 10 filters and stride of 10, followed by BatchNorm and a ReLU activation or LeakyReLU (Xu et al., 2015) for the Atari 2600 and physics environments. This layer has 16 feature maps and is followed by channel-wise linear transformation (i.e., 1 X 1 convolution). with feature maps as output. For both models, we choose a sigmoid activation function after the last layer obtain object masks with values in (0, 1). We use the same two-layer architecture for the Atari 2600 environments and the 3-body physics environment, but with 9 filters (and zero-pate paro-padding) in the first layer, and 5 5 filters with a stride of 5 in the second layer.
feature maps,
D.2 OBJECT ENCODER
After reshaping/flattennng the output of the object extractor, we obtain a vector representation per object (2500-dim for the 3D eubes environment, 25-dim for the 2D shapes environment, and 1000- dim for Atari 2600 and physics environments). The object encoder is an MLP with two hidden layers of 512 units and each, followed by ReLU activation functions. We further use LayerNorm (Ba et al., 2016) before the activation function of the second hidden layer. The output of the final output layer is 2-dimensional (4-dimensional for Atari 2600 and physies environments), reflecting the ground truth object state, i.e., the object coordinates in 2D (although this is not provided to the model), and velocity (if applicable).
hhttps://gii thub. http://ggthhubb.ooom/uu////a 14/ppp
19

---

Published as a conference paper at ICLR 2020
D.3 TRANSITION MODEL
Both the node and the edge model in the GNN-based transition model are MLPs with the same architecture number of hidden units as the object encoder model, i.., two hidden layers of 512 units each, LayerNorm and ReLU activations.
D.4 LOSS FUNCTION
The margin in the hinge loss is chosen as 7 1. We further multiply the squared Euclidean distance d(z, y) in the loss function with a factor of 0.5/o² with a 0.5 control the spread of the embeddings. We use the same setting in all experiments.
D.5 BASELINES
World Model Baseline The World Model baseline is trained in two stages: First, we train an auto- encoder or VAE with a 32-dim latent space, where the encoder a CNN with the same architecture as the object extractor used in the C-SWM model, followed by an MLP with the same architecture as the object encoder module in C-SWM on the flattened representation of the output the encoder CNN. The decoder exactly mirrors this architecture where we replace convolutional layers with deconvolutional layers. We verified that this architecture can successfully build representations of single frames.
Example reconstructions from the latent code are shown in Figure 13. We experimented both with mean squared error and binary cross entropy (using the continuous channel values in [0, 1] targets) reconstruction loss in the (V)AE models, both of which are typical choices in most practical implementations. We generally found binary cross entropy to be more stable to optimize and produce better results, which why we opted for this loss in all the baselines using decoders considered in the experimental section.
(a) 3D Cubes
(b) 2D Shapes
Figure 13: Reconstructions from the la- tent code of trained VAE-based World Model baseline.
In the second stage, we freeze the model parameters of the auto-encoder and train transition model with mean- mean-squared error on the latent representations. For the VAE model, we use the predicted mean values of the latent representations. This transition model takes the form of an MLP with the same architecture and number of hidden units as the node model in C-SWM. We experimented both with translational transition model (i.e., the transition model only predicts the latent state difference. instead of the full next state) and direct prediction of the next state. We generally found that the translational transition model performed better and used throughout all the reported results.
The World Model baselines trained with smaller batch size of 512, which slightly improved performance and simplified memory management.
Ablations We perform the following ablations: 1) we replace the latent GNN with an MLP (per object, i.e., we remove the edge update function) to investigate whether structured transition model necessary, 2) we remove the state factorization and embed the full scene into a single latent vari- able of higher dimensionality (original dimensionality number of original object slots) with an MLP as transition model, and 3) we replace the contrastive loss with pixel- based reconstruction loss on both the current state and the predicted next state (we add decoder that mirrors the archi- tecture of the encoder).
Physice-as-Invers--Graappp (PAIG) For this baseline, we train the PAIG model from Jaques et al. (2019) with the code provided by the authors? settings recommended by the authors for this particular task, namely: model=PhysiesNet, epochs=500, batch size=100, base Ir=le-3, autoencoder loss=5.0 00, anneal Ir=true color=true, and cell_type-gravity. _ele_cell. We use the same input size as C-SWM, i.e., frames are of shape
dataset with the standard
Phttps://0g wwt.:/cc///eeaaaa 11/ppai
20

---

Published as a conference paper at ICLR 2020
50 50 x 3. We further set input_steps=2, sted_ste and extrap-steps=0 of predicting for total of 10 frames while conditioning on pair initial frames to obtain ini- tial position and velocity information. We train this model with four different random seeds. For evaluation, we extract learned position representations for all frames of the test set and further run the latent physics prediction module (conditioned the first two initial frames) to obtain model predictions. We further augment position representations with velocity information by taking the difference between two consecutive position representations and concatenating this representation with the 2-dim position representation, which we found to slightly improve results. In total, we obtain a 12-dim (4-dim objects) representation for each time step, i.e., the same as C-SWM. From this, we obtain ranking metrics. We found that one of the training runs (seed=2) collapsed all latent representations to single point. We exclude this run the results reported in Table i.e., we report average and standard error over the three other runs only.
match setting
21


Reviews:
[{"review_id": "rJgl-c_P_B", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The construction and learning of structured world models is an interesting area of research that could in principle enable better generalisation and interpretability for predictive models. The authors overcome the problem of using pixel-based losses (a common issue being reconstruction of small but potentially important objects) by using a contrastive latent space. The model otherwise makes use of a fixed number of object slots and a GNN transition model, similarly to prior approaches. The authors back up their method with nice results on 3D cubes and 3-body physics domains, and reasonable initial results on two Atari games, with ablations on the different components showing their contributions, so I would give this paper an accept.\n\nThe comparisons to existing literature and related areas is very extensive, with interesting pointers to potential future work - particularly on the transition model and graph embeddings. As expected, the object-factorized action space appears to work well for generalisation, and could be extended/adapted, but setting a fixed number of objects K is a clearly fundamentally limiting hyperparameter, and so showing how the model performs under misspecification of this hyperparameter is useful to know for settings where this is known (2D shapes, 3D blocks, 3-body physics). The fact that K=1 is the best for Pong but K=5 is the best for Space Invaders raises at least two questions: can scaling K > 5 further improve performance on Space Invaders, and is it possible to make the model more robust to a greater-than-needed number of object slots? On a similar note, the data collection procedure for the Atari games seems to indicate that the model is quite sensitive to domains where actions rarely have an impact on the transition dynamics, or the interaction is more complex (e.g. other agents exist in the world) - coming up with a synthetic dataset where the importance of this can be quantified would again aid understanding of the authors' proposed method.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly."}, "extra_text": {}}, {"review_id": "S1gKxxeZ5B", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper aims to learn a structured latent space for images, which is made up of objects and their relations. The method works by (1) extracting object masks via a CNN, (2) turning those masks into feature vectors via an MLP, (3) estimating an action-conditioned delta for each feature via a GNN. Learning happens with contrastive losses, which ask that each feature+delta is close to the true next feature, and far away from other random possibilities. Experiments in simple synthetic environments (e.g., 2D geometric shapes moving on a black background) show encouraging results. \n\nThis paper has a simple, well-motivated method. It is clearly written, and easy to understand. The evaluation is straightforward also: the paper merely shows that this model's nearest neighbors in featurespace are better than the nearest neighbors of World Model (2018) and PAIG (2019). Also, some visualizations indicate that for these simple directional manipulations (up/down/left/right motion), PCA compressions of the model's states have a clean lattice-like structure.\n\nIt is impressive that the model discovers and segments objects so accurately. Perhaps this could actually be evaluated. However, I do not understand why results are so sensitive to the number of object slots (K). This seems like a severe limitation of the model, since in general we have no idea what value to set for this. \n\nAlthough I like the paper, I am not sure that there is sufficient evidence for the method being something useful. Yes, H@1 and MRR are high, but as the paper itself implies, the real goal is to improve performance (or, e.g., sample efficiency) in some downstream task. Given how simple these domains are, and the fact that data is collected with purely random exploration, it is difficult to imagine that there is any significant difference between the training set and the test set. For example, if you make 1000 episodes of 10 steps each in Space Invaders, you practically get 1000 copies of the same 10 frames. I worry that all the evaluation has shown so far is that this model can efficiently represent the state transitions that it has observed.\n\nThe authors note that it was beneficial to only use the hinge on the negative energy term. This seems unusual, since a hinge on the positive term allows some slack, which intuitively makes the objective better-formulated. Can the authors please clarify this result, at least empirically? \n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "H1l81UwCYS", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\nThey perform an extensive assessment of their model, with many ablations, on 2 gridworld environments, one physical domain, and on Atari.\n\nThe 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.\n\nApart from some issues presented below, I feel that this work is of good quality and would recommend it for acceptance.\n\n1.\tThe 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:\n\ta.\tWhy 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. \n\tThis seems overly constraining to me, as this restricts the network to only output “1 bit” of information per “object”.\n\tHowever, maybe being able to represent other factors of these objects might be necessary to make better predictions?  \n\tThis 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.\n\t Did you try allowing the encoder to output more features per objects? \n\tObviously 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.\n\tb.\tIt 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.  \n\tWhat happens if D is larger? Could you find the optimal D by some means?\n2.\tThe 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.\n3.\tSimilarly, 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.\n4.\tIn 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).\n5.\tIt 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.\n6.\tRelated 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.\n7.\tFigure 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.\n8.\tThere 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.\n9.\tIt 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.\n\n\nReferences:\n[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\n[2] Nicholas Watters, Loic Matthey, Sebastian Borgeaud, Rishabh Kabra, Alexander Lerchner, “Spriteworld: A Flexible, Configurable Reinforcement Learning Environment”, https://github.com/deepmind/spriteworld/ \n\n", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents an approach to learn state representations of the scene as well as their action-conditioned transition model, applying contrastive learning on top of a graph neural network. The reviewers unanimously agree that this paper contains a solid research contribution and the authors' response to the reviews further clarified their concerns.
Decision: Accept (Talk)