Interaction-Based Disentanglement of Entities for Object-Centric World Models
Authors: Akihiro Nakano
Masahiro Suzuki
Yutaka Matsuo
Venue: ICLR 2023 poster (2023)
Keywords: object-centric
object-oriented
world models
self-supervised learning
probabilistic deep learning
structured models
video prediction
physics prediction
planning
variational autoencoders
model-based reinforcement learning
VAEs
unsupervised

Abstract: Perceiving the world compositionally in terms of space and time is essential to understanding object dynamics and solving downstream tasks. Object-centric learning using generative models has improved in its ability to learn distinct representations of individual objects and predict their interactions, and how to utilize the learned representations to solve untrained, downstream tasks is a focal question. However, as models struggle to predict object interactions and track the objects accurately, especially for unseen configurations, using object-centric representations in downstream tasks is still a challenge. This paper proposes STEDIE, a new model that disentangles object representations, based on interactions, into interaction-relevant relational features and interaction-irrelevant global features without supervision. Empirical evaluation shows that the proposed model factorizes global features, unaffected by interactions from relational features that are necessary to predict outcome of interactions. We also show that STEDIE achieves better performance in planning tasks and understanding causal relationships. In both tasks, our model not only achieves better performance in terms of reconstruction ability but also utilizes the disentangled representations to solve the tasks in a structured manner.

Published as a conference paper at ICLR 2023
INTERACTION-BASER DISENTANGLEMENT OF ENTI- TIES FOR OBJECT-CENTRE WORLD MODELS
Akihiro Nakano, Masahiro Suzuki, and Yutaka Matsuo
Graduate School of Engincering
The University of Tokyo 7-3-1 Hongo, Bunkyo-ku, Tokyo-to, Japan (nakano. akihi ro, masa, oat.:atooobbbb.......oooo......
ABSTRACT
Perceiving the world compositionally in terms of space and time is essential to un- derstanding object dynamics and solving downstream tasks. Object-centric learn- ing using generative models has improved in its ability to learn distinct represen- tations of individual objects and predict their interactions, and how utilize the learned representations to solve untrained, downstream tasks is a focal question. However, as models struggle to predict object interactions and track the objects accurately, especially for unseen configurations, using object-centric representa- tions in downstream tasks still challenge. This paper proposes STEDIE, a new model that disentangles object representations, based on interactions, into interaction-relevant relational features and interaction-irrelvann global features without supervision. Empirical evaluation shows that the proposed model factor- izes global features, unaffected by interactions from relational features that are necessary to predict the outcome of interactions. We also show that STEDIE achieves better performance in planning tasks and understanding causal relation- ships. In both tasks, our model not only achieves better performance in terms of reconstruction ability but also utilizes the disentangled representations solve the tasks structured manner.
INTRODUCTION
Humans have evolved to perceive the world in a structured manner, such that we can infer unseen objects, predict their interaction with the environment, and plan to use them to perform certain tasks. Previous works have emphasized achieving similar levels of systematic generalization in deep learning (Goyal & Bengio, 2020) requires new set of inductive biases that would enable model to perceive the world as composition of objects and their relationships (Greff al. 2020). Recent works object-centric learning have devised generative models to achieve spatial disentanglement; i... these methods decompose images into individual objects in an unsupervised manner (Kosiorek)
al., 2018, Greff et al.] 2017] Van Steenkiste al.) 2018 Hsieh et al. 2018 Lin al., [2020] Kossen et al.) 2020). This is achieved with the introduction of various inductive biases, such as object propagation and discovery (Kosiorek et al. 2018 Greff et al., 2017: Van Steenkiste et al.,
2018), temporal decomposition (Hsich et al.) (018)) background modeling (Lin et al. 2020), and
object-interaction modeling (Kossen al. 2020).
Given the success of object-centric learning, several studies have investigated its effectiveness in learning world models (Ha & Schmidhuber, 2018)) to solve downstream tasks (Veerapaneni al., 2020 Watters et al.) 2019a] Kossen et al. [2020] Min et al. (202T). As object-centric learning rep- resents the scene as composition of objects and interactions, it would drastically reduce the complexity of modeling the temporal evolution of the scene and, therefore, help in predicting the future and planning. For example, Veerapaneni et al. (2020) proposes OP3, the first fully proba- bilistic object-centric and action-conditioned video prediction model. They show that it can model interactions between entities and plan interactions to solve simple block-stacking tasks. Neverthe- less, prior works have still struggled with solving downstream tasks, modeling object interactions accurately and keeping track of objects as individuals are common challenges. Also, few studies have proposed models conditioned actions (Veerapaneni et al. 2020 Kossen et al. 2020).

---

