Unsupervised Object-Based Transition Models For 3D Partially Observable Environments
Authors: Antonia Creswell
Rishabh Kabra
Christopher P. Burgess
Murray Shanahan
Venue: NeurIPS 2021 Poster (2021)
Keywords: Objects
transition model
alignment
entities
MONet

Abstract: We present a slot-wise, object-based transition model that decomposes a scene into objects, aligns them (with respect to a slot-wise object memory) to maintain a consistent order across time, and predicts how those objects evolve over successive frames. The model is trained end-to-end without supervision using transition losses at the level of the object-structured representation rather than pixels. Thanks to the introduction of our novel alignment module, the model deals properly with two issues that are not handled satisfactorily by other transition models, namely object persistence and object identity. We show that the combination of an object-level loss and correct object alignment over time enables the model to outperform a state-of-the-art baseline, and allows it to deal well with object occlusion and re-appearance in partially observable environments.

Unsupervised Object-Based Transition Models for 3D Partially Observable Environments
Antonia Creswell
Rishabh Kabra
DeepMind
DeepMind London rkabra@deepiind.com
London
tonicreeseel1@@eeppinad
Christopher Burgess
Murray Shanahan
DeepMind
DeepMind
London cgburgasssddeepmind com
London
mshanahan@deepminnd com
Abstract
We present slot-wise, object-based transition model that decomposes a scene into objects, aligns them (with respect to slot-wise object memory) to maintain a consistent order across time, and predicts how those objects evolve over successive frames. The model is trained end-to-end without supervision using transition losses at the level of the object-structured representation rather than pixels. Thanks the introduction of our novel alignment module, the model deals properly with two issues that are not handled satisfactorily by other transition models, namely object persistence and object identity. We show that the combination of an object- level loss and correct object alignment over time enables model to outperform state-of of teffthe-art baseline, and allows it to deal well with object occlusion and re-appearance in partially observable environments.
1
Introduction
In spite of their well-documented ability to learn complex tasks, today's deep reinforcement learning agents are still far from matching humans at out-of-distribution generalisation or few-shot transfer
[6 21 23 Two architectural features commonly proposed to remedy this are (1) transition models
that enable the agent to internally explore paths through state space that it has never experienced (29 25 [2], and (2) compositionally structured representations that enable the agent to represent meaningful states that has never encountered [7]. These two features not exclusive; transition models that operate on compositionally structured representations are a potent combination, and these the subject of the present paper. Specifically, our focus is transition models that operate the level of objects, which are the most obvious candidates for the structural elements of representations likely to be useful for artificial agents inhabiting 3D worlds such
28
current models do not deal satisfactorily with object persistence (the concept that objects typically continue to exist when they are no longer perceptible (24) or with object identity (the concept that a token object at one time-step is the same token object at a later time-step [I]. As show, transition models that neglect object persistence tend to perform badly in complex, partially observable environments, while models that neglect object identity interactions) over time in a way that generalises to future time-steps. By proposing and incorporating novel module for aligning objects across time using slot-based memory, model handles both these concepts, and exhibits better performance as result.
While there has been progress in object-based transition models 34
35
unable to integrate information about single object (and its
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

t=
t=2
b=3
t=4
t=10
t=12
t=14
t= 16
t=
188
Imput
Target
Unroll
OAT (ours)
Unroll Unroll
Memory Neig
OP3 (Veerepaneni current
state-of-art model)
Comparing our model to baselines. Each model was trained with four input steps and to unroll for six steps, here we unroll for 15 steps. Our model, OAT, performs significantly better than OT (our model without the AlignNet) and current state-of-the-an model, OP3 [3] See Figure in the Appendix for additional OP3 roll-outs.
Figure 1:
An important feature of our transition model is that makes predictions and computes losses in representation space that is divided into objects. This contrasts with existing models that make predictions in an unstructured representation space [13 12). Making predictions and com- puting losses in an object-structured rep- resentation space facilitates learning, not only because the representation space lower dimensional than pixel space, but also because the model can exploit the fact that dynamics tend to apply objects a whole, which simplifies learning. How- compute prediction losses directly over distinct object representations (rather than first mapping predictions back to pix- els [34] 32 the objects in a predicted rep- resentation must be matched with those in the target representation, which again re- quires a proper treatment of object persis- tence and identity. The result is a model that can roll-out accurate predictions for significantly more steps than seen during training, outperforming sate-of-the-ar for comparable models.
m
z
M++
za
h
h
h
Transition module
Transition module
LOSS
za 1-1
a
z
a
Alignment module
Alignment module
Z-1 m
m
Scene decomposition and representation module
Scene decomposition and representation module
X
X
Figure 2: Encoding steps. The scene decomposition and representation module extracts K object represen- tations, 4, E RF from images, 2t R x3 These are aligned using the aligument module to obtain of R Aligned objects are fed the slot-wise transition model with the action, at and hidden state, (, h2) E Mx to predict the object representa- tions at the next time-step, of E M well the updated memory, m E XF We also depict the transition model loss, Transition model. computed between z and 2 (shown as LOSS in the figure).
To achieve this, model, Objects- Align- Transition (OAT), combines (1) scene decomposition and representation module, MONet 2]. that transforms a raw image into slot-wise object-based representation, (2) an novel alignment module which, with the aid of a slot-wise memory, ensures that each object is represented in the same slot time, even if has temporarily disap- peared from view, and (3) novel slot- wise transition model that operates on the object representations to predict future states. All three components are differentiable, and the whole model trained end-to-end without supervision. We evaluate the model on two sequential image datasets. The first is collected from pre-trained agent moving simulated 3D environment, where we show that our model outperforms the current state-of-the the-art object-beer bosect-bbaed based transition model (OP3 32 The second dataset is collected from a real

---

robot arm interacting with various physical objects, where we demonstrate accurate roll-outs over significantly longer periods than used in training. An ablation study shows that the model's 's success depends on the combination of correct object alignment across time and the use of transition loss over object-level representations instead of over pixels.
2
Our Model: Objects-Align-Traaniiion
Our model, Objects-Aligni Transition module, in this case MONet, with our slot-wise transition module, via our alignment module, which ensures that each slot the transition model receives objects corresponding to the same token object (or identity) across time. The whole model is trained end-to-end. We now provide details of each of these modules (see Figure
Aliaaaaaiinnnnnnniiiininnnn (OAT), combines scene decomposition and representation
2.1 Scene Decomposition and Representation Module
The input to OAT is sequence of RGB images, E [0, 1] with T sime- width, W, and height, H. We leave out the batch dimensions for simplicity. Each image in the sequence, Iz, is passed through a scene decomposition and representation moduld] in this case MONet, to obtain a slot-wise object representation, [21, occupying K slots, where 21, E R is an object representation vector with F features.
MONet consists of an attention module (a U-net H4, E [0, x] #xx1, for each slot and slot-wise VAF. The encoder of the slot-wise VAE is fed the predicted object segmentation masks and the input image, Zt, and outputs object representations, 2 "R The decoder reconstructs the masks, uu,,,. and each object's pixels, Tt,k. MONet's computations can be summarised by the following equations: URRt = NitentionNNttoor((t:) Zkett MONet_encoder(z;; (l.,)) and ji,.,, Tt
which predicts object segmentation masks,
(eONee__ccode((1)) (z)..
may be tempting to feed the wise object representations, [= k= at time, t, directly to slot-wise transition model. However, MONet's representations [2] are not stable across time, meaning that specific object may appear in different slots different times [35] (see Figure B3). This leads to two major problems when training slot-wise transition models on slot-wise object-based representations. Firstly, difficult to compute losses between predicted and target object slots. Computing slot-wise object-level loss requires us to know how the previous object representations (and thus the predicted object representations) correspond with the target object representations. Secondly, if object representations do not occupy consistent slots it makes it harder to integrate information about a single object (and its interactions with other objects) across time (35), and makes it harder to predict the reappearance of that object (as show in Section 4.2 To address the issue of slot stability we introduce Memory AlignNet, an alignment module.
2.2 Alignment Module
Our alignment module, Memory AlignNet, must play two key roles. The first is align objects across time, enabling us to compute slot-wise object-level losses for training. The second is to learn a slot-wise memory that encodes the history of each slot across time, allowing our transition model to operate effectively in partially observable environments by keeping track of objects as they go in and out of view. This is especially important for embodied agents that take actions in 3D environments, where the agent's 's looking around frequently causes objects to move in and out of its field-of-view.
We propose Memory AlignNet, which takes slot-wise memory, m. E X with M > K slots and the (stacked) output of the scene decomposition and representation module, 2, E R R returning the aligned object representations, 2 M To perform alignment, the Memory that aligns objects, 2t, with the memory, mi- This adjacency matrix allows to compute soft alignment 20.t = A, hard alignment, 2a A za,hard Hungarian( (A:) We will refer to 2 a.herd as z throughout the rest of the paper. Hungarian(-) denotes the application of the Hungarian algorithm, a non-different non-differentiable algorithm which computes permutation matrix given an adjacency matrix. The soft version of the alignment
AlignNet predicts an adjacency matrix, A₁ E R MM
We tried some variants where depended on 2 However these did not yield significantly beneficial results.
3

---

Alignment module losses (Memory AlignNet). The Memory AlignNet trained on a recon- struction loss between the softly aligned object representations, 2, 20,soft and the output of the transition model, z for the corresponding time-step. There are also regularisation losses includ- used to train the Memory AlignNet while the hard version of the alignment is passed to the slot-wise transition model. which has the same number of slots, M, the memory. Further details can be found in Section A.2 of the Appendix.
The hard-aligned objects may also be used update the memory using recurrent slot-wise model. For simplicity we use our transition model, (-).). to predict deltas for both the object representations, zf, and the memory, my. We will detail this in the next section.
2.3 Slot-Wise Transition Module
The transition model operates in both an encoding and an unrolling phase. During the encoding phase the transition model is fed aligned, observed object representations, 24, and actions, dg, to predict aligned object representations at the next time-step, =t+1 E RMMX.. During the unroll phase, the transition model is fed the predictions, 2f, and actions, de, from the current time-step predict the object representations, z[+1 at the next time-step.
More concretely in the encoding steps, our transition model, To()), takes aligned object representa- tions, of, and a hidden state, h1 E 3M** and predicts deltas for both the object representations, 24 and memory, my, The output of our transition model given by [A] Ar"), hg+1 = To at hf). The memory and the object representations are then updated as follows: M++1 = m. + A," and 2d
2141 = 2f + A, In the unroll steps, next step predictions are given by 2 4++1 = 2f + A where [A Ar",, h = (=( ag, h.). By using the transition model to predict deltas, 2 is aligned by default and does not need be re-aligned when unrolling the model.
A key feature of our transition model is that weights are shared across object slots and can be instantiated in many different ways. One simple option is independently to each slot, sharing weights between slots. An alternative, novel instantiation which found work well is first apply a transformer and Section[A.3 This allows the model to capture interactions with other objects (via the transformer) and integrate that information over time (via the SlotLSTMM, Further details are in Section A.2 of the Appendix.
a SlorLSTM; LSTM applied
KI and then the SlotLSTM (see Figure TQ
2.4 Training
When training OAT, we jointly learn the parameters of the scene decomposition and representation module, the aligmment module and the transition module. Each module has its own losses that contribute to downstream gradients and updates. For example, both the transition model and the aligmment module losses can influence the scene decomposition and representation module's 's updates. Let us define the losses for each module.
Scene decomposition and representation module losses (MONet). We train MONet using stan- dard MONet losses, |MONet. see Equation 3 in [2]. The first term in the MONet loss is scene reconstruction term. This is a spatial mixture loss parameterised by dbg and og which are the standard deviation used for each slot's component likelihood (for the first slot and remaining slots, respectively) that go into the mixture loss. The remaining terms are regularisation terms that 1) induce latent information bottleneck needed for good representation learning (scaled by B, the latent KL loss weight), and 2) ensure that predicted and target masks are similar (scaled by y, the mask KL loss weight). MONet losses computed only in the encoding phase and not used in the unroll phase. Additionally, MONet loss gradients do not affect the transition module or the alignment module weights.
ing an entropy loss, [H!()), on the adjacency matrix, A., which encourages values towards zero and one; and a loss that penalises columns that sum to more that one, avoiding the case where multiple objects are assigned to the same memory slot. The AlignNet loss, (AligaNer is defined as VeiggnNer = Ei || - =i =g. + +H (A,) + Σド' max(0, (ΣK Ar,. 1)), where T is the total number of encoding and unroll steps. We use t/ = 0.01 for experiments presented this paper.

---

Slot-wise transition module losses. The slot-wise transition model is trained with reconstruction loss between the outputs of the transition model, 2 and the aligned observations for the same time-step,
z.. The slot-wise transition model loss, Transition model is defined as Transition model
Ez || =f 1|. Importantly, compute loss directly between object representations without the need for decoding them to obtain pixels.
OAT is trained end-to-end to minimise IMONet + AAlignNet + Çl Transition model. using = 10 for all experiments presented in this paper.
3 Related Work
Our work addresses the challenging topic of learning abject-based transition models in 3D partially observable environments without supervision. We identify two phenomena that are not dealt with adequately in current object-based transition models, namely object persistence and identity. As consequence, current models are not able to perform well in partially observable environments, and are often trained without semantically meaningful losses.
Firstly, we acknowledge previous work on temporally extended scene decomposition and representa- tion models [35 [4] (typically used for video representation learning), object-based transition models [20]7 and action conditional object-based transition models such as OP3 (32) and C-SWM (1)]. We consider Veerapaneni al.'s 's OP3 32
be the current state-of-the-art object-based transition model since C-SWM requires privileged information about exactly which object an action was applied to, while OP3 (like our model) simply requires the action taken by the agent. OP3 alternates between refinement steps and dynamics (or prediction) steps. One drawback of OP3 is that its refinement steps require access to the observation, which means refinement cannot be applied when rolling out the model to unobserved time-steps. For this reason Veerapaneni et al. [22 only demonstrate unrolls for a limited number of steps.
3.1 Current Models are not Designed for Partially Observable Environments
The world that we operate in, and which we intend our agents to operate in, is partially observable. A severe limitation of existing object-based transition models [ [9] and temporally extended object scene decomposition and representation models [T) 14 is their inability to cope with partially observable environments. A promising approach proposed by He et al. [14] uses an external memory for object tracking (but does not use transition model). Their mechanism is different to that of our proposed alignment module, AlignNet, which we use to perform slot alignment in OAT. He al. 14 train their model, TBA, for reconstruction while AlignNet incorporates dynamics and is trained using prediction, meaning that AlignNet can resolve ambiguities using dynamic while TBA cannot. Additionally, TBA can only cope with static backgrounds and so it is not applicable here. Goyal et al. (8] also develop a general model with a notion of object persistence. However they only show results in 2D environments.
3.2 Current Models are not Trained Using Semantically Meaningful Losses
Another problem with current approaches to learning object- transition models is how they are trained. In most scenes, the background accounts for most of the pixels while the objects account for only small fraction them. So while it may tempting to train transition and video--rpresenntation models using a pixel-level loss [4 K5
32
20
17
[4], we show that it is preferable to compute
losses directly between predicted and target objects (see Section and Figure T.. Furthermore,
pixel-level losses require the object representations to be decoded into an image 34 35 is often computationally expensive.
32
which
An alternative way to find object-level losses is to compute a minimum assignment loss 4 employing the Hungarian algorithm, between predicted and target object representations. However, doing so can be problematic because you need to first compute a similarity matrix on which to apply the Hungarian algorithm. We show in Section 4.2 that computing a loss using the Hungarian algorithm (using L2
For example, when two visually similar objects collide with one another, the AlignNet can use dynamics to resolve which object is which after the collision.
5

---

1
Figure 3: Outputs of our scene decomposition and representation (MONet) module across time. MONet predicts K 10 latent object representation vectors {=r,)i= K each time-step. Here we visualise those vectors, in columns CI to C10, using MONet's decoder. Notice that objects switch slots across time. This makes it difficult (a) compute losses between predictions and targets and (b) integrate information about an object across time.
loss between all object pairs as the similarity matrix) leads to poor generalisation in transition models when performing longer unrolls than those seen during training.
Furthermore, at the start of training, comparing predictions with targets to compute a loss may not be meaningful, since predictions will start off essentially random. our approach, align the inputs with the targets, and condition the slot-wise prediction on the slot-wise history. This reinforces slot stability and allows us accurately predict changes, Ai, in each object representation.
Interestingly, Lowe et al. 22 use Deep Sets [36 to encode predicted and target object sets and use contrastive loss between encodings. Their approach has only been demonstrated for very simple datasets. Kipf et al. (C-SWM) [19] do compute losses their representation space. However, they avoid the object correspondence problem because they extract representations spatially, keeping the spatial ordering and using data where object movement was limited. Their approach is unlikely to scale well to partially observable environments with significant movement of objects across the field of view, especially when being trained to predict multiple steps into the future.
OP3 322 and other models 35 attempt induce weak, implicit object alignment by conditioning predicted object representations on those from the previous time-step. This does not guarantee alignment (or slots with consistent identity) over time, especially in partially observable environments. Moreover, in contrast to OAT, they do not use their implicit alignment for computing losses.
Finally, we note that while there are existing models that learn to predict future states, given actions, directly from pixels [13] 12]. we have focused primarily on related work that predicts the future states of objects, because we are particularly interested in developing models that may support future work towards object-based based agents.
Experiments and Results
In this section we (1) demonstrate OAT's performance in the 3D Playroom environment and compare the current state-of-the-ar object-based transition model, OP3 (32) (Section 4.1
[
16
investigate the benefit of alignment and object-level losses when training transition models (Section 4.2). and (3) apply OAT in real world robotics environment and show it accurately predicts both the motion of the robotic arm and its physical interaction with objects (Section
4.1 Objects-Align- Transition Results in Playroom
We train and test our model, OAT, on data collected by an agent taking actions a Playroom environment [6]. Each procedurally generated room in the dataset contains between 10 and 45 objects from 34 different classes 10 different colours and three different sizes. dataset of observation-action trajectories, (Iz, o,) is generated by agent taking actions according to learned policy in procedurally generated room. We collect 100, 000 trajectories with a 7:2:1 train:validation.test split. The top row of Figure 5/oows example trajectory.
We train OAT with four encoding steps (i.e. the model sees four observations for the first four time-steps) and six unrolling steps. MONet outputs K object representations, ( We

---

slot
slot
slot 2
slot
slot 4
slot
slot
slot 7
slot
slot
slot 10 slot
Σ
Figure 4: Aligned inputs (first time-steps) and targets (last 6 time-steps). Our alignment module outputs object representation vectors which visualise here using MONet's entity decoder. Most of the objects are now in consistent slots across time, making easier to compute losses directly between object representations using simple L2 loss. Notice that while MONet outputs 10 slots, the AlignNet has 12 slots. Additional time-steps shown in Figure 11 in the Appendix.
K 10 objects, with F 32 features, and a memory with M slots. Figure visualises the MONet outputs for the first three time steps. Our model achieves good segmentation (see the Adjusted Rand Index, ARI, in Table [ segmentation metrics). Importantly, notice that the slots are not stable across time. For example, the purple object in slot C7 at t = 0 switches slots at t = 2. 0 switches slot C2 at 1 and slot C4
Similarly the yellow object in slot C5
2.
The object representations, 2, R K F), output by MONet are fed to the alignment module. The outputs of the alignment module, 2 E 3R(Mx F), are visualised in Figure 4 The outputs of the aligmment module object representations, and Figure His visualisation these representations using MONet's decoder. Our alignment module successfully keeps objects in consistent slots across time.
The output the alignment module is used to generate both inputs and targets for training the transition model. The transition model predicts latents, 2f E xF] In Figure 5 we visualise roll-outs from our model using MONet's decoder; the reconstructed scene visualisations as the mask weighted sum of the each slot's pixels, Σx (, While the model is trained to unroll for six steps, here unroll for 15 steps. The transition model only sees the first four frames. In the top example, we see that the model is able to predict the appearance of the avatar well, and in general notice that the model, given the agent's actions, able to predict the position of objects well and without the representations degrading. (For comparison to baselines see Figure T some of the examples, the targets appear to have more objects than those seen in the unroll. This because the model has only seen the first four frames and has not seen those other objects in the room, and therefore does not have enough information predict where unseen objects will appear.
generated
Ianget
Unai
target
Ingget
Figure 5: Unrolling OAT for more time-steps than seen during training. The slot-wise transition model makes predictions in latent space (one latent per object), which we visualise using MONet's decoder. The slot-wise transition model was trained with 4 inputs steps and to unroll for 6 steps. Here the model takes the first 4 steps as input and is unrolling for 15. On the top row at f = 16, it is impressive that the model able predict the appearance of the avatar's base (the red circle with the yellow stripe) when looking down. (For per-object visualisation see Figure 13 in the Appendix.)
7

---

To conclude these experiments, we compare our model to the current state-of-the-an object- object-based transition model, OP3 [32 and our ablated model, OT, without the alignment module in Table and Figure [ Qualitatively, from Figure ① we see that OAT significantly outperforms both baselines. Unrolls using OT, trained without alignment, lead to objects merging towards grey cloud in the middle of the frame. Unrolls from OP3 lead to objects fading into the background (we explore the cause of each of these pathologies in the next Section 4.2). Results were consistent for each model across multiple runs.
To quantitatively compare models we consider three metries: Encoding ARI, Unroll Pixel Error and Unroll ARI. The Encoding ARI (Adjusted Rand Index. [18 26 measures the accuracy of the object segmentation masks learned by the scene decomposition model. The Unroll Pixel Error and Unroll ARI evaluate the quality of the transition model's 's unrolls. To compute each of these we decode the latents, z predicted by unrolling the transition model, to produce images, æy and masks, Mi.. The Unroll Pixel Error is the mean-squared-emvee between the ground truth images and xy for the unroll steps only. The Unroll ART is the accuracy of the decoded masks, 1& compared with the ground truth masks. For both ARI scores we exclude background pixels from the score since accurate decomposition of the objects is the main concern here. The Unroll ARI is more meaningful evaluation of the unrolls than the Unroll Pixel Error since is not affected by the background which often accounts for most of the pixels. Note that we only use ground truth masks for evaluation purposes.
The results in Table furthhr demonstrate the crucial role that alignment (AlignNet) plays when performing unrolls. Without alignment the Unroll ART is significantly lower because we are not able to compute a semantically meaningful object-level loss for training. Additionally, we see that our model significantly outperforms OP3 on all metrics (see Figure In light of the results in Table I in the next section we will more concretely look at the role of alignment and the object-level loss (between object representations) when training transition models.
Encoding ARI
Unroll Pixel Error
Unroll ARI
OAT (ours)
0.62
0.0121
0.42
OT (ours, no AlignNet)
0.59
0.0143
0.12
OP3
0.32
0,0132
0.33
Table 1: Comparing
model OAT to baselines. We trained three OAT and OT models (three seeds each), and report the ARI score, Unroll Pixel Error and Unroll ARI score for the model with the best Unroll ARI score. Results for OP3 were obtained similarly, except that ten models were trained (ten seeds), since more variance was observed in the OP3 results (see Appendix D for details).
4.2 What Matters in Object-Centric Transition Models?
In this section demonstrate the need for both (1) alignment, which ensures that each slot in the transition model receives the same object consistently across time, and (2) object-level loss, computed between predicted and target object representations, rather than a pixel-level loss. We also compare different transition model cores and find a transformer followed by a slot-wise LSTM to be best.
In this section (Section 4.2 only, we use a MONet model that is trained using ground-truth masks instead of learning its own masks. We do this here directly analyse the benefits of aligned object representations and object-level losses for the transition model without other confounds. Using ground-truth masks allows us know each object's true identity across time and directly measure the role of alignment in transition models without confounding errors from AlignNet or MONet's segmentation quality. We use a pre-trained MONet (with fixed weights) to compute object representations to keep the model similar to the full end-to-end model described in the rest of the paper.
We train slot-wise transition models under four different conditions: feeding aligned or unaligned object representations (and target) object representations the transition mode and training the
'For the unaligned inputs. we shuffle the order of the ground-truth masks E W,H,5, that are fed to MONet along the K axis. For aligned inputs, ensure that each slot contains consistent object across time.
8

---

t=10
1=11
t=27
t=30
t=31
[=40
t=48
Targets
Aligned inputs, object-level loss
Aligned inputs, pixel-level loss
Unaligned inputs, object-level loss
Unaligned inputs, pixel-level loss
Figure 6: Comparing slot-wise transition models trained with aligned vs. unaligned input latents and pixel-level vs. object-level loss. These models were trained with 8 inputs steps and 12 unroll steps. Here they are being unrolled for 40 time-stepy The output of the transition model is slot-wise object-based representation for each time-step. Here, we visualise the object representation vectors using MONet's is decoder. Only models trained with aligned inputs were able to predict the reappearance of the chair at = 30. Only the model trained with aligned inputs and object-level loss is able to predict the appearance of the avatars "base". In models trained with pixel-level loss predictions becomes very blurred.
Reporting object-level error for models trained using:
Object-level loss
Pixel-level loss
Aligned inputs
0.0929 1 0.00469
14.4 18.4
Unaligned inputs
> 20.3 + 0.757
3312 + 4050
Table 2:
The role of alignment and object-level loss when training transition models. This table reports the std latent error between predicted and ground truth latents (for 5 seeds). For models trained using unaligned input latents we compute the object-level loss using the Hungarian algorithm which is a lower bound the actual value. Models trained with pixel-level loss become very unstable resulting in high variance between runs. Individual seeds are shown in Figure
model with pixel-level or object-level loss. Results in Table 2 clearly demonstrate the benefits of (1) using aligned object representations for training transition models and (2) training transition models with object-level loss. These results are critical for the futuree development of object-based transition models and, in particular, demonstrate the need for alignment. Note that for the experiments using the unaligned inputs computed the object-level loss using the Hungarian algorithm which is a lower bound estimate of the true object-level loss (since the minimum assignment in L2 may not be the correct assignment).
Figure coomppres models trained under the four conditions listed above. We see that models trained using aligned inputs (and targets) are able to predict the re-appearance of objects while those without aligned latents are not (Figure 6 t = 30). We also see that models trained with an object-level loss and without alignment deteriorate quickly. Additionally, for models trained with pixel-level loss notice a "ghosting" effect where objects fade into the background across time, similar to the effects seen in the OP3 results (Figure I7p. Pixel-level loss leads to this "ghosting" effect because most pixels in the observation are background pixels and so the background pixels dominate the loss.
We found these results to be consistent across multiple runs (see Figure 9 in the Appendix) and different choices for slot-wise transition model architectures. For the results shown this section we used a transformer with a slot-wise LSTM. Figure 10 in the Appendix, compares object and pixel errors for different transition module cores. Training with object-level loss and aligned inputs, the transformer with slot LSTM worked best.

---

t=0
t=2
t=3
t=4
t=10
t=17
t=33
t=44
t=54
t=71
t=103
Target
Uhroll
Figure 7: Unrolling for many more steps than seen during training, OAT is trained to take four input steps and unroll for six steest Here we demonstrate the model unrolling for 100 time-steps. Impressively, we see that the model learns both to predict the behaviour of the robot arm accurately, as well as how the arm interacts with the objects. We also see that the model able to unroll for significantly more steps than seen during training. Note that our model makes predictions latent space; here we are visualising those latents using MONet's 's decoder. Additional results in Figure 15 in the Appendix.
4.3 Application to Robotics
Here we demonstrate the application of OAT to a real world dataset B of robot trajectories. These trajectories involve robot arm interacting with three objects of varying shapes and colours. The dataset is particularly challenging because our model must learn to predict not just the motion of the robot arm, given the actions, but also how the arm interacts with objects requiring some understanding of intuitive physics.
We train OAT with four input steps and to unroll for six steps. Our model achieves excellent segmentation results, shown in Figure 14 of the Appendix. Figure [ shows impressive results obtained when unrolling the OAT model for significantly more steps than seen during training. What more, in the third sample from the top. we see that the model correctly predicts the reappearance of the red object after it had been fully occluded at t = 3 for 51 time-steps, reappearing fully at t 54 in both the prediction and the target frame. Our model is able to accurately predict the reappearance of objects, even after long-term occlusion, because explicitly captures the history of each object, endowing the model with a notion of object persistence.
Conclusion
We presented Dojects-Align-- Transition (OAT), an object-centric transition model that combines scene decomposition and object representation module (MONet) with slot-wise transition module, via an alignment module. The alignment module plays two key roles. Firstly, it ensures that the slot-wise transition model receives slot-consistent object representations across time. Secondly, allows us to compute object-level loss rather than a pixel-level loss which is commonly used when training transition models.
In an ablation study, we demonstrated the essential role that alignment and object-level losses play when training transition models. Additionally, we significantly outperform existing state-of-the-an object-centriy transition models in 3D partially observable environment, and we applied our model robotics dataset, predicting many steps further into the future than seen during training.
There is room to improve our model in further work, for example by making stochastic predictions about the future, and by better modelling the uncertainty about the objects in the environment. Meanwhile, our work paves the way for future object-centric agent research, for example, enabling agents to plan over future trajectories in object representation space.
Acknowledgments and Disclosure of Funding
We would like to thank Adam Kosiorek for feedback on early drafts of the paper as well as Arunkumar Byravan, Dushyant Rao, Yusuf Aytar and Markus Wulfmeier for their help running experiments on robotics data. This work was carried out within DeepMind, who are wholly owned by Alphabet, the parent company of Google.
10

---

References
[1] R. Baillargeon, E. S. Spelke, and S. Wasserman. Object permanence five-month-old infants. Cognition, 20(3):191-208 1985.
[2] C. P. Burgess, L. Matthey, N. Watters, R. Kabra, Higgins, M. Botvinick, and A. Ler- chner. MONet: Unsupervised scene decomposition and representation. arXiv preprint arXiv: 1390, 2019.
[3] S. Cabi, S. Gómez Colmenarejo, A. Novikov, K. Konyushkova, S. Reed, R. Jeong, K. Zolna, Y. Aytar, D. Budden, M. Vecerik, et al. Scaling data-driven robotics with reward sketching and batch reinforcement learning. arXiv, pages arXiv-1909, 2019.
[4] D. Co-Reyes, R. Veerapaneni, M. Chang, M. Janner, C. Finn, J. Wu, Tenenbaum, and S. Levine. Discovering, predicting, and planning with objects. International Conference Machine Learning Workshop, 2019.
[5] Z. Dai, Z. Yang, Y. Yang, J. Carbonell, Q. V. Le, and R. Salakhutdinov. Transformer-xl: Attentive language models beyond fixed-length context. arXiv preprint 1rXIv: 1901:02861.0 2019.
[6] M. Garnelo, K. Arulkumaran, and Shanahan. Towards deep symbolic reinforcement learning. arXiv preprint arXiv: 1609,05588, 2016.
[7] M. Garnelo and M. Shanahan. Reconeiling deep learning with symbolic artificial intelligence: representing objects and relations. Current Opinion in Behavioral Sciences, 29:17-23, 2019.
[8] A. Goyal, A. Lamb, P. Gampa, P. Beaudoin, S. Levine, C. Blundell, Y. Bengio, and M. Mozer. Object files and schemata: Factorizing declarative and procedural knowledge in dynamical systems. arXiv preprint arXiv:2006. 16225, 2020.
[9] A. Graves, M. Liwicki, S. Fernández, R. Bertolami, H. Bunke, and J. Schmidhuber. A novel connectionist system for unconstrained handwriting recognition. JEEE transactions on patter analysis and machine intelligence, 3 31(5):855-88) 868, 2008.
[10] K. Greff, R. L. Kaufmmnn, R. Kabra, N. Watters, C. Burgess, D. Zoran, L. Matthey, M. Botvinick, and A. Lerchner. Multi-object representation learning with iterative variational inference. In International Conference Machine Learning, pages 2424 2433. PMLR, 2019.
[11] K. Greff, A. Rasmus, M. Berglund, Hao, H. Valpola, and Schmidhuber. Tagger: Deep unsupervised perceptual grouping. In Advances in Neural Information Processing Systems, pages 4484 4492, 2016.
[12] D. Ha and J. Schmidhuber. World models. arXiv preprint 1013, 2018.
[13] D. Hafner, T. Lillicrap, Ba, and M. Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv: 19X.v:9912.016000 2019.
[14] Z. Hc, Li, D. Liu, He, and Barber. Tracking by animation: Unsupervised learning of multi-object attentive trackers. arXiv preprint
2018.
[15] F. Hill, A. Lampinen, R. Schneider, S. Clark, M. Botvinick, J. L.. McClelland, and A. Santoro. Environmental drivers of systematicity and generalization in a situated agent. In International Conference on Learning Representations, 2020.
[16] Hill, S. Mokra, N. Wong, and T. Harley. Human instruction-follloiing with deep reinforcement learning via transfer-learning from text. arXiv preprint arXiv:2005.09338, 2020.
[17] J.-T. Hsich, B. Liu, D.-A. Huang. L. F. Fei-Fei, and J. C. Niebles. Learning to decompose and disentangle representations for video prediction. In Advances in Neural Information Processing Systems, pages 517-526, 2018.
[18] L. Hubert and P. Arabie. Comparing partitions. Journal of classification, 2(1):193-213 1985.
11

---

[19] T. Kipf, E. van der Pol, and M. Welling. Contrastive learning of structured world models. arXiv preprint arXiv: 1911. 2019.
[20] A. Kosiorek, H. Kim, Y. W. Teh, and Posner. Sequential attend, infer, repeat: Generative modelling of moving objects. In Advances in Neural Information Processing Systems, pages 8606-8616 2018.
[21] B. M. Lake, T. D. Ullman, J. B. Teeenbaum, and S. Gershman. Building machines that learn and think like people. Behavioral and Brain Sciences, 40, 2017.
[22] S. Löwe, K. Greff, R. Jonschkowski, A. Dosovitskiy, and T. Kipf. Learning object-centric video models by contrasting sets, arXiv preprint arXiv:2011.10280 2020.
[23] G. Marcus and E. Davis. Rebooting Al: Building Artificial Intelligence We Can Trust. Ballantine Books Inc., 2019.
[24] Z. W. Pylyshyn. The role of location indexes in spatial perception: A sketch of the finst spatial-index model. Cognition, 32:65-995 97, 1989.
[25] S. Racanière, T. Weber, D. Reichert, L. Buesing, A. Guez, D. Jimenez Rezende, A. Puig- domènech Badia, O. Vinyals, N. Heess, Y. Li, R. Pascanu, P. Battaglia, D. Hassabis, D. Silver, and D. Wierstra. Imagination-augmented agents for deep reinforcement learning. In Advances in Neural Information Processing Systems 30, pages 5690-5701, 2017.
[26] W. M. Rand. Objective criteria for the evaluation of clustering methods. Journal of the American Statistical association, 66(336):846 850, 1971.
[27] O. Ronneberger, Fischer, and T. Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and computer-assisted intervention, pages 234-241. Springer, 2015.
[28] M. Shanahan, B. Beyret, M. Crosby, and L. Cheke. Artificial intelligence and the common sense of animals. Trends in Cognitive Sciences, 24(11):862-8) -72, 2020.
[29] D. Silver, A. Huang, C. Maddison, A. Guez, L. Sifre, G. Van Den Driessche, Schrittwieser, I. Antonoglou, V. Panneershelvam, M. Lanctot, et al. Mastering the game of Go with deep neural networks and tree search. Nature, 529((587):484
2016.
[30] S. van Steenkiste, M. Chang, K. Greff, and J. Schmidhuber. Relational neural expecta- tion maximization: Unsupervised discovery of objects and their interactions. arXiv preprint arXiv:1802.10353, 2018
[31] A. Vaswani, N. Shazeer, N. Parmar, Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. In Advances in neural information processing systems, pages 5998-6009 6008, 2017.
[32] R. Veerapaneni, J. D. Co-Reyes, M. Chang, M. Janner, C. Finn, J. Wu, J. Tenenbaum, and S. Levine. Entity abstraction in visual model-based reinforcement learning. In Conference Robot Learning, pages 1439-1456. PMLR, 2020.
[33] X. Wang, R. Girshick, A. Gupta, and K. He. Non-local neural networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 7794-7803, 2018.
[34] N. Watters, L. Matthey, M. Bosnjak, C. P. Burgess, and Lerchner. Cobra: Data-efficient masel-based rl through unsupervised object discovery and curiosity-driven exploration, 2019.
[35] M. A. Weis, K. Chitta, Y. Sharma, W. Brendel, M. Bethge, A. Geiger, and A. S. Ecker. Unmasking the inductive biases of unsupervised object representations for video sequences, 2020.
[36] M. Zaheer, S. Kottur, S. Ravanbakhsh, B. Poczos, R. R. Salakhutdinov, and A. Smola. Deep sets. In Advances in neural information processing systems, pages 3391-3401, 2017.
12


Reviews:
[{"review_id": "grCcBf1fymZ", "reviewer": "Reviewer_NjnN", "summary": "The paper proposes an architecture for unsupervised object discovery and tracking and shows competitive results on two datasets consisting of video of 3D scenes.", "questions": "", "limitations": "", "rating": 6, "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": "The paper proposes an architecture for unsupervised object discovery and tracking. The paper proposes a solution to the problem that frame-based object decomposition architectures outputs sets of objects vectors, which are not aligned across time.\n\nThe architecture consists of three modules, that are trained jointly, end-to-end. The three modules are \n\n1) A \"standard\" MONet image object decomposition module, which decomposes images into sets of object vectors\n2) An alignment module \"AlignNet\", which outputs an alignment matrix A, that aligns new object vectors, z_t with previously seen object vectors.\n3) A transition module that predicts new object vectors given past ones.\n\nThe main novel part in my understanding is the addition of the AlignNet which serves as a kind of translator between MONet and the transition module. In a sense, it's quite similar to [1], except that [1] doesn't have an explicit alignment module, but rather has a big Transformer model that internally learns to perform the needed alignment between objects between timesteps. I'd encourage the authors to discuss the relation to [1] in their paper.\n\nThe paper is fairly well written. There are a few things that could be made more clear. I would appreciate an algorithm describing how the many parts are trained. It's easy to get lost in the text, so summarizing it in pseudo-code would be beneficial. \n\nIt wasn't clear to me how the model handled disappearing objects. In the appendix, it seems the object remembers an occluded object for many timesteps. While this is impressive, it does make me wonder if the model is capable of forgetting objects. There was some text in the appendix on how missing objects are handled, but this wasn't clear to me.\n\n[1] - Ding, David, et al. \"Object-based attention for spatio-temporal reasoning: Outperforming neuro-symbolic models with flexible distributed architectures.\" arXiv preprint arXiv:2012.08508 (2020).", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "VDdbMOiFJY8", "reviewer": "Reviewer_ayAi", "summary": "Current state-of-the-art unsupervised scene decomposition models cannot be directly applied to learn a dynamics model from video. This is because the order of decomposed objects is not consistent across video frames. This paper proposes an alignment module by predicting an affinity matrix that map an orderless output to an ordered output, driven by the consistency loss of the aligned output between timesteps. It can be naturally used to learn the transition model in the level of object, rather than the global scene. They demonstrate the results both quantitatively and qualitatively on synthetic and real world robotics datasets.", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "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": "Strength:\n\nThe major problem this paper aims to solve is important. The authors propose a surprisingly simple yet effective method to deal with it. I appreciate the simplicity. Meanwhile, in the related work section, the authors make a detailed comparison with previous work, which is helpful to identify the contribution.\n\nWeakness and Suggestions:\n\n1) I have concern about whether the contribution is enough for Neurips. Although the AlignNet is novel and effective, the other two parts of the proposed OAT (Objects-Align-Transition) model are adopted from previous work: the object model is a MONet and the transition model is a simple LSTM. Therefore, I would highly suggest adding experiments on planning / model-based RL tasks for two reasons. Firstly, the main compared work, OP3, did such experiments. It can make a better and clearer argument if the proposed prediction method is better when applied in the RL tasks. Secondly, as the authors argue, one of the most important usage of prediction model is the planning tasks, and showing the planning performance can make a greater impact.\n\n2) The writing needs to be improved. \nFrom the technical perspective, the authors seems to assume the readers are very familiar with MONet. For example, in section 2.1, it is better to give a concise and self-contained summary of MONet. The current form can easily make readers lost.\nFrom the language perspective: \n- It is mentioned many times the loss is applied at the level of object-structure representation rather than pixels. I can vaguely understand what the authors want to express but I think this sentence itself is confusing: which “pixel-level” method are the authors referring to? If the method is MONet and the loss computation refers to the loss in output space, it seems to be incorrect because the proposed method also use pixel-level loss in the objects model. If the loss computation refers to the loss in the feature space, then which previous work compute feature space loss in the ‘pixel-level’ (or the global level)?\n- Some typos: (line 356) out performs; (line 204) because [you] need \n\n3) In Figure 5, how can the model predict the red object after t=16? It seems the model has no input information about the red object. And What if there is a new object appearing in the scene?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I didn't find any potential negative societal impact of this work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6.5 hours"}, "extra_text": {}}, {"review_id": "DnEqyKF6Pz", "reviewer": "Reviewer_ApBk", "summary": "The paper proposes an object-centric transition model that can learn to decompose scenes into object slots and predict action-conditioned slot-wise dynamics end-to-end without supervision. The model has a novel memory-based alignment module that aligns the unordered input object slots with a slot-wise object memory at each time step, so that each memory slot can capture the same object across time and integrate information consistently. This alignment module is trained with an object-level loss rather than a pixel-level loss. The model is evaluated on two datasets featuring partially observable 3D environments. It is shown that the model can outperform OP3 in terms of segmentation quality, can unroll for far more steps than the training regime, and can predict re-appearance of objects. Ablation study shows that the alignment module and the object-level loss are both crucial.", "questions": "", "limitations": "", "rating": 4, "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": "**STRENGTHS**\n\n- Object permanence plays an important role in human cognition, and is a desirable ability for agents to have so as to effectively interact with 3D environments from first-person views. The paper demonstrates this ability by showing that the learned model can predict re-appearance of objects even after a long period of occlusion.\n- The alignment module is well motivated by experiments showing that MONet outputs inconsistently ordered object slots and that with unaligned inputs the transition model performs significantly worse.\n- The paper is well written and easy to follow.\n\n**WEAKNESSES**\n\n- The quantitative evaluation seems inadequate:\n    - The ARI seems to be computed independently for each time step, which does not measure the slot consistency across time. While there are qualitative results showing the consistency, it would be better to also quantitatively show that the alignment module indeed achieves consistency.\n    - It is shown qualitatively that the model can roll out many steps into the future. However, it would be better to also plot a curve showing quantitatively how the performance changes with increasing number of rollout steps.\n- The paper failed to discuss the novelty and strength over some relevant work on unsupervised object-centric video models [1-4]. In particular, SCALOR[3] shows ability to maintain consistency, roll out stochastically, and deal with dynamic backgrounds, even when there are around 50 objects. In comparison, the proposed model is only shown to keep track of at most $M=12$ objects. While scalability is not the main focus here, this raises concern about the significance of the paper. Below are two possible ways I can think of to strengthen the paper:\n    - Compared to SCALOR, the strength of this paper appears to be that it is action-conditioned and it can predict re-appearance of objects. However, it seems straightforward to extend SCALOR to be action-conditioned, and without experiments, it is hard to tell whether SCALOR will fail to detect objects or predict re-appearance of objects on the datasets in this paper. Hence, I think the paper can be improved by including action-conditioned SCALOR as a baseline.\n    - Another potential advantage of the proposed model is that it may be able to recognize the same object when it re-appears at an encoding step. SCALOR cannot do this since the propagation-discovery models will treat objects that re-appear as new objects during inference. Notice that this is more challenging than predicting the re-appearance during rollout, because this requires the memory to indeed not forget the object even after multiple steps of memory updates during which the object is not seen (L128-129 seem to suggest that the memory is fixed during rollout). Therefore, it would be impressive to show results with more (e.g. ~20 in Playroom) encoding steps.\n\n**MINOR COMMENTS**\n\n- It is unclear how many unroll steps are used in Table 1. Is it 6 or 15?\n- Footnote 3 is missing (L307).\n\n\n[1] Exploiting Spatial Invariance for Scalable Unsupervised Object Tracking (AAAI 2020).\n\n[2] Structured Object-Aware Physics Prediction for Video Modeling and Planning (ICLR 2020).\n\n[3] SCALOR: Generative World Models with Scalable Object Representations (ICLR 2020).\n\n[4] Learning 3D Object-Oriented World Models from Unlabeled Videos (OOL Workshop at ICML 2020).", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The model may have scalability issues when there are many objects, and it is unclear whether the model can handle object persistence when there are many more encoding steps.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "ABY-ENDU0Q", "reviewer": "Reviewer_sjbW", "summary": "This paper introduced a novel way to learn transition models for deep reinforcement learning agents in an unsupervised way, to improve generalization in unseen (out-of-distribution) environments. The main idea is to build a transition model not for individual pixels in the image, which is hard to model or learn, but for a higher-level object representation, because transitions of objects are much more natual and simpler to model and learn for typical 3D environments. To extract object instances from the scene in an unsupervised way, the method ulitizes the existing MONet model (Multi-Object Network) [2] and adds two components:\n\n(1) An alignment module: \"Memory AlignNet\", which associates frame-by-frame object instances from MONet over time to guarantuee object persistence and identity\n\n(2) A transition module: Transformer + SlotLSTM, which learns to predict transitions of objects, given an action.\n\nWith the presented combination of modules, it is possible to roll out future states (given a sequence of actions) more precisely than with existing object-based transition models, specifically in partially observable environments, where knowledge of object persistence and identity is important.\n", "questions": "", "limitations": "", "rating": 8, "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": "Pro:\n- The paper is generally very well written, the line of argumentation was always clear and the pictures support the reading and message.\n- The method outperforms the state-of-the-art baseline (OP3)\n- Experiments generally support the arguments well, e.g. why good alignment is important and why a semantically meaningful object-level loss outperforms pixel-level loss\n- The results are relevant for the wider RL community\n\nCon:\n- Minor: From the very beginning of the paper the term \"slot-wise\" is used many times, starting in the abstract. To me the term \"slot\" wasn't clear at all, so maybe it's worth introducing it early in the paper to avoid confusion. I only understood it after reading the original MONet paper.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, limitations have been discussed briefly in the conclusion.", "review_text": "", "extra_scores": {"ethical_concerns": "none", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents work on learning an object-centric transition model.  The reviewers appreciated the well-written paper and the proposed alignment module.  However, concerns were raised regarding the quantitative evaluation and novelty with respect to related papers.  The reviewers considered the authors' response in their subsequent discussions.  The reviewers engaged in discussion.  While the problem area is important, the comparisons with respect to other solutions (including SCALOR), the lack of demonstration of object permanence, and overall contribution of the paper remained concerns of the reviewers.

Based on this feedback the paper is not yet ready for publication in NeurIPS.

Decision: Accept (Poster)