Contrastive Learning of Structured World Models

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 . 1 INTRODUCTION . Compositional reasoning in terms of objects , relations , and actions is a central ability in human cognition ( Spelke & Kinzler , 2007 ) . This ability serves as a core motivation behind a 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 a strong inductive bias for generalization to 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 a completely unsupervised or self-supervised manner ( Eslami et al. , 2016 ; Greff et al. , 2017 ; Nash et al. , 2017 ; van Steenkiste et al. , 2018 ; Kosiorek et al. , 2018 ; Janner et al. , 2019 ; Xu et al. , 2019 ; Burgess et al. , 2019 ; Greff et al. , 2019 ; Engelcke et al. , 2019 ) . These methods follow a generative approach , i.e. , they learn to discover object-based representations by performing visual predictions or reconstruction and by optimizing an objective in pixel space . Placing a loss in pixel space requires carefully trading off structural constraints on latent variables vs. accuracy of pixel-based reconstruction . Typical failure modes include ignoring visually small , but relevant features for predicting the future , such as a bullet in an Atari game ( Kaiser et al. , 2019 ) , or wasting model capacity on visually rich , but otherwise potentially irrelevant features , such as static backgrounds . To avoid such failure modes , we propose to adopt a discriminative approach using contrastive learning , which scores real against fake experiences in the form of state-action-state triples from an experience buffer ( Lin , 1992 ) , in a similar fashion as typical graph embedding approaches score true facts in the form of entity-relation-entity triples against corrupted triples or fake facts . We introduce Contrastively-trained Structured World Models ( C-SWMs ) , a class of models for learning abstract state representations from observations in an environment . C-SWMs learn a set of abstract state variables , one for each object in a particular observation . Environment transitions are modeled using a graph neural network ( Scarselli et al. , 2009 ; Li et al. , 2015 ; Kipf & Welling , 2016 ; Gilmer et 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 at this formulation by adapting methods for learning translational graph embeddings ( Bordes et al. , 2013 ; Wang et al. , 2014 ) to our use case . By establishing a connection between contrastive learning of state abstractions ( François-Lavet et al. , 2018 ; Thomas et al. , 2018 ) and relational graph embeddings ( Nickel et al. , 2016a ) , we hope to provide inspiration and guidance for future model improvements in both fields . In a set of experiments , where we use a novel ranking-based evaluation strategy , we demonstrate that C-SWMs learn interpretable object-level state abstractions , accurately learn to predict state transitions many steps into the future , demonstrate combinatorial generalization to novel environment configurations and learn to identify objects from scenes without supervision . 2 STRUCTURED WORLD MODELS . Our goal is to learn an object-oriented abstraction of a particular observation or environment state . In addition , we would like to learn an action-conditioned 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.1–2.2 , and in the following describe a variant that utilizes object-factorized state representations , which we term a Structured World Model . 2.1 STATE ABSTRACTION . We consider an off-policy setting , where we operate solely on a buffer of offline experience , e.g. , obtained from an exploration policy . Formally , this experience buffer B = { ( st , at , st+1 ) } Tt=1 contains T tuples of states st ∈ S , actions at ∈ A , and follow-up states st+1 ∈ S , which are reached after taking action at . We do not consider rewards as part of our framework for simplicity . Our goal is to learn abstract or latent representations zt ∈ Z of environment states st ∈ S that discard any information which is not necessary to predict the abstract representation of the followup state zt+1 ∈ Z after taking action at . Formally , we have an encoder E : S → Z which maps observed states to abstract state representations and a transition model T : Z × A → Z operating solely on abstract state representations . 2.2 CONTRASTIVE LEARNING . Our starting point is the graph embedding method TransE ( Bordes et al. , 2013 ) : TransE embeds facts from a knowledge base K = { ( et , rt , ot ) } Tt=1 , which consists of entity-relation-entity triples ( et , rt , ot ) , where et is the subject entity ( analogous to the source state st in our case ) , rt is the relation ( analogous to the action at in our experience buffer ) , and ot is the object entity ( analogous to the target state st+1 ) . TransE defines the energy of a triple ( et , rt , ot ) as H = d ( F ( et ) + G ( rt ) , F ( ot ) ) , where F ( and G ) are embedding functions that map discrete entities ( and relations ) to RD , where D is the dimensionality 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 , we replace G ( rt ) with T ( zt , at ) , i.e. , with the transition function , conditioned on both the action and the ( embedded ) source state via zt = E ( st ) . The overall energy of a state-action-state triple then can be defined as follows : H = d ( zt + T ( zt , at ) , zt+1 ) . This additive form of the transition model provides a strong inductive bias for modeling effects of actions in the environment as translations in the abstract state space . Alternatively , one could model effects as linear transformations or rotations in the abstract state space , which motivates the use of a 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 ( zt + T ( zt , at ) , zt+1 ) + max ( 0 , γ − d ( z̃t , zt+1 ) ) , ( 1 ) defined for a single ( st , at , st+1 ) with a corrupted abstract state z̃t = E ( s̃t ) . s̃t is sampled at random from the experience buffer . The margin γ is a hyperparameter for which we found γ = 1 to be a good choice . Unlike Bordes et 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 zt , 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 to learn a relational and object-oriented model of the environment that operates on a factored abstract state space Z = Z1× . . .×ZK , whereK is the number of available object slots . We further assume an object-factorized action space A = A1 × . . .×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 object-factorized representations is shown in Figure 1 . Encoder and Object Extractor We split the encoder into two separate modules : 1 ) a CNN-based object extractor Eext , and 2 ) an MLP-based object encoder Eenc . The object extractor module is a CNN operating directly on image-based observations from the environment with K feature maps in its last layer . Each feature mapmkt = [ Eext ( st ) ] k can be interpreted as an object mask corresponding to one particular object slot , where [ . . . ] k denotes selection of the k-th feature map . For simplicity , we only assign a 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 to produce multiple feature maps per object slot . After the object extractor module , we flatten each feature map mkt ( object mask ) and feed it into the object encoder Eenc . The object encoder shares weights across objects and returns an abstract state representation : zkt = Eenc ( m k t ) with z k t ∈ Zk . We set Zk = RD in the following , where D is a hyperparameter . Relational Transition Model We implement the transition model as a graph neural network ( Scarselli et al. , 2009 ; Li et al. , 2015 ; Kipf & Welling , 2016 ; Battaglia et al. , 2016 ; Gilmer et al. , 2017 ; Battaglia et al. , 2018 ) , which allows us to model pairwise interactions between object states while being invariant to the order in which objects are represented . After the encoder stage , we have an abstract state description zkt ∈ Zk and an action akt ∈ Ak for every object in the scene . We represent actions as one-hot vectors ( or a vector of zeros if no action is applied to a particular object ) , but note that other choices are possible , e.g. , for continuous action spaces . The transition function then takes as input the tuple of object representations zt = ( z1t , . . . , z K t ) and actions at = ( a 1 t , . . . , a K t ) at a particular time step : ∆zt = T ( zt , at ) = GNN ( { ( zkt , akt ) } Kk=1 ) . ( 2 ) T ( zt , at ) is implemented as a graph neural network ( GNN ) that takes zkt as input node features . The model predicts updates ∆zt = ( ∆z1t , . . . , ∆z K t ) . The object representations for the next time step are obtained via zt+1 = ( z1t + ∆z 1 t , . . . , z K t + ∆z K t ) . The GNN consists of node update functions fnode and edge update functions fedge with shared parameters across all nodes and edges . These functions are implemented as MLPs and we choose the following form of message passing updates : e ( i , j ) t = fedge ( [ z i t , z j t ] ) ( 3 ) ∆zjt = fnode ( [ z j t , a j t , ∑ i 6=j e ( i , j ) t ] ) , ( 4 ) where e ( i , j ) t is an intermediate representation of the edge or interaction between nodes i and j . This corresponds to a 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 isO ( K2 ) . 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-th object as ∆zkt = T k ( zt , at ) in the following . Multi-object Contrastive Loss We only need to change the energy function to take the factorization of the abstract state space into account , which yields the following energy H for positive triples and H̃ for negative samples : H = 1 K K∑ k=1 d ( zkt + T k ( zt , at ) , z k t+1 ) , H̃ = 1 K K∑ k=1 d ( z̃kt , z k t+1 ) , ( 5 ) where z̃kt is the k-th object representation of the negative state sample z̃t = E ( s̃t ) . The overall contrastive loss for a single state-action-state sample from the experience buffer then takes the form : L = H + max ( 0 , γ − H̃ ) . ( 6 )

Summary: 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.