Published as a conference paper at ICLR 2023
Table 1: Comparison of relevant work done on object-centric generative models and its utilization for model-basee RL. The models can be categorized based on whether it models interactions, is action conditoend or not, and type of factorization of object representations. OP3 [Veerapaneni et al. (2200) is similar to our work as it models object interactions and actions by an agent, but it represents objects as a single latent variable. PARTS (Zoran et al., (021) is another work that, like OP3, uses single latents, but the generative process is not conditioned on actions. DDPAE (Hsieh) ct al. 2018) is also close to our work as it learns factorization implicitly, but interactions are not modeled and is not action-conditionee For thorough comparison, see Appendix G
Action-Condee Conditioned (Plannable)
Method
Models interactions
Type of factorization
VIMON
Single
OP3, PARTS
Single
SQAIR, SCALOR
Explicitly factorized
GSWM
Explicitly factorized
STOVE
Explicitly factorized
DDPAE
X
X
Implicidl factorized
STEDIE (ours)
Implicitly factorized
To improve the generation quality of videos and downstream performances, we introduce interaction-based disentanglement, which aims to factorize the representations of entities into interaction-relovamt relational features and interaction-irrelevaa relevance refers to a feature affecting the future properties of other objects via interaction; e.g., the weight of objects is essential to determining whether an object would be moved after contact oc- curred. At the same time, some objects' features would not be affected via interaction, e.g., the shape of a rigid body would remain unchanged. Importantly, these factorizations must be fully un- supervised to develop model that can handle various causal relationships, hand-crafting decom- position for each task is an infeasible burden. (Table 1 summarizes relevant earlier research (Hsieh) al., 2018: Veerapaneni et al. 2020; Zoran et al. 202T; Kossen al., 2020 Although learn- ing single representation or explicit factorizing object representations for model-based RL has been introduced in prior investigations, implicit factorization has not been explored.
global features. Here, interaction
To enable interaction-based disentanglement, we propose Spatio T'Eporral Disentanglement from Interaction of Entities (STEDIH): a probbbbpobbiiiiiiccc object-centric, and action-conditioond video prediction model. By designing a generative and inference model that implement interaction-bass disentanglement, STEDIE disentangles videos both spatially into objects and temporally into re- lational features and global features. Whereas previous works (Li & Mandi 2018 Hsich et al.) 2018) have factorized into time-varying and time-invariant features, we model object interactions using neural networks to motivate temporal factorization into interaction-eelevan and interaction- irrelevant features. The model can be trained using only raw input video data in an end-to-end fashion with standard evidence lower bound (ELBO). In our experiments, we first demonstrate the model's ability to disentangle videos spatiotemporally. Furthermore, to verify if decompos- ing such a way helps solve downstream tasks, we evaluate the trained model on the planning task and causal relationship understanding task. In the planning task, combined with the cross-entropy method (Rubinstein & Kroese, 2014), we show that STEDIE can perform complex tasks (building block towers) better (up to 13%) and approximately 2x more efficiently than OP3. For causal under- standing, we use CausalMBRL (Ke et al. 2021), )) a recently proposed benchmark video prediction requiring understanding of causal relationships. We show that it outperforms various baselines, including a standard VAE-based model, object-centric method trained with pixel- based loss function (CSWM (Kipf et al. 2020 and OP3 by large margin.
RELATED WORKS
There has been many research on using Variational Autoencoders (VAE) (Kingma & Welling. 2013) to learn object-factorized representations for both images and videos with multiple objects. An object-centric generative model trained on videos should learn disentangled representations about not only the objects but also its dynamics, which would lead to better generation of unseen objects or dynamics and understanding of the scenes. Expanding the data domain temporally has been con- fronted by two major challenges: keeping track of each object and modeling object interactions, such collision. To overcome these hurdles, there have been mainly three approaches in designing
2

---

Published as a conference paper at ICLR 2023
object representations: (1) learning a single representation per object, (2) learning explicit factoriza- tion of representation, and (3) learning factorization of representations unsupervised.
First, many studies have used architectures shown to be effective for static images to extract object-
level representations from vidcos (Weis et al.) 2020; Veerapaneni et al. 2020 Watters et al. 2019a
Creswell et al. 2020; Min et al.
2021iZoral et al. 2021). For example, VIMÓN (Weis et al. 2019 to videos by adding a gated recurrent unit after the encoder expands MONet (Burgess et al.
2020)
to aggregate information over time for each object. OP3 (Veerapaneni et al., 2020) uses IODINE
as the inference model, with additional modules to capture interactions between objects and action interventions.
Secondly, another branch of research has derived from AIR (Eslami et al. 2016), a model that learns several predefined factors of objects ("what", "where". and "presence") ") for a static image to obtain further disentangled object representations (Kosiorek et al.) 2018: Hsieh et al. 2018 [iang et al.) 2019 Kossen et al.) 2020; Crawford & Pineau] [020; [iin et al.] 2020). These models, similar to AIR, learn set of latent variables per object, designed by humans before training. for better object tracking. Introduction of "velocity" and physical interaction (Kossen et al., 2220) or "layer" and
background (Jiang al., 2019 Lin et al. 2020) has improved AlR-derived models' ability to keep track of objects even in cases of occlusion or interaction.
Thirdly, unsupervised time-wise disentanglement of videos into content and pose of objects has been studied (Li & Mandt 2018 [Hiihh et al.) 2018 Zablotskaia et al. [222 Li ct al. 2021). Using VAEs, |Hsieh et al. (2018) proposed learning learning a content vector and a low- oiw-dimenioonal vector predict dynamics for each object in video, but do not model interactions between objects. The work by Li & Mandt 2018 is closest to ours; it proposes a model that learns to disentangle the scene temporally using bI-LSTM (Graves & Schmidhuber (005) architecture, but it is limited to single object per scene. Temporal disentanglement into üime-varying and time-invariant features has been studied in various domains as well, including texts into topics and details (Bowman et al. 2015), speeches into speaker and linguistic features (Hsu et al. 2017). and images into content and detailed textures (Chen) et al. 2016). Unsupervised disentanglement of object representations is also related to how humans
recognize the "sameness" of an object (Baillargeon) 1987; Spelke 1990 Aguiar & Baillargeon 1999 Spelke et al. 1992 Hespos & Baillargeon 2001 SSeelk & Kinzler 2007 2013).
Among these three categories, our proposal falls in the third category. Besides, we aim to learn object-centric world models (action-conditioned video prediction models), and to develop a model that can be used on downstream tasks, such as multi-step planning. As explained in (Table I while single-representation modeling and spplicitffactorizateet
modeling methods have been extended for
usage in model- RL (Veerapaneni et al. 2020; Watters et al. 2019a| Kossen et al.) 2020 Min modeling methods has not been explored. In this paper, we reconsider and broaden temporal disentanglement to ob- ject interactions as the decomposition into time-invariant, intercttioniirelevanl features and time- dependent, interaction-relevand features.
al.] 202T, Zadaianchuk et al.] [021][2022 using unsupervised-factoree
Apart from generative models, models trained using expectation maximization (Greff et al. 2016
2016 Watters et al. 2017 2017 Van Steenkiste al. 2018), pixel-based losses (Battaglia et al.
Xu et al. 2019 Piloto et al. 2022). and contrastive learning (Kipf et al. 2020) have been proposed.
More recently, slot attention (Locatello et al. 2020) has been used to extract object-centric features and predict its dynamics (Kabra et al. |021 Kipf et al. 2021 Elsayed et al., 2022). For temporal disentangiement, approaches using generative adversarial networks (Goodfellow et al., 20200) and autoencoders have been studied (Tulyakov et al. 2018 Villegas et al.) [017]] Denton et al.) 2017).
As we are interested in learning disentangled representations that can be utilized for downstream tasks, we do not consider competing reconstruction ability with these models as the main objective of the present work.
3 SPATIOTEMPORAL DISENTANGLEMENT FROM INTERACTIONS OF ENTITIES
Our goal is to build a probabilistic model for a sequence of images, x (0:T), using Markovian latent variables for each timestep. which are further decomposed into object-centric representations for K distinct objects. As we are interested in utilizing the trained model for model-based RL to solve downstream tasks, we assume an invisible agent with fixed point of view that can interact with the environment. The actions taken by the agent is expressed a (0:7-1); each action consists of moving an object at a certain position to another coordinate. In this paper, we refer the object itself in the real world as objects and its latent representations acquired through the model as entities.
3

---

Published as a conference paper at ICLR 2023
pli
M steps refinement
slots paralle!
Go
K slots parallel
Across timastops
Dynamics Mudule
8/
<
a
Across refinements
Enzeden
x(-1)
(a) Graphical model overview
(b) Detailed model architecture
Figure I: ø1KK and of-k denote global features and relational features respectively. (a) Graphical model visualization of STEDIE. Solid black arrow indicates generation process and dashed red arrow indicates inference procedure. (b) Model architecture of STEDIE. In the inference process, images and actions up to the timestep of interest are used as inputs. To generate at a certain timestep, the relational feature that timestep is concatenated with the global feature. As two features are inferred independently, the global feature inferred at a previous timestpp can be reused.
Figure T (a) shows the graphical model of our proposed model, STEDIE. STEDIE learns to dis- entangle entities into interaction-irelevant (global features, TI.K) and interaction-relevan features
(relational features, o(.k) We hypothesize that entities can be factorized based on interactions as object dynamics can be predicted from the object's partial information. Other features should be learned as they operate as "keys" to identify the same object across time. It is essential to learn to model this factorization implicitly to reduce assumption of the data. For example, the assumption that object size remains unchanged is valid for 2D data but fails for 3D data objects can move towards or away from the observing agent.
As our model builds OP3 Veerapaneni al. 2020), first explain about OP3 and then proceed to explain the extension we incorporated to learn global and relational features for each entity.
3.1 PRELIMINARY: ORJECT-CENTRIC PERCEPTION, PREDICTION, AND PLANNING (OP3)
OP3-Object-centricc Perception, Prediction, and Planning-is a generative model that can de- compose videos (0:T) into K objects (0:T). predict the effect of actions (0:T-1) by agent and object interactions, predict future states, and use obtained representations to solve down- stream planning tasks (Veerapaneni al. (020). The model consists of observation mod- ulc, (x (0:)) |2 ((::2) and a dynamics module, po(2
(t-1) ,a(t-1)), which are trained using (0:7) (r((U:)) (0:2)) IODINE (Greff et al. variational inference and inference model, go(=1. 2019) architecture is used as the inference model. The dynamics module is reused for infer- ence to reduce the number of parameters and training time. For better object-wise decomposi- tion. OP3 proposes an interactive inference algorithm. which updates the posterior and the belief over time. OP3 claims that interactive inference encourages the model disambiguate objects and ensure temporal consistency. Overall, the generative model and inference model can be ex- (0:2")
pressed as (0:T)),), 11 |a (0:T
((2 II P0 (=1
(t-1) ((111) IT p(x(()) (t) 1=1k))
(1) (t-1) (1))) ,.))(())) respectively. Although iterative and yo (= |a. (0:T), a (0:T-1))- II_o Yo (=1:k) 1i:k
(0:T))
(=
inference of IODINE and the proposed interactive inference motivates the model to recognize and keep track of objects, experimental results (Appendix E} show failure cases in which objects move around slots or its color changes over time. We think this is due to dynamics module using full rep- resentation of objects; small variances accumulated over time cause changes to object features that should remain unaffected by the interactions.
3.2 STEDIE
Generative Model. The generation process of our model is included in Figure I (b). STEDIE
assumes that image at timestep i, (t)), is generated from relational features, $10) and global
4

---

