SPECTRA: Sparse Entity-centric Transitions

Learning an agent that interacts with objects is ubiquituous in many RL tasks . In most of them the agent ’ s actions have sparse effects : only a small subset of objects in the visual scene will be affected by the action taken . We introduce SPECTRA , a model for learning slot-structured transitions from raw visual observations that embodies this sparsity assumption . Our model is composed of a perception module that decomposes the visual scene into a set of latent objects representations ( i.e . slot-structured ) and a transition module that predicts the next latent set slot-wise and in a sparse way . We show that learning a perception module jointly with a sparse slot-structured transition model not only biases the model towards more entity-centric perceptual groupings but also enables intrinsic exploration strategy that aims at maximizing the number of objects changed in the agents trajectory . 1 INTRODUCTION . Recent model-free deep reinforcement learning ( DRL ) approaches have achieved human-level performance in a wide range of tasks such as games ( Mnih et al. , 2015 ) . A critical known drawback of these approaches is the vast amount of experience required to achieve good performance . The promise of model-based DRL is to improve sample-efficiency and generalization capacity across tasks . However model-based algorithms pose strong requirements about the models used . They have to make accurate predictions about the future states which can be very hard when dealing with high dimensional inputs such as images . Thus one of the core challenge in model-based DRL is learning accurate and computationally efficient transition models through interacting with the environment . Buesing et al . ( 2018 ) developed state-space models techniques to reduce computational complexity by making predictions at a higher level of abstraction , rather than at the level of raw pixel observations . However these methods focused on learning a state-space model that doesn ’ t capture the compositional nature of observations : the visual scene is represented by a single latent vector and thus can not be expected to generalize well to different objects layouts . Extensive work in cognitive science ( Baillargeon et al. , 1985 ; Spelke , 2013 ) indeed show that human perception is structured around objects . Object-oriented MDPs ( Diuk et al. , 2008 ) show the benefit of using object-oriented representations for structured exploration although the framework as it is presented requires hand-crafted symbolic representations . Bengio ( 2017 ) proposed as a prior ( the consciousness prior ) that the dependency between high-level variables ( such as those describing actions , states and their changes ) be represented by a sparse factor graph , i.e. , with few high-level variables at a time interacting closely , and inference performed sequentially using attention mechanisms to select a few relevant variables at each step . Besides , a recent line of work ( Greff et al. , 2017 ; van Steenkiste et al. , 2018 ; Eslami et al. , 2016 ; Kosiorek et al. , 2018 ; Greff et al. , 2019 ; Burgess et al. , 2019 ) has focused on unsupervised ways to decompose a raw visual scene in terms of objects . They rely on a slot-structured representation ( see Figure 1 ) of the scene where the latent space is a set of vectors and each vector of the set is supposed to represent an “ object ” ( which we refer to as “ entity ” ) of the scene . However , to the best of our knowledge , Watters et al . ( 2019 ) is the only work that investigates the usefulness of slot-structured representations for RL . They introduced a method to learn a transition model that is applied to all the slots of their latent scene representation . Extending their work , we go further and posit that slot-wise transformations should be sparse and that the perception module should be learned jointly with the transition model . We introduce Sparse Entity-Centric Transitions ( SPECTRA ) , an entity-centric action-conditioned transition model that embodies the fact that the agents actions have sparse effects : that means that each action will change only a few slots in the latent set and let the remaining ones unchanged . This is motivated by the physical consideration that agent interventions are localized in time and space . Our contribution is motivated by three advantages : − Sparse transitions enable transferable model learning . The intuition here is that the sparsity of the transitions will bias the model towards learning primitive transformations ( e.g . how pushing a box affects the state of a box being pushed etc ) rather than configurationdependent transformations , the former being more directly transferable to environments with increased combinatorial complexity . − Sparse transitions enable a perception module ( when trained jointly ) to be biased towards more meaningful perceptual groupings , thus giving potentially better representations that can be used for downstream tasks , compared to representations learned from static data . − Sparse transitions enable an exploration strategy that learns to predict actions that will change the state of as many entities as possible in the environment without relying on pixels error loss . 2 RELATED WORK . Unsupervised visual scene decomposition . Learning good representations of complex visual scenes is a challenging problem for AI models that is far from solved . Recent work ( Greff et al. , 2017 ; van Steenkiste et al. , 2018 ; Eslami et al. , 2016 ; Kosiorek et al. , 2018 ; Greff et al. , 2019 ; Burgess et al. , 2019 ) has focused on learning models that discover objects in the visual scene . Greff et al . ( 2019 ) further advocates for the importance of learning to segment and represent objects jointly . Like us they approach the problem from a spatial mixture perspective . van Steenkiste et al . ( 2018 ) and Kosiorek et al . ( 2018 ) build upon Greff et al . ( 2017 ) and Eslami et al . ( 2016 ) respectively by incorporating next-step prediction as part of the training objective in order to guide the network to learn about essential properties of objects . As specified in van Steenkiste et al . ( 2019 ) we also believe that objects are task-dependent and that learning a slot-based representations along with sparse transitions bias the perception module towards entity-centric perceptual groupings and that those structured representations could be better suited for RL downstream tasks . Slot-based representation for RL . Recent advances in deep reinforcement learning are in part driven by a capacity to learn good representations that can be used by an agent to update its policy . Zambaldi et al . ( 2018 ) showed the importance of having structured representations and computation when it comes to tasks that explicitly targets relational reasoning . Watters et al . ( 2019 ) also show the importance of learning representations of the world in terms of objects in a simple model-based setting . Zambaldi et al . ( 2018 ) focuses on task-dependent structured computation . They use a selfattention mechanism ( Vaswani et al. , 2017 ) to model an actor-critic based agent where vectors in the set are supposed to represent entities in the current observation . Like Watters et al . ( 2019 ) we take a model-based approach : our aim is to learn task-independent slot-based representations that can be further used in downstream tasks . We leave the RL part for future work and focus on how learning those representations jointly with a sparse transition model may help learn a better transition model . 3 SPECTRA . Our model is composed of two main components : a perception module and a transition module ( section 3.1 ) . The way we formulated the transition implicitly defines an exploration policy ( section 3.3 ) that aims at changing the states of as many entities as possible . Choice of Environment . Here we are interested in environments containing entities an agent can interact with and where actions only affect a few of them . Sokoban is thus a good testbed for our model . It consists of a difficult puzzle domain requiring an agent to push a set of boxes onto goal locations . Irreversible wrong moves can make the puzzle unsolvable . Each room is composed of walls , boxes , targets , floor and the agent avatar . The agent can take 9 different actions ( no-op , 4 types of push and 4 types of move ) . Fully Observed vs Learned Entities . The whole point is to work with slot-based representations learned from a raw pixels input . There is no guarantee that those learned slots will effectively correspond to entities in the image . We thus distinguish two versions of the environment ( that correspond to two different levels of abstraction ) : − Fully observed entities : the input is structured . Each entity corresponds to a spatial location in the grid . Entities are thus represented by their one-hot label and indexed by their x-y coordinate . This will be referred to as the fully observed setting . There is no need for a perception module in this setting . − Raw pixels input : the input is unstructured . We need to infer the latent entities representations . This will be referred to as the latent setting . 3.1 MODEL OVERVIEW . The idea is to learn an action-conditioned model of the world where at each time step the following take place : − Pairwise Interactions : Each slot in the set gathers relevant information about the slots conditioned on the action taken − Active entity selection : Select slots that will be modified by the action taken − Update : Update the selected slots and let the other ones remain unchanged . Ideally , slots would correspond to unsupervisedly learned entity-centric representations of a raw visual input like it is done by Burgess et al . ( 2019 ) ; Greff et al . ( 2019 ) . We show that learning such perception modules jointly with the sparse transition biases the perceptual groupings to be entity-centric . Perception module . The perception module is composed of an encoder fenc and a decoder fdec . The encoder maps the input image x to a set of K latent entities such that at time-step t we have fenc ( x t ) = st ∈ RK×p . It thus outputs a slot-based representation of the scene where each slot is represented in the same way and is supposed to capture properties of one entity of the scene . Like ( Burgess et al. , 2019 ; Greff et al. , 2019 ) we model the input image xt with a spatial Gaussian Mixture Model . Each slot stk is decoded by the same decoder fdec into a pixel-wise mean µik and a pixel-wise assignment mtik ( non-negative and summing to 1 over k ) . Assuming that the pixels i are independent conditioned on st , the conditional likelihood thus becomes : pθ ( x t|st ) = D∏ i=1 ∑ k mtikN ( xti ; µtik , σ2 ) with µtik , mtik = fdec ( stk ) i . As our main goal is to investigate how sparse transitions bias the groupings of entities , in our experiments we use a very simple perception module represented in Figure 1 . We leave it for future work to incorporate more sophisticated perception modules . Pairwise interactions . In order to estimate the transition dynamics , we want to select relevant entities ( represented at time t by the set st ∈ RK×p ) that will be affected by the action taken , so we model the fact that each entity needs to gather useful information from entities interacting with the agent ( i.e . is the agent close ? is the agent blocked by a wall or a box ? etc .. ) . To that end we propose to use a self-attention mechanism ( Vaswani et al. , 2017 ) . From the k-th entity representation stk at time t , we extract a row-vector key K t k , a row-vector query Q t k and a row-vector value V t k conditioned on the action taken such that ( aggregating the rows into corresponding matrices and ignoring the temporal indices ) : s̃ = softmax ( KQT√ d ) V where the softmax is applied separately on each row . In practice we concatenate the results of several attention heads to use it as input to the entity selection phase . Entity selection . Once the entities are informed w.r.t . possible pairwise interactions the model needs to select which of these entities will be affected by the action taken at . Selection of the entities are regulated by a selection gate ( Hochreiter & Schmidhuber , 1997 ; Cho et al. , 2014 ) computed slotwise as : f tk = σ ( MLP ( [ s̃ t k ; a t ] ) ) ( 1 ) where f tk can be interpreted as the probability for an entity to be selected . Update . Finally , each selected entity is updated conditioned on its state stk at time-step t and the action taken at . We thus simply have : st+1k = f t kfθ ( [ s t k , a t ] ) + ( 1− f tk ) stk fθ is a learned action-conditioned transformation that is applied slot-wise . We posit that enforcing the transitions to be slot-wise and implicitly sparse will bias the model towards learning more primitive transformations . We verify this assumption in next subsection in the simpler case where the entities are fully observed ( and not inferred with a perception module ) .

Summary: This paper introduces a model that learns a slot-based representation, along with a transition model to predict the evolution of these representations in a sparse fashion, all in a fully unsupervised way. This is done by leveraging a self-attention mechanism to decide which slots should be updated in a given transition, leaving the others untouched. The model learns to encode in a slot-wise and is trained on single step transitions.