Published as a conference paper at ICLR 2023
features, di:K of K distinct objects. Note that because global features are unaffected by interactions, it remains the same regardless of timestep. Hence, the generative model of observing TT' frames of scenes x. (0(111) and latent entities 6 6(0,7) oIk from taking actions a (0:T-1) consists of (1) observation module po (x((1))6 1:K 'I:K) that renders the image at timestep t conditioned on relational features features d(tkk and global features d'K and (2) dynamies module (0 " (a(t)) that models object interactions in latent space. Overall, the generative model is expressed as,
(o(x o00:7)
Jl. lal/aao000 ())
T
= =po(o 1:K (
II po 101x2 a
po II t=0
1000 1:K, OLx),
(0) where p(o is Gaussian distribution with mean 0 and variance I, and po (0_-x ( modeled with an MLP. We condition relational features on global features as these two features are not com- pletely independent. In most cases, static feature affects how the objects move, as the 2D weighted- block pushing dataset (Ke et al., (021) we use in our experiments is an explicit example. We model the observation module using a BroadcastCNN (Watters et al. 2019b). Each entity representations (()) . are decoded into pixel-wise mean, M4x and segmentation mask, my =p(C' = k1o%% ,o)), to model the likelihood, (=(x(() ((1) T1.k] = E my (N pua )²), with the same vari- ance E² for all pixels. Following OP3, we approximate the effect of the intervention on the entities by an action at)), (6 1:K "fxx") a(t-1)),, by breaking into pairwise interactions. The entire dynamics module is modeled as, using MLP-based function, E, (o (0( (a(-))) = (66""16611 (t-1),, 8% Sinteract) where Sinteract Σi ( We apply the same function, doo: for each le 1,.... K to calculate pairwise entity interactions. While full latent variables were used to predict the dynamics previous studies, our model only takes the relational features, 0 This motivates the model to learn object features unaffected by interactions global features.
Inference Model and Learning. (b). The inference model is, The inference process of our model is also included in Figure
as (6 (.:k( (0:21))
=0 a (1:T1)) ya (6 (o) I
(10) Pik 200) a((-1)),
which is factorized into two parts, (1) inferring the global features from the entire sequence, and (2) inferring the relational features per timestep. As single inference in-
(00:)) (0:T), (0:7), fer the parameters at certain timestep is insufficient to break the symmetry to decompose scene into entity representations (Kipf et al. 2 we incorporate iterative inference in several aspects. Following OP3, we use iterative inference Marino et al. 2018) to refine the parameters at each timestep using IODINE (Greff et al. 2019). Iterative inference updates the posterior parameters using gradient information about the ELBO obtained by the current estimate of the parameters. The gradients pass information of what other slots have not learned, therefore encouraging slots to cap- ture different pixels of the image. We also incorporate interactive inference OP3 and update the prior and posterior at each timestep. This ensures keeping track of the "same" object using interac- tive inference for both global and relational features (Appendix F).
However, this causes difference in the expressive power of global and relational features as all timesteps up to the current timestep are used to infer global features while relational features inferred from the current timestep only. Hence, to compensate this problem, we infer global features only at the first refinement step at each timestep. Note that must infer at the first step, as IODINE requires gradient information about the ELBO. To infer the global variable, we process the outputs of IODINE through a bi-LSTM network (Graves & Schmidhuber 2005). To infer object dynam-
ics, we reuse po (t) (t-11) (t-11)) to reduce the number of parameters and training time (ossen et al. 2020 Veerapaneni al. 2020).
k.
As the inference global and relational features are independent at each timestep. it is also possible to update only either of the features. Inferring features independently allows efficient use of the
5

---

Published as a conference paper at ICLR 2023
model during test-time; infer both features initially, update only the relational features to predict interactions, and then combine the features to reconstruct a scene at any timestep.
Finally, parameters of generative model, pg. and inference model, 46, optimized via varia- tional inference to maximize the evidence lower bound (ELBO):
E ((0 ooukrqol0i1k (0022) ox (xx112000)) 2000 log Po(x
1:K T1.K
go(0 (0:2) 11 (x007) (al::1)))) all:T))
(3)
c ( Dkalasla (0:0 a (Ip
where c. E
0 oi 610x-1) 2000)) 11) log po (1x1)) σ1:k)] and -
DKi(go (6(0) 18(t 1) (z)), (tt llp ((t) 1) a 1)) See appendix for derivation of the loss i.K function (ppendixx A and details on model architecture (Appendix B
(go
EXPERIMENTS
We evaluate STEDIE using two datasets each with different characteristics: (1) 3D block moving dataset (Blocks) from Janner et al. (2019) and Veerapaneni et al. (2020). and (2) 2D weighted- pushing dataset (Weights) from Ke et al. (2021). Using these datasets, we conduct three experiments aimed to answer the following questions, (I) Have the global and relational features learned different information? and (2) Does modeling learn disentangled features lead to better performance in downstream tasks, such as planning and causal understanding?
To answer the first question, we conducted latent swapping experiment (subsection 4.1). Specifi- cally, given two unseen samples, we swapped each slot's relational feature accordingly and recon- structed the samples. Since all objects are identical except for their color, we expect features such as position to be learned relational features and features such as color to be learned as global features. If the model has learned to disentangle, generating with the swapped relational features will result in sequences with the same object colors but different positions.
We then analyzed the impact of learning disentangled representations on two downstream tasks (subsection 4.2). First, evaluated the learned representations planning. We used the block stacking task included in Blocks dataset, and used the models to plan the stacking process com- bined with interactionirreeeveee information should help the model plan better, because objects follow the same dynamics regardless of their interaction irrelevant characteristics. Secondly, we used the Weights dataset to evaluate whether learning global and relational features contribute to better understanding of underlying causal relationships between objects in scene. In this experiment, we are interested in investigating whether interaction based disentanglement is effective for a more complex environ- ment. The trained model should capture object color as an interaction-releann feature to utilize to learn causal structure present in the data. We compared STEDIE with previous works in this dataset only, as AIR- -lReb models (ex. STOVE) are known to struggle with data that include back- ground information (Lin et al. 2020). See appendix for further details on training setup, datasets, and downstream tasks (Appendix C).
control (MPC) (Finn & Levine 2017] planning. Intuitively, separating
anleractioniireeevaant
4.1 SWAPPING LATENT VARIABLES
Figure visualizes an example of generated sequences with originally inferred and swapped global features. By reconstructing the sequence using the swapped global features, while retaining the relational features results, the color of the objects were exchanged, but their positions remained the same. This suggests that STEDIE successfully disentangled representations into relational fea- tures that change through interactions, such as position, and global features that remain unchanged throughout interactions. such as color. We also observed that, due to the disentanglemennt STEDIE was able to keep track of objects better even when rolling out longer timesteps (Appendix E).
We further evaluated quantitatively how much object color and position match after swapping the global and relational variables respectively (Table 2}. We calculated mean squared er- (MSE) of average RGB color and intersection over union (loU) of the slots" masks be-

---

Published as a conference paper at ICLR 2023
1=3
1=4
1=2
Reconstruction
Enity
8 (00-))
Entiyy
1:K
Entiy
Eniiy
Reconstruction
Enttyy
Entity
8 (0:7) I:K
Entity
Entiyy
Ground truth sequences
(b) Reconstructed sequenices
(c) Reconstructed sequences swapped global fentures
Figure 2: Visualization of swapping global features between two sequences. The columns represent timesteps. Blue arrows indicate swapped global features and red arrows indicate unswapped rela- tional features. In (b) and (c), top represents the reconstructed whole image and the second to fifth rows represent the the reconstructed image for each slot. Slots are sorted in ascending order of mask size and global features are swapped between the slots with the same slot #.
tween the reconstructed sequences before and after swapping the latent variables. As the objects are the elements that differ between sequences, we calculate the metrics for only the slots which contain the object representations by using mask size as threshold. As Table 2] color MSE of the original sequence and the sequence with the global vari- able swapped is low, whereas having the same relational variable resultss in higher IoU,
Overall, these results show that STEDIE learns disentangle sequences both spatially into ob- ject representations and temporally into time- varying and inve-invariant features. The exper- iment showed that object colors are preserved between sequences with same global variable and object configurations are preserved be- tween sequences with same relational variable. Since entities are learned two different latent variables, the model is capable freely com- bining variables between sequences, suggesting how STEDIE captures the scene in a structured manner.
Table 2:
Color MSE and mask loU on swap- ping global relational latent variables between two sequences the test dataset using STEDIE. Lower MSE indicate that color similar before and after swapping the feature. Higher IoU indi- cates similar position.
Swapping feature
MSE (4)
IoU (t)
Swap relational
0.0838
47.4
Swap global
0.1574
62.8
4.2 EFFECT ON DOWNSTREAM TASKS
Planning. We compare planning results with OP3 VVeeapaneni et al., 2020). Given representa- tions of current timestep. h (t) Concat(a1:k3 (, and inferred goal representations, hi. cost function can be defined as, C(N A(xx)) =E EkEst) minteso c(h(.) ,,f), where S(1) and $a denote predicted entities of current timestep and goal image, respectively and c(.) is some distance function. Assuming that single step is enough move certain object to its goal position, we can greedily choose the actions that achieve minees e(h!.) h17). The actions are sampled and optimized using the cross-entropy method (CEM) (Rubinstein & Kroese 2014). We evaluated the model using (1) accuracy of # of blocks in the correct position where correct position is based goot
threshold of the distance error (Janner et al. 2019). (2) MSE of the reconstructed planned image, and (3) average CEM steps taken when sampling actions. The model was trained using sequences that contained up two objects.

---

Published as a conference paper at ICLR 2023
Table 3: Accuracy (%), MSE, and average CEM steps until reaching the goal state of multi-step planning for building block towers by OP3 and our STEDIE model.
Blocks
Model
Acc. (1)
MSE (U)
Average CEM steps (L)
OP3
85 %
0.0014
1.17
STEDIE (ours)
84%
0.0012
1.00
OP3
58%
0.0049
1.63
2
STEDIE (ours)
63% %
0.0041
1.45
OP3
42%
0.0T02
3.50
3 STEDIE (ours)
55 %
0.0075
1.82
Table 4: Rollout MSE for the Weights dataset calculated at = 1,5, 10. Except for OP3 and STEDIE, the models were trained in two stages, which the encoder and the decoder were trained first without transition and the transition module was trained afterwards. For STOVE, we report the MSE at 5, 10 only as the latents representing object velocities are estimated by conditioning on the first two timesteps.
MSE (L)
Dataset
Model
Loss
(1 1)
(t (t= 5)
(t 10)
VAE
ELBO
167.75
329.83
365.01
CSWM
Pixel MSE
121.91
282.18
307.91
Observed
OP3
ELBO
20.80
275.44
444.14
STOVE
ELBO
N/A
301.49
448.93
STEDIE (ours)
ELBO
68.44
110.76
225.30
VAE
ELBO
111.37
269.79
310.84
CSWM
Pixel MSE
95.57
221.61
234.32
FixseUnobeerved
OP3
ELBO
270.56
267.61
267.34
STOVE
ELBO
N/A
362.23
496.70
STEDIE (ours)
ELBO
81.51
102.18
181.05
VAE
ELBO
110.11
263.43
318.17
CSWM
Pixel MSE
92.57
215.49
229.85
Unobserved
OP3
ELBO
59.08
340.00
616.76
STOVE
ELBO
N/A
118.10
313.22
STEDIE (ours)
ELBO
39.68
47.94
93.18
Figure 3 shows some examples of planning results, and Table 3
shows the quantitative results. As the results show, our model ex- ceeds OP3 except for accuracy for when the number of goal blocks STEDIE achieves higher equivalent accuracy even for unseen numbers of objects, which suggests that disentangling interaction- relevant and interaction-irrelevee predicting object dynamics, enabling the model to accurately pick and place the objects. Our model also results in lower reconstruction error, indicating that the model "imagines" more accurately how the objects' layout will change given an action. Fewer CEM steps until completion indicate that model is able pick the correct actions efficiently. The results especially show improvement in lowering av- erage CEM steps for unseen object configuration (# Blocks=1.3), This supports our hypothesis that separating object representations into global and relational features enhances the model solve tasks efficiently. By separating irrelevant features, STEDIE finds opti- mal actions efficiently by predicting interactions and tracking objects more accurately.
image
Start
Goal image
OPS
STEDIE (ours)
features reduces redundancy when
Figure 3: Visualization of start, goal, and reconstruc- tion of reached state by OP3 and STEDIE. STEDIE reaches the goal state more accurately.
summarizes the rollout performance compared to other methods. We compared STE- DIE against OP3 and STOVE as well as two baseline models provided by Kipf et al. (020) with action vectors incorporated metric. For OP3, STOVE, and STEDIE. the models were trained to predict up to STOVE was trained by modeling the scene as a composition of objects. We also used action vectors to condition the dynamics model without training the reward network. For other baseline
Understanding Causal Relationship. Table 4
the dynamics module, using reconstruction MSE as the evaluation timesteps.
8

---

Published as a conference paper at ICLR 2023
Object
Object
Object
Do nothing Move Up Move Move Down Move Left Move Up Move FightMoove Down Move Lett. Move Up Move Migh Down Move Left
Ground
VAE
CSWM
OP3
Figure 4: Reconstruction example of every possible single rollout on the Weights dataset (Observed). STOVE result is not included as it only uses the dynamics module starting from t = 1. OP3 and our model achieves realistic reconstruction for most cases.
models, the same dataset was broken down into five different single-step rollouts. As Table 4 shows, method achieved the lowest MSE for most cases, outperforming baseline models which were trained using reconstruction-aase
loss. STOVE, albeit its explicit factorization of entities, per- formed worse than STEDIE. This is because STOVE only infers object positions, velocities, and sizes from images (0:7) while their colors are not explicitly modeled, leading to its inability to understand the causal relationship of objects accurately. Although OP3 achieves the low MSE at case, the error increases largely at f 5, 10. This suggests that OP3 has learned to infer a certain timestep accurately using iterative inference, but did not capture the causal relationship and its affect on the dynamics. On the other hand, STEDIE's MSE is consis- tently lower than the other methods when rolled out further into the future. This result indicates that STEDIE able to reconstruct well by capturing the causal relationship of objects.
t = 1 except for
Figure 4 shows the outcome of single rollout of each object into all possible directions. The figure shows the ability of OP3 and STEDIE to understand causal relationships. Here, the model needs to comprehend the relationship for pushing (1) the dark blue square (object 1) towards the left or (2) the light blue circle (object 2) towards the right. The figure shows that both OP3 and STEDIE have captured this causal relationship. Comparing OP3 and STEDIE, Table 4 and Figure 6 show that STEDIE understands the causal relationship better as it tracks object shape and color throughout the sequence and therefore better overall reconstruction. However, its reconstruction can be improved in the future, as it reconstructs object as a square instead cirele in the example above.
5 CONCLUSION
In this paper, we proposed STEDIE, a model that learns to disentangle an image-action sequence spatially into each object and temporally into interaction-irclevaan global features and interaction- relevant relational features, which trained using only raw input data without any supervision (Equation 3). The first experiment, latent swapping experiment, showed our model's ability to dis- entangle object representations in meaningful way. From purely input images and actions, the model decomposed entities into features based on interactions. The next two experiments confirmed that decomposing scenes spatiotemporally leads to better performance in solving downstream tasks. In the planning task, separating irrelevant information enabled the model to use only positional infor- mation predict interactions. In causal understanding task, the results showed that STEDIE learns the causal relationship between the objects better than other models. OP3 did not factorize entities and STOVE factorized them explicitly, but STEDIE's implicit factorization a feasible implemen- tation that can be applied to various datasets, through which it achieves better understanding of object interactions. However, as AlR-based STOVE and spatial-mixture based STEDIE's differ in terms of architecture, implementing explicit factorization using spatial-mixture based models should be explored. Although we experimented on two datasets with different characteristics, further work on designing more natural and complex dataset should be considered to explore the effects of learning multiple object features without supervision. We also leave extending this work to other architectures for future exploration.

---

Published as a conference paper at ICLR 2023
REFERENCES
Andréa Aguiar and Renée Baillargeon. 2.5-month-old infants' reasoning about when objects should and should not be occluded. Cognitive psychology. 39(2):116-157, 1999.
Renee Baillargeon. Object permanence in 31/2-and 41/2-month-old infants. Developmental psy- chology, 23(5);655, 1987.
Peter Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, et al. Interaction networks for learning about objects, relations and physics. Advances in neural information processing systems, 29, 2016.
Samuel R Bowman, Luke Vilnis, Oriol Vinyals, Andrew M Dai, Rafal Jozefowicz, and Samy Ben- gio. Generating sentences from a continuous space. arXiv preprint arXiv:7511.06349 2015.
Christopher P. Burgess, Loic Matthey Nicholas Watters, Rishabh Kabra, Irina Higgins, Matt Botvinick, and Alexander Lerchner. Monet: Unsupervised scene decomposition and represen- tation. arXiv preprint arxiv.1911.11 390, 2019.
Xi Chen, Diederik P Kingma, Tim Salimans, Yan Duan, Prafulla Dhariwal, John Schulman, Ilya Sutskever, and Pieter Abbeel. Variational lossy autoencoder. arXiv preprint arXiv:661.02231,, 2016.
Eric Crawford and Joelle Pineau. Exploiting spatial invariance for scalable unsupervised object tracking. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 3684- 3692, 2020.
Antonia Creswell, Kyriacos Nikiforou, Oriol Vinyals, Andre Saraiva, Rishabh Kabra, Loic Matthey, Chris Burgess, Malcolm Reynolds, Richard Tanburn, Marta Garnelo, et al. Alignnet: Unsuper- vised entity alignment. arXiv preprint arXiv:2007.000 089733, 2020.
Emily I. Denton et al. Unsupervised learning of disentangled representations from video. Advances in neural information processing systems, 30, 2017.
Gamaleldin F Elsayed, Aravindh Mahendran, Sjoerd van Steenkiste, Klaus Greff, Michael C Mozer, and Thomas Kipf. Savi++: Towards end-to-end object-centric learning from real-world videos. arXiv preprint arXiv:2206.07764 2022.
S. M. Eslami, Nicolas Heess, Theophane Weber, Yuval Tassa, David Szepesvari, Koray Kavukcuoglu, and Geoffrey E. Hinton. Attend, infer, repeat: Fast scene understanding with generative models. In Advances in Neural Information Processing Systems., volume 29, 2016.
Chelsea Finn and Sergey Levine. Deep visual foresight for planning robot motion. 2017 IEEE International Conference on Robotics and Automation (ICRA), pp. 2786-2793, 2017.
Ian Goodfellow, Jean Pouget-Abadie, Mchdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 631114:139-1 2020.
Anirudh Goyal and Yoshua Bengio. Inductive biases for deep learning of higher-level cognition. arXiv preprint arXhv:2011.1509) 2020.
Alex Graves and Jürgen Schmidhuber. Framewise phoneme classification with bidirectional Istm and other neural network architectures. Neural networks, 18(5-6):602-610 2005.
Klaus Greff, Antti Rasmus, Mathias Berglund, Tele Hao, Harri Valpola, and Jürgen Schmidhuber. Tagger: Deep unsupervised perceptual grouping. Advances in Neural Information Processing Systems, 29, 2016.
Klaus Greff, Sjoerd Van Steenkiste, and Jürgen Schmidhuber. Neural expectation maximization. Advances in Neural Information Processing Systems, 30, 2017.
Klaus Greff, Raphaël Lopez Kaufman, Rishabh Kabra, Nick Watters, Chris Burgess, Daniel Zoran,
Loic Matthey, Matthew Botvinick, and Alexander Lerchner. Multi-object representation learning with iterative variational inference. International Conference Machine Learning, 2019.
10

---

Published as a conference paper at ICLR 2023
Klaus Greff, Sjoerd van Steenkiste, and Jürgen Schmidhuber. On the binding problem in artificial neural networks. arXiv preprint wwxh.2012.05208, 2020.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 100.:00200.001,,... 2018.
Susan Hespos and Renée Baillargeon. Infants' knowledge about occlusion and containment events:
A surprising discrepancy. Psychological Science, 12(2):141-144 2001.
Jun-Ting Hsich, Bingbin Liu, De-An Huang, Li F Fei-Fei, and Juan Carlos Niebles. Learning to decompose and disentangle representations for video prediction. Advances in neural information processing systems, 31, 2018.
Wei-Ning Hsu, Yu Zhang, and James Glass. Unsupervised learning of disentangled and interpretable representations from sequential data. Advances in neural information processing systems, 30, 2017.
Michael Janner, Sergey Levine, William T. Freeman, Joshua B. Tenenbaum, Chelsea Finn, and Jiajun Wu. Reasoning about physical interactions with object-oriented prediction and planning. arXiv preprint 1tx.iv:1812.1097. 2019.
Jindong Jiang, Sepehr Janghorbani, Gerard De Melo, and Sungjin Ahn. Scalor: Generative world 2019.
models with scalable object representations. arXiv preprint
Rishabh Kabra, Daniel Zoran, Goker Erdogan, Loie Matthey, Antonia Creswell, Matt Botvinick, Alexander Lerchner, and Chris Burgess. Simone: View-invariant, temporally-abstracce, object representations via unsupervised video decomposition. Advances in Neural Information Process- ing Systems, 34:201 34:20146-20155 2021.
Nan Rosemary Ke, Aniket Didolkar, Sarthak Mittal, Anirudh Goyal, Guillaume Lajoic, Stefan Bauer, Danilo Rezende, Yoshua Bengio, Michael Mozer, and Christopher Pal. Systematic evaluation of causal discovery visual model based reinforcement learning. arXiv preprint a0x.h:2107.0088, 2021.
Diederik P. Kingma and Max Welling, arXiv:1312.6114 2013.
Auto-encoding variational bayes.
arXiv preprint
Thomas Kipf, Elise van der Pol, and Max Welling, Contrastive learning structured world models. In International Conference on Learning Representations, 2020.
Thomas Kipf, Gamaleldin F Elsayed, Aravindh Mahendran, Austin Stone, Sara Sabour, Georg Heigold, Rico Jonschkowski, Alexey Dosovitskiy, and Klaus Greff. Conditional object-centric learning from video. arXiv preprint urXiv:2111.P22 12594, 2021.
Adam Kosiorek, Hyunjik Kim, Yee Whye Teh, and Ingmar Posner. Sequential attend, infer, repeat: Generative modelling of moving objects. Advances Neural Information Processing Systems, 31, 2018
Jannik Kossen, Karl Stelzner, Marcel Hussing, Claas Voelcker, and Kristian Kersting, Structured object-aware physics prediction for video modeling and planning. International Conference on Learning Representations, 2020.
Nanbo Li, Muhammad Ahmed Raza, Wenbin Hu, Zhaole Sun, and Robert Fisher. Object-centric representation learning with generative spatial- temporll factorization. Advances in Neural Infor- mation Processing Systems, 34:10772-10788 2021.
Yingzhen Li and Stephan Mandt. Disentangled sequential autoencoder. arXiv preprint
1803v:1803.022 02991, 2018
Zhixuan Lin, Yi-Fu Wu, Skand Peri, Bofeng Fu, Jindong Jiang, and Sungjin Ahn. Improving gener- ative imagination in object-centric world models. International Conference on Machine Learning, 2020.
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot atten- tion. Advances in Neural Information Processing Systems, 13:11525-11538, 2020.

---

Published as a conference paper at ICLR 2023
Joe Marino, Yisong Yue, and Stephan Mandt. Iterative amortized inference. In International Con-
ference on Machine Learning, pp. 3403-3412. PMLR, 2018.
Cheol-Hui Min, Jinseok Bae, Junho Lee, and Young Min Kim. Gatsbi: Generative agent-centric spatio-temporal object interaction. arXiv preprint arxiv. 104. 1eee:2104404222 04275, 2021.
Luis S Piloto, Ari Weinstein, Peter Battaglia, and Matthew Botvinick, Intuitive physics learning in deep-learning model inspired by developmental psychology. Nature human behaviour, pp. 1-11, 2022.
Reuven Y. Rubinstein and Dirk P. Kroese. The cross-entropy method. Information Science and
Statistics, 2014.
Elizabeth S Spelke. Principles of object perception. Cognitive science, 14(1):29-56 1990.
Elizabeth S Spelke. Where perceiving ends and thinking begins: The apprehension of objects in infancy. In Perceptual development in infancy, pp. 209-246. Psychology Press, 2013.
Elizabeth S Spelke and Katherine D Kinzler. Core knowledge. Developmental science, 10(1):89-96, 2007.
Elizabeth S Spelke, Karen Breinlinger, Janet Macomber, and Kristen Jacobson. Origins of knowl- edge. Psychological review, (904) 1992.
Sergey Tulyakov, Ming-Yo Yu Liu, Xiaodong Yang, and Jan Kautz. Mocogan: Decomposing motion and content for video generation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1526-1535, 2018.
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.
Rishi Veerapaneni, John D. Co-Reyes, Michael Chang, Michael Janner, Chelsea Finn, Jiajun Wu, Joshua Tenenbaum, and Sergey Levine. Entity abstraction in visual model-based reinforcement learning. In Proceedings of the Conference on Robot Learning, volume 100, pp. 1439-1456.
PMLR, 2020.
Ruben Villegas, Jimei Yang, Seunghoon Hong, Xunyu Lin, and Honglak Lee. Decomposing motion and content for natural video sequence prediction. arXiv preprint arXiv: UrXiv:1706.000 1706.0803, 2017.
Nicholas Watters, Daniel Zoran, Theophane Weber, Peter Battaglia, Razvan Pascanu, and Andrea Tacchetti. Visual interaction networks: Learning a physics simulator from video. Advances in neural information processing systems, 30, 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 arXiv: /005: 09275, 2019a.
Nicholas Watters, Loic Matthey, Christopher P. Burgess, and Alexander Lerchner. Spatial broadcast decoder: A simple architecture for learning disentangled representations in vaes. arXiv preprint arxiv.1911.770 2019b.
Marissa Weis, Kashyap Chitta, Yash Sharma, Wieland Brendel, Matthias Bethge, Andreas Geiger, and Alexander S. Ecker. Benchmarking unsupervised object representations for video sequences. arXiv preprint arxiv.2006.077 000404 2020.
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: 2019.
Polina Zablotskaia, Edoardo A Dominici, Leonid Sigal, and Andreas M Lehrmann. Provide: A probabilistic framework for unsupervised video decomposition. In Uncertainty in Artificial Intel- ligence, pp. 2019-2028. PMLR, 2021.
12

---

Published as a conference paper at ICLR 2023
Andrii Zadaianchuk, Maximilian Seitzer, and Georg Martius. Sulf-supervised visual reinforcement learning with object-centric representations. In International Conference on Learning Represen- tations, 2021.
Andrii Zadaianchuk, Georg Martius, and Fanny Yang. Self-supervised reinforcement learning with independently controllable subgoals. In Conference on Robot Learning, pp. 384-394. PMLR, 2022.
Daniel Zoran, Rishabh Kabra, Alexander Lerchner, and Danilo J Rezende. Parts: Unsupervised seg- mentation with slots, attention and independence maximization. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 10439-10447, 2021.
13

---

means Published as a conference paper at ICLR 2023
A DERIVATION OF LOss FUNCTION
Here we provide derivation of the ELBO. As we define generative model and inference model as Equation and Equation respectively, the ELBO (Equation 3 can be factorized across time,
log (0(x(0:7)) (0:2)) (1kK(a(TTT)))
(OT o (077) max (
go (117) 01. (x(0.T) (a(0:)))
E B o m (0 (0k po
1.k II:) ( (=))) (o1:)) +
1 ox F 1 -000 -
E +
log po
Po (z (z) T1:k]]
Dru (()) (((x=) x 1lpa (1.. 18 ()-
D (0:1) (1)
T
Σ (E40) c())
Dku(golo1x((2000) ( "1
(0:T
1=0
(4)
where
4. E Egey)
,000000000,0000 a00n myllog po (o(2:1)) po ((1) 11k 01.)]
Bi x 1o 1i ,,1000
(5)
C DKL (ggl0011 I:K (0))
1)
a 1))
B DETAILED ARCHITECTURE
We use similar architecture as OP3 (Veerapaneni et al. 2020). with learning global and relational features for each object as the key difference. In this section, we provide details on the model architecture.
Following IODINE (Greff et al. 2019) and OP3, we use the following inputs to IODINE, where LNN means Layernorm and SG means stop gradients. The following image-sized inputs are concatenated and fed to the corresponding convolutional network:
Description
Formula
LN
SG
Ch.
image
3
/4 mrg
3
mask
mask-logits
mi
mask posterior
p(m, [x, pA) V&L
gradient of means gradient of mask pixelwise likelihood leave-one-out likelih. coordinate channels
3
V Vm
p(æe h)) p(ae (c(())))))
2
total: 17
We concatenate the posterior parameters, λu:K E (ō1:k, Ji:k }. and their gradients with the output of the refinment networks's convolutional part and use the result as input to the LSTM network inside IODINE:
Description gradient posterior
Formula
LN
SG
Xx L
posterior
Ap
14

---

Published as a conference paper at ICLR 2023
Below, we denote Ra, de. dg. and dg dimensions of deterministic components of latent variables, action vectors, global features, and relational features, respectively. The architecture of the inference model is a combination of IODINE and bi-LSTM network as follows:
Type
Size/Ch.
Act. Fn.
Comments for relational for global +actions, for global
MLP
de
MLP
de
Leaky ReLU
biLSTM
+ da
Tanh
MLP
128
Linear
LSTM
128
Tanh
Concat [A, Vx,]
2R,
MLP
128
ELU
Avg. Pool
R
ELU
Conv 3 x
R
ELU
Conv 3
32
ELU
Conv 3 x 3
32
ELU
Inputs
Latent variables concatenated along an axis represent the latent state at a certain timestep. Hence, the architecture of the observation module is as follows:
Decoder
Type
Size/Ch.
Act. Fn.
Comments
Latents
d d8 + Ra
Broadcast
d, + ds + Rg +
+coordinates
Conv 5 x
32
ELU
Conv 5 X
32
ELU
Cony 5 x
32
ELU
Conv 5 x 32
ELU
Conv 5 x 4
Linear
RGB+Mask
We follow implementation by OP3 model object interactions using MLP follows:
Dynamics Network
Function
Size/Ch.
Act. Fn.
d,
128
ELU
da
32
ELU
dao-elf
128
ELU
dao-act
128
Sigmoid
doo-eff
256
ELU
d
256
Sigmoid
deomb
256
ELU
fact
128
fsto
128
The architecture is modification of the interaction function proposed by Van Steenkiste et al. (2018) by conditioning on actions. Dynamics composed of effects of the action and effects from other objects, modeled as summation of pairwise interactions. Overall, it can be expressed as,
dz do(2(())
à (a!
jact daa () a)
(6)
5 ginteract
Σ doa (5gee (act)
(1+1) doomb (8get, ,6 ginteract)
2k
15

---

Published as a conference paper at ICLR 2023
where deo and doo are modeled using two neural networks to predict how and what de- gree other actions or objects affect an object as, d a) d do--en (ōk, a) d ao-an( a) and (oo og = d oo-en( (6tet, ,8ar)) d (0 8grt), respectively. 2 is the concatenation of the deterministic component and stochastic component (relational features). The output, 2k ((-1) is split into deterministic component and relational features using separate networks, Saot and sam. respectively.
C TRAINING AND EXPERIMENT SETUP
For both datasets, we used Adam optimizer and a batch size of 80. We used gradient clipping where if the norm of global gradient exceeds 5.0, then the gradient is scaled down to that norm. As OP3 (Veerapaneni et al., 2020), for better convergence, we divided the latent distribution into stochastic component (size R,) and deterministic component (size Rą), and set the dimensions for both 64. Dimensions of global and relational features were both set to 64. We implemented curriculum training scheme that increases prediction horizon throughout training.
C.I BLOCKS DATASET
The 3D block moving dataset, proposed by Janner et al. (2019), is composed from 10000 sequences of moving around two different-colored objects. Each sequence consist of not only images but also action vectors that represent the picking and dropping coordinates. The actions are biased such that 30% of actions will pick up a block and place it in a random position, 40% of actions will pick up a block and place it on top of another block, and 30% of actions and place it in another random position. The dataset also has a downstream planning task to evaluate the quality (how "good") of the representations that the model has learned is. The task consists of a start image (or sequence) and goal image, and we use the model to plan the actions necessary to reach the goal state. The task is split by the number of goal objects, n n objects, we plan and execute n steps in total as a single step is sufficient to move each object to its goal configuration. We repeat the planning and executing for 2n times or until it has reached the goal state. Accuracy is computed as of blocks in the correct position where correct position is # goal blocks
try to pick from random place
2, 3]. For goal with
based on threshold value of the distance error.
Training. We trained the model with a fixed learning rate of 0.0005 for 400 epochs. The coefficients of KL divergence terms in the loss function were set to 0.01 for relational feature and for global feature. Iterative inference was conducted with 4 refinement steps. The number of slots were set to
Latent Swapping Task. In order to evaluate whether the model has learned to disentangle global and relational features in a meaningful manner, we swapped either variables between sequences and reconstructed them. More precisely, for any two sequences in the test set, we first inferred their latent variables. Using the mask size of the reconstructed images of each slot, we can identify whether a slot includes the object representations. As assume that certain object is tracked in the same slot through time, we used the first timestep to identify the slots with objects and use the same indices for all timesteps in the sequence. Then, swapped either the global or the relational features of the slots that include object representations between the sequences. The features of the remaining two slots were kept unswapped. Finally, we reconstructed the images using the swapped features.
Throughout the sequences, objects' color remain unchanged while objects' position were affected by the actions. Therefore, the model should have learned disentangle colors as global features and positions as relational features. In order to evaluate this quantitatively, we used MSE of average RGB of the slots and loU of slots' masks for global and relational features, respectively. Note that the slots without object representations were not used because their representations were not swapped.
Planning Task. As the model has learned object-centric representations, it can be used to identify the pick location. order to determine the position of a certain object, k, we sampled coordinates from a uniform distribution over the (z. y) space and then calculated the weighted average, using its K attention coefficient p(ōx Jx, v) Eiyz & d an (8pet, faet) from the dynamics module its weight. For each action, we sampled actions 5 times using the CEM method. Given representations of
16

---

Published as a conference paper at ICLR 2023
3D visualization of attention values
2D visualization of attention values
Ground Truth
Reconstruction
14
Start image
1.0
35
(a)
r=
25
Ground Truth
Reconstruction
(b)
Goal image
An example of planning results on the Blocks dataset. (a) Each row corresponds to each timestep during planning (t 0, 1, The images in each column show, from left to right, ground truth image, reconstructed image, three-dimensional visualization of attention values where each color corresponds each slot, and two-dimensional visualization of attention values. For the visualization of attention values, the axes match the environment's 's coordinates. Points were randomly sampled. Ground truth position of the blocks are shown in green squares in the two- dimensional visualization. Black dots with color crosses show the actual pick location for each slot. Results show that attention values change corresponding to changes objects' positions. In general, attention values are higher where the object is and smaller in other locations. We also see that objects are tracked in the same slots, yellow and purple. (b) Ground truth and reconstructed images of the goal configuration. Compared with the images the left (f 2), we see that STEDIE correctly reached the goal configuration.
Figure 5:
current timestep, (t) Concat(O1.K 11: and inferred goal representations, h we used cost function, (( h(.x)) Exesoo minjes c(h (*) .f)) where S(t) and S9 denote predicted entities of current timestep and goal image, respectively, calculate how "good" certain action is. c(-) can be any distance function, but we found that reconstruction MSE results the best performance of both OP3 and STEDIE.
D VISUALIZATION OF PLANNING
In the planning task, calculated attention coefficients, p(& (r, v) find the pick location. Figure 5) shows one example of the planning result and how the location of attention changes in alignment with the objects' positions.
Exes d a (Jrect .8 Jact), to
D.1 WEIGHTS DATASET
The 2D weighted-block pushing dataset, proposed Ke et al. (202)) consists of 2D scenes with multiple, different-coloredd objects. The colors of the objects represent their "weights", such that only heavier object move lighter object while pushing, but vice versa will result in both
17

---

Published as a conference paper at ICLR 2023
objects not moving all. The dataset consists of three variations, namely, Observed, FixedUnob- served, and Unobserved datasets. In the Observed dataset, all objects have the same shade of color and its intensity represents weight, such that darker-colored objects jects, but not vice versa, In the other datasets, objects' colors are sampled from a discrete colormap. For Unoseeeede dataset, the shapes are coupled with colors and can be interpreted hint to identifying its weight. Mcanwhile, in the Unobserved dataset, the shapes are distractors and the color defines objects' weights.
move lighter-colored ob-
Training. We trained the model with a fixed learning rate of 0.00067 for 1000 epochs. The coef- ficients of relational feature and global feature's KL divergence terms in the loss function were set to O initially and linearly increased to 0.01 and 0. respectively, after the 500th epoch in a range of 100 epochs. Iterative inference was conducted with refinement steps. Number of slots were set to 4. To train OP3, we conducted grid search on the hyperparameters with batch size of [16, 32, 80), learning rate of [0.0001, 0.0003), and KL coefficient of (0.001, 0.01]. From the search, we found that batch size of 16, learning rate of 0.0003, and KL term's 's coefficient of 0.001 resulted in the best results.
Causal Understanding Task. To evaluate if the trained model has learned the causal relationship of objects, we generated a random sequence and visualized the rollout results for all possible actions. At any timestep, the agent could choose to move certain object one grid horizontally vertically, or move nothing.
We saw that the attention values were higher around the place where the objects were positioned and lower otherwise. In this example, the slots shown in yellow and purple correspond to the red and gray cubes, respectively. The results also showed that STEDIE was able track objects in the same slot, which was crucial to solve the planning task because the model needed to compare for each object whether it had reached the goal position or not.
E ROLLOUT ANALYSIS
Although OP3 motivates the model to keep track of objects using interactive inference, we found that it fails to keep track of temporally invarying object features (e.g. color). This is because it used the full representation to predict object interactions; object features that are unaffected by interactions were processed through the neural network, and caused the accumulation of small variances. We also found that this caused the object to move between slots across time. These problems have large impact downstream task performance as the model cannot identify the "same" object between the current timestep and the goal state. Figure 6 and Figure show examples comparison between OP3 and STEDIE in terms of rollout, and Figure 8 shows MSE of mean color of slots, which include the objects.
We saw that STEDIE was more capable of tracking objects better as it learned to disentangle Interaction-imrllvwee features as global features and use only relational features to predict inter- actions. On the other hand, we observed that OP3 lacks both spatial and temporal consistency.
F ABLATIONS
We provide ablation results interactive inference to infer global features and relational features in Table 5) Following OP3, we updated the parameters at each timestep to ensure timewise consis- tency for both features. Although the model can learn reconstruct the video without interactive inference, this leads to poor planning performance as object identity is swapped between slots across time.
G COMPARISON WITH OTHER RELATED WORKS
In this section, we provide more detailed comparison with related works on object-centric learning for multi-object videos (Table 6).
18

---

Published as a conference paper at ICLR 2023
Timestep
Ground Truth
Racoratructin
Erdty
Endty
Endty
OP3
Reconatruction
Enblyy
Endly
STEDIE (ours)
Figure 6: Long-term rollout results of OP3 and STEDIE on the Blocks dataset. Columns represent timestep. Both models were trained with sequences with a maximum length of Although STEDIE does not reconstruct all timesteps accurately, it is able to track the same object.
Timestep
Ground Truth
Reconsimuution
:se
Entty 1
Exttyy
Extiy
Exttyy
OP3
Reconstnuction
Entty
Entiy 2
Entty
STEDIE (ours)
Figure 7: Long-term rollout results of OP3 and STEDIE on the Weights dataset. Columns represent timestep. Both models were trained with sequences with maximum length of 5. Similar to the rollout example on the Blocks dataset, STEDIE is able to track the same object, leading to better causal understanding.
19

---

Published as a conference paper at ICLR 2023
Mean color MSE
OP3 STEDIE (ours)
50
MSC
Timestep
Figure 8: MSE of mean color for slots which include the objects in the Blocks dataset. Mask sizes of the slots were used identify whether slot includes information about object or STEDIE is able to reduce color MSE by large margin, as it ensures temporal consistency by disentangling based on interactions.
Table 5: Ablation results multi-step planning for building block towers non-/usage of interac- tive inference (int. inf.). Evaluation metrics are accuracy (%), MSE, and average CEM steps until reaching goal state. N/A indicates that the model could only predict not to move objects already at their goal position.
# Blocks
STEDIE
Acc. (f)
MSE (L)
Average CEM steps (4)
w/o int. inf.
26%
0.0035
1.00
w/ int inf.
84%
0.0012
1.00
w/o int. inf.
15%
0.0071
1.00
2
w/ int. inf.
63%
0.0041
1.45
w/o int. inf.
19%
0.0141
N/A
3
w/ int. inf.
55%
0.0075
1.82
20

---

Published as a conference paper at ICLR 2023
Table 6: Detailed comparison of relevant prior works on object-centric and its utilization for model- based RL. The models can be categorized based on whether models interactions, is action- conditioned or not, type of factorization of object representations, and objective function. t: PARTS (Zoran et al. 2021) expands OP3 by using transformer-based dynamics module instead of modeling pair-wise interactions and replacing the prior Gaussians that independent of the timestep, past latents, and actions. In order for the global and relational features to capture interaction-relevant and interaction-irrelevva generative process is conditioned on the actions. Our experimental results show that this also relates with solving OP3's limitations which were raised in PARTS, that "interactions may cause the latent variable to effectively lose its slotted structure and slot independence as time passes in the sequence".. We show in our experiments that STEDIE disentangles object representations by conditioning on ac- tions, and that it leads to better reconstruction, object tracking, and scene understanding.
features, respectively, we think it crucial that the
Models Interactions
Action-Conditoon Conditioned (Plannable)
Type of factorization single single implicitly factorized
Objective fn.
Method
ViMON
ELBO
TBA
X
Pixel-based
X
DDPAE
ELBO
SQAIR, SCALOR
X
explicitly factorized
ELBO
CSWM, SAVI, +AVI++
single
Pixel-based
OP3
single
ELBO
PARTS
V
single
ELBO
GSWM
x
explicitly factorized
ELBO
STOVE
explicitly factorized
ELBO
STEDIE (ours)
implicitly factorized
ELBO
21


Reviews:
[{"review_id": "yd5TCFmIZM", "reviewer": "Reviewer_Hvkp", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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": "The paper propose STEDIE, a new model that disentangles object representations based on interactions, into interaction-relevant relational features and interaction-irrelevant global features without direct supervision. Across different experiments, the authors showcases the model's effectiveness.", "summary_of_the_review": "Please see my strength and weakness.", "strength_and_weaknesses": "I want to preface this by saying that I'm not really familiar with the field of object-centric learning.\n\nStrength:\n- I think the model is well designed and has some interesting novelties. The authors do a good job explaining the model.\n\nWeakness:\n- Equation one \\sigma has a (0) superscript, which seems to be a typo.\n- I think overall, the downstream task evaluation seems a bit underwhelming to me.\n1. I think the evaluation metrics and visualizations are not exactly informative. As far as I can tell from the paper, the action and measurements are done in 3D. The evaluation metrics of accuracy is based off a distance metric that's unspecified in the paper and appendix. Not showing direct measurement in 3D limits my confidence in the model. MSE is calcuated in pixel space, which is weird, because say we move the cube along the camera direction, it results in not so much change in pixel mse, but the 3D location can be far off.\n2. From what I can tell, OP3 seems like a model from quite awhile ago, while STOVE seems to be a very relevant baseline. I can't really tell from the paper which reference is STOVE, and why the authors are not comparing with it.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ku5sZVuNmtH", "reviewer": "Reviewer_wuf5", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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": "This paper proposes a method that learns disentangled representations of objects determined by interaction-relevant and interaction-irrelevant latents. The proposed method is based on OP3, but with a formulation that encourages the desired disentanglement. In experiments, the authors show that the interaction-relevant and interaction-irrelevant latent variables are in fact disentangled by swapping them between two different sequences. They also show that the disentangled representations lead to improved performance in downstream tasks against the baselines.", "summary_of_the_review": "In conclusion, I find this paper interesting regardless of the toy setup the conceptual contribution is shown. I have a couple of questions that I have stated in the weakness section that I would like the authors to clarify. In the mean time, I am leaning towards accepting this paper, but I am looking forward to discussions to determine the final assessment.", "strength_and_weaknesses": "Strengths:\n+ New method for disentangled object centric representation learning\n+ Experiments showing the disentanglement is present\n+ Experiments showing the disentanglement leads to better downstream performance\n\n\nWeaknesses:\n\n- Intersection over union evaluation:\nThe authors present an intersection over union evaluation to show that interaction-relevant latents contain location information. However, I am not clear how the authors are able to get locations from the latents, and so, can't really see how this evaluation is possible. Can the authors clarify how this is done? Or point out where in the paper this is described. If this is in fact somewhere in the paper, and I cannot find it, it would be useful if this is revised later on so it's easier to find.\n\n\n\n- Pairwise interactions mechanism visualizations:\nIn page 5, the authors mention that they use a d_{o,o} function to calculate pairwise interactions between objects. However, there is no analysis of how these interactions are happening after a model is trained. If this is possible, it would be good to have some type of visualization showing how these become active when objects are interacting in the pixel space.\n\n\n\n- Typos in the writing:\nThere are many typos throughout the text. For example, in section 3.2, first paragraph, line 7, module is misspelled as moduele. In addition, there is misuse of the verb is. For example, in page 6, 4th paragraph, the authors say \"Finally, parameters of a generative model pθ and an inference model qφ is\" while it should be \"Finally, parameters of a generative model pθ and an inference model qφ are\". These type of typos are happening in multiple places in the manuscript. I suggest the authors proof read the manuscript and fix these writing issues.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "AeGI7t4LWi", "reviewer": "Reviewer_sHVk", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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": "This paper proposes a structured world model which learns to disentangle image sequences into interaction-irrelevant global features and interaction-relevant relational features without supervision. The design of the model continues the ideas from IODINE, which iteratively inference latent representations to generative target images. Differently, the global features and relational features are designed to be separated and learned implicitly. The model is evaluated on two downstream tasks: planning and casual understanding, experiments show the effectiveness of the disentangled representation.", "summary_of_the_review": "The paper presents a method to disentangle representations based on interaction. The ideas are validated quantitatively and qualitatively. Yet, I am concerned about the insufficient evaluation. I am open to adjusting my rating if the authors do sufficiently address the points raised in weaknesses.", "strength_and_weaknesses": "Strength:\nThe paper presents a very interesting idea of disentangling representations based on interaction.\n\nExperiments show that the learned representations achieve better performance in planning tasks and understanding causal relationships.\n\nWeakness:\nMissing discussion on type of factorization, especially between implicitly factorized representation and explicitly factorized representation.\n\nPlease discuss why the performance of OP3 on the Weights dataset is so poor even when doing nothing, because the result on the block moving dataset seems fine.\n\nI would like to see the video prediction performance comparison between OP3 and STEDIE on the non-action-conditioned datasets, like CLEVRER.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "24gG2hjFKq", "reviewer": "Reviewer_qwR6", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "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": "This paper proposes a probabilistic model to learn disentangled object representation in an unsupervised fashion. The generative model is expressed through an observation module and a dynamic module, where the representations of entities are factorized into interaction-relevant relational features and interaction-irrelevant global features. The model is optimized with variational inference by maximizing the ELBO. Latent swapping experiments and downstream tasks are conducted on Block and Weights datasets.", "summary_of_the_review": "My ratings are given based on my evaluation of the strengths and weaknesses above. The authors are expected to address my concerns over the experiments.", "strength_and_weaknesses": "Strengths:\n- This paper differentiates itself from prior work by disentangling the object-centric representation to interaction-relevant feature and the global feature, which is a promising direction for unsupervised representation learning.\n- The experiments show that the model can learn disentangled representations, and the learned representations can facilitate downstream tasks in planning and roll-out reconstruction.\n\nWeaknesses:\nThe major concern lies in the experiments. 1) the latent swapping experiments are conducted on the block moving datasets, and the involved features are relatively simple, i.e., color for global feature and position as relational feature. Prior work also demonstrates such disentangled representation even without explicit modeling of the dynamic and static features. Without further quantitative comparison with other baselines or in more complex scenarios, it’s hard to justify the significance of the proposed model. 2) The downstream experiments are mainly compared against OP3, which does not disambiguate global or relational features. Are there justifications why not compare to factorized models like STOVE?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster