COMPAS: Compose Actions and Slots in Object-Centric World Models

Abstract

In this paper, we propose a novel approach, COMPAS (COMPose Actions and Slots), which leverages the strengths of state-of-the-art object-centric approaches for modeling the dynamics of an environment. Our method encodes the environment’s state into symbol-like, object-centric representations, known as slots, where each slot corresponds to an individual object. This approach offers a structured and interpretable way to model complex environments by combining slots with action representations for accurate next-state prediction. The primary contribution of our work is an efficient world model with a dynamics predictor capable of predicting accurate trajectories in action-dependent environments. Additionally, our slot extractor module enhances the predictive capabilities by extracting deterministic slots that remain consistent both within a single trajectory and across episodes. Unlike slots sampled from a trainable distribution, deterministic slots are generated from a single trainable parameter together with slot positional embeddings. This design improves the consistency across episodes, which in turn leads to more accurate dynamics prediction. We present a comprehensive evaluation of our approach in various environments, demonstrating that our proposed method outperforms competing models in environments with discrete and continuous action spaces.

1. Introduction

The ability to generalize compositionally is the key to learning faster, solving new problems, and understanding new concepts with limited experience [1]. The difficulty in compositional generalization is caused by the so-called binding problem [2]—the inability of modern artificial neural networks to dynamically and flexibly bind information distributed over the network, which arises in learning on unstructured input data.

A possible solution to this problem could be the use of symbol-like representations. Such representations can be slot representations [3], where, instead of a single latent embedding, the input data are encoded by a set of such embeddings (slots). Slots compete with each other to describe a portion of the input data. Such representations have been successfully used for object-centric tasks such as set property prediction and object detection [3] in an image, learning visual dynamics from video [4], image generation [5], and unsupervised object-centric representation learning for real-world data [6].

World models are essential in embodied machine learning tasks such as robotics, autonomous vehicles, game AI, and video synthesis [7,8,9]. They help agents to understand, predict, and interact with dynamic environments. However, modeling high-dimensional, continuous, and time-varying data remains difficult. Traditional models represent the environment holistically, without distinguishing individual objects [10,11], which limits their ability to capture object relationships [12] and evolving dynamics [13]. These models often struggle to generalize and scale to complex, multi-object scenes. In contrast, object-centric world models offer interpretable, factored representations that improve generalization, scalability, and data efficiency by representing scenes in terms of individual objects and their properties.

Previous object-centric world models typically address only part of the overall problem. On the one hand, many slot-based extractors [14,15,16,17] struggle to maintain episodic and long-horizon consistency, since the same object may be assigned to different slots across time or across episodes. On the other hand, existing dynamics models [4,18] often do not provide an efficient action-conditioned transition mechanism for accurate prediction in action-dependent environments. Our main goal in this work is to address both limitations by learning consistent object-centric slot representations and combining them with an efficient transition model for future-state prediction.

In this paper, we propose a novel approach, COMPAS (COMPose Actions and Slots), to encode these dynamics, focusing on object-centric world models that integrate actions and slots. The main contributions of our work are as follows:

• We introduce a video slot extractor with a deterministic slot initialization procedure that enforces slot consistency both within trajectories and across episodes. Consistent slot acquisition is illustrated in Figure 1, with additional visualizations provided in Appendix C.
• We apply a temporal consistency block that enforces trajectory consistency in a parallel manner, significantly speeding up inference and training compared to recurrent approaches.
• We utilize stacked slot attention blocks, enabling scalability and more efficient inference and training compared to the original iterative approach.
• We present an efficient transition model that incorporates causal slot information, slot interactions, and actions to predict next states effectively.

2. Related Works

2.1. Object-Centric Learning

Many modern object-centric representation models are based on the slot attention module [3]. It implements an iterative attention mechanism, based on soft k-means clustering, for autoregressive slot refinement. Recent improvements of this method are based on better optimization [19], learnable slot initialization [20], slot structure augmentation [21], and improving the underlying clustering algorithm [22] or simplifying it [23]. However, these methods use simple convolutional neural network (CNN) encoders and decoders, which result in worse image reconstructions. The authors of SLATE [5] proposed to use a discrete latent space, which is extracted from the image by dVAE [24]. The computed slots are then passed through a transformer that predicts the latent token. This token is used in the dVAE decoder for image reconstruction. This results in better-quality reconstructions than in other models.

Classical object-centric models encounter issues with slot permutations. Specifically, when the same image is passed through the model multiple times, the same object may be allocated to a different slot each time. To address this problem, the SAVI [15,16] and STEVE [14] models propose the use of a predictor model that predicts the initialization of slots for use in slot attention for the subsequent frame. This predictor can be implemented as either a multi-layer perceptron (MLP) or a transformer.

Recent progress in object-centric learning has been closely linked to the adoption of pre-trained image transformers. Notably, DINOSAur [6] leverages pre-trained DINO transformers [25]. The primary training objective in this paradigm is to reconstruct patched-feature maps from slots using an MLP as a decoder. An extension of this method, VideoSAur [17], adapts DINOSAur for video data. However, achieving slot consistency within a trajectory through predictor models is challenging due to overfitting and gradient vanishing caused by the recurrent nature of training. Additionally, the iterative processing of slot attention models limits their inference speeds and scalability.

2.2. Object-Centric World Models

Previous object-centric world models have used other algorithms besides slot attention for object state extraction. SQAIR [26] uses a special detection model to detect objects in the scene and track their trajectories through RNNs. SCALOR [27] and SILOT [28] are able to scale the number of objects in SQAIR due to a parallel inference mechanism. STOVE [29] introduces a graph neural network (GNN) [30] for dynamics prediction and per-object interactions.

One of the most notable world models is the Generative Structured World Model (G-SWM) [31]. It treats foreground objects and the background separately by encoding them into two different latent vectors. For next-state prediction, it uses two separate RNNs—one for the foreground latent vector and the other for the background. One of the limitations of this model is that it does not consider actions taken in the environment when predicting future trajectories.

Another important approach is Contrastive Learning of Structured World Models (C-SWM) [32]. It uses a contrastive loss function at the slot level, instead of the basic image reconstruction loss. The model predicts the trajectory using a GNN. A notable improvement upon C-SWM is negative sampling [33], which improves the loss function by selecting negative samples from either different time steps in the same episode or the same time step in different episodes. Another improvement is represented by two types of action attention [34]. Soft attention uses simple self-attention with the single head of a transformer [35]. Hard attention calculates the expectation of all possible assignments to objects, takes the index of the object with the highest probability, and maps the action to this object. The quality of predictions in models utilizing contrastive loss is significantly influenced by the quality of negative samples [33]. The necessity of data deduplication imposes additional constraints on the dataset sampling process, complicating it further in complex, real-world environments. However, since our model employs only the mean squared error (MSE) loss for training, it does not impose strict requirements on the datasets.

Slotformer [4] uses either STEVE or SAVI to extract slots from a sequence of images. The slots are then passed through a transformer to predict future frames of the video. This approach has achieved better results in dynamics prediction than previous models. A further improvement is OCVP [18], which decouples slot interaction and causal dynamics into two attention blocks. However, the main limitation of this method is that, unlike our approach, the model cannot work efficiently in action-dependent environments. Since the taken action can significantly change the environment or agent state, predicting the future based only on previous states may prove to be challenging.

3. Materials and Methods

The COMPAS world model consists of two interconnected components (as shown in Figure 2): the slot extractor and the transition model. The purpose of the slot extractor is to generate object-centric slot representations that remain consistent both within a trajectory and across episodes. Consistency ensures that each slot corresponds to a specific object in the scene, maintaining its association without permuting with other slots over time or across episodes. The transition model predicts the next time step based on a context window of length T containing previous slot representations and the actions taken. The formal problem statement of the COMPAS learning task is presented in Appendix D.

Figure 2. The method starts by inputting a trajectory of T observations into a frozen DINO encoder to extract features. Trainable tokens are duplicated and combined with positional encodings to produce deterministic slot initializations $S_t^{\mathrm{init}}$ (Figure 3a). These, along with the features, are fed into a slot extractor with temporal sinusoidal encoding. The extractor has two components: a shared key slot attention block that maps slots to feature patches and a temporal consistency block to maintain coherence across time. After $L_{\mathrm{extr}}$ layers, slot representations $S_t$ are produced for each step. A transition model then predicts future states using discrete $a_t^{\mathrm{desc}}$ or continuous $a_t^{\mathrm{cont}}$ actions, processed into vectors. Slots and actions, embedded with temporal encodings, are passed through the predictor block (Figure 3c). After $L_{\mathrm{dyn}}$ layers, the model outputs next-step slots $S_{T+1}$.

Figure 2. The method starts by inputting a trajectory of T observations into a frozen DINO encoder to extract features. Trainable tokens are duplicated and combined with positional encodings to produce deterministic slot initializations $S_t^{\mathrm{init}}$ (Figure 3a). These, along with the features, are fed into a slot extractor with temporal sinusoidal encoding. The extractor has two components: a shared key slot attention block that maps slots to feature patches and a temporal consistency block to maintain coherence across time. After $L_{\mathrm{extr}}$ layers, slot representations $S_t$ are produced for each step. A transition model then predicts future states using discrete $a_t^{\mathrm{desc}}$ or continuous $a_t^{\mathrm{cont}}$ actions, processed into vectors. Slots and actions, embedded with temporal encodings, are passed through the predictor block (Figure 3c). After $L_{\mathrm{dyn}}$ layers, the model outputs next-step slots $S_{T+1}$.

3.1. Slot Representation Extraction

Contrary to iterative approaches used in the classical slot attention algorithm (or its variations), we utilize sequentially stacked blocks. This design enables increased parallelization, scalability, and faster inference and training. The overall architecture of the slot extractor module is presented in Figure 2.

Firstly, we initialize slots by duplicating a trainable token of size D for $T\times N$ times, resulting in a slot initialization ${\widehat{S}}_t^{init}$. T represents the dimension of the starting time window, D is the size of a single slot, and N is the number of slots. By adding a trainable slot position embedding $p_s$, which encodes each token with specified slot position and sinusoidal time position embeddings $p_t$ to the corresponding slots, we obtain the slot initialization at time step t:

$$S_t^{init}={\widehat{S}}_t^{init}+p_s+p_t$$

(1)

The slot’s positional encoding remains constant across time steps but differs for each slot. Similarly, the time positional encoding is the same for all slots within a single time step but varies between time steps. This deterministic initialization, combined with a trainable token and slot position embeddings, ensures slot consistency across episodes.

After initialization, we feed $S_t^{init}$ into the slot extraction block. To generalize the notation, let $S_t^l$ represent the slots from the l-th layer of the slot extractor block, where $S_t^0=S_t^{init}$. This block extracts slot representations from encoder features $F_t\in {\mathbb{R}}^{M\times J}$, where M is the number of patches, and J denotes the encoder feature dimensions. We utilize a modified version of the shared key slot attention [23]. The encoder features $F_t$ are extracted from input image frames using a pre-trained DINO [25] encoder:

$$\begin{array}{c}\hfill \begin{array}{cc}\hfill & Q=\mathrm{Linear}\left(S_t^{l-1}\right)\hfill \\ \hfill & K=\mathrm{Linear}\left(F_t\right)\hfill \\ \hfill & W=\mathrm{Softmax}\left({\displaystyle \frac{K{Q}^T}{\tau }},\mathrm{axis}=\mathrm{slots}\right)\hfill \\ \hfill & R=W^{T}K\hfill \\ \hfill & {\tilde{S}}_t^l=\mathrm{MLP}\left(R\right)+S_t^{l-1}\hfill \end{array}\end{array}$$

(2)

Object-centric models for video, such as the SAVI or Videosaur models, implement slot trajectory consistency through a recurrent predictor model, which predicts slot initialization for the next time step based on the slots from the previous time step. However, this approach may be less high-performing and is prone to gradient decay. To address these issues and improve the training and inference speed, we introduce a temporal consistency block. This block processes slots in parallel and is implemented as a transformer encoder layer with a mask $M_{\mathrm{temp}}$. The mask promotes attention to the same slots from previous time steps causally, while ignoring other slots. In other words, the i-th slot ${\tilde{s}}_{t,i}^l\in {\tilde{S}}_t^l$ can only attend to the same slot ${\tilde{s}}_m^{i,l}$ from a previous time step m ($m<t$). The visualization of the matrix is shown in Figure 3b.

By feeding the sequence of extracted slots from the l-th layer in a time window of size T, ${\left\{{\tilde{S}}_t^l\right\}}_{t=1}^T$, into the temporal consistency block, we obtain the final slots:

$${\left\{S_t^l\right\}}_{t=1}^T=\mathrm{TransfEncLayer}({\left\{{\tilde{S}}_t^l\right\}}_{t=1}^T,\mathrm{mask}=M_{\mathrm{temp}})$$

(3)

Assuming that the slot extractor model comprises L layers, sequentially implementing these operations results in the final slots ${\left\{S_t\right\}}_{t=1}^T={\left\{S_t^L\right\}}_{t=1}^T$.

Since the context length of the model is constrained to a fixed time window T, which is typically small, challenges arise regarding the scalability of the model for longer videos. To address this issue, the slot extractor module is trained in an autoregressive manner. The initialization of the next slots, $S_{T+1}^{init}$, is performed as described in Equation (1). Subsequently, the slots are extracted using the procedure detailed in Equation (2).

The extracted slots, ${\tilde{S}}_{T+1}^l$, are then appended to the shifted time window of previously extracted slots, ${\left\{S_t\right\}}_{t=2}^T$, resulting in the sequence $\left\{S_2,S_3,\dots ,{\tilde{S}}_{T+1}^l\right\}$. These slots are then fed into the temporal consistency block described in Equation (3). By sequentially processing them through all L layers, the final slots, $S_{T+1}$, are obtained. This procedure can be repeated for any desired horizon beyond the initial context length T.

The training approach follows the methodology outlined in DINOSAUr [6]. After extracting the slots $S_t$ across the entire training trajectory $T+Z$, where Z represents the prediction horizon for autoregressive extraction, these slots are processed through a shared MLP decoder to generate feature reconstructions ${\widehat{F}}_t$. The training objective is to minimize the mean squared error (MSE) loss: ${\mathcal{L}}_{\mathrm{extr}}={\displaystyle \frac{1}{T+Z}}{\sum }_{t=1}^{T+Z}{\parallel F_t-{\widehat{F}}_t\parallel }^2$.

The overall time complexity of a single forward pass of the extractor model is $O(T·L·K·N·D)$, where T is the initial time window of slots, L is the number of extractor layers, K is the number of frame feature patches, N is the number of slots, and D is the slot dimension.

3.2. Transition Model

The transition model is inspired by OCVP. During the early phases of training of the slot extractor, slots often lack useful information, which can push the transition model away from the optimal solution or cause it to become stuck in a local minimum. To address these challenges and accelerate the training process, the transition model is trained separately from the slot extractor on pre-computed slots.

At the start, slots $S_t\in {\mathbb{R}}^{N\times D},t=0,\dots ,T$ within a time window of size T are projected into the transition model’s dimension. Since the transition model operates with the same dimensions D as the slots,

$$S_{\mathrm{prep},t}=\mathrm{Linear}\left(S_t\right)+p_t$$

(4)

The transition model supports both continuous and discrete actions. Continuous actions $a_t^{\mathrm{cont}}\in {\mathbb{R}}^B$ are projected into the same dimension as the transition model:

$${\widehat{a}}_t\in {\mathbb{R}}^D=\mathrm{Linear}\left(a_t^{\mathrm{cont}}\right)$$

(5)

For discrete actions, we use trainable embeddings for each possible action $a_t^{\mathrm{desc},i}\in A_t^{\mathrm{desc}}=\{a_t^{\mathrm{desc},1},a_t^{\mathrm{desc},2},\dots ,a_t^{\mathrm{desc},J}\}$:

$${\widehat{a}}_t=\mathrm{EmbeddingsDictionary}\left(a_t^{\mathrm{desc},i}\right)$$

(6)

To incorporate temporal information, we add sinusoidal time embeddings $p_t$ to each slot and action at time step t:

$$\begin{array}{c}\hfill \begin{array}{cc}\hfill & s_{\mathrm{prep},t}^i=s_t^i+p_t,s_{\mathrm{prep},t}^i\in S_{\mathrm{prep},t}\hfill \\ \hfill & a_t={\widehat{a}}_t+p_t\hfill \end{array}\end{array}$$

(7)

Subsequently, actions and slots are passed into an action attention block, relation attention block, and temporal attention in parallel to the original OCVP layer. The architecture of the predictor block is presented in Figure 3c.

Action attention is implemented as a simple multi-head cross-attention. In this layer, actions $a_t$ serve as keys and values, while slots $S_t$ act as queries.

After projecting the features from the final time step into slot dimensions, we obtain the predicted features for the next time step, denoted as $S_{T+1}^{\prime }$. By shifting the input slot time window and appending $S_{T+1}^{\prime }$ as the last element, along with obtaining ${\widehat{a}}_{T+1}$, we can autoregressively predict future states for any horizon H.

The training objective for the COMPAS dynamics model is defined as the ${\mathcal{L}}_2$ reconstruction loss between the ground truth slots $s_{t,i}\in S_t$ and the model’s predictions $s_{t,i}^{\prime }\in S_t^{\prime }$:

$${\mathcal{L}}_{dyn}={\displaystyle \frac{1}{H·N}}\sum _{t=T+1}^H\sum _{i=1}^N\left|\right|s_{t,i}^{\prime }-s_{t,i}{\left|\right|}^2$$

(8)

The single-forward-pass time complexity of the transition model is $O(L·D·T·N·(N+2·T\left)\right).$

4. Results

4.1. Environments and Datasets

To evaluate the effectiveness of our proposed model, we perform experiments using trajectories sampled from three environments: ViZDoom [36], Causal World [37], and RoboSuite [38]. By conducting experiments in the selected environments, we aim to demonstrate the versatility and scalability of our model. A detailed description of the datasets is presented in Appendix B.

We primarily focus on comparing the predictive accuracy of our model against that of recent state-of-the-art (SOTA) methods, such as Slotformer and OCVP. Additionally, we include comparisons with the GNN transition model from CSWM, despite it being outdated and relying on a lower-quality slot extraction procedure.

Additionally, we evaluate the COMPAS extractor model in terms of the quality of object detection and slot consistency performance against the VideoSAUR [17] model across 20 and 100 steps.

4.2. Metrics

Other frameworks, such as CSWM [34], use ranking metrics like HITS@1 and MRR for dynamics prediction evaluation. However, this approach is not entirely representative, as these metrics depend on the scaling of the actual slot values and can yield high scores even if the predictions are inaccurate. To address this, we evaluate the models using pre-computed slots $S_t$ and decode them back into observations $O_t$ with sizes $H\times W$. If a model can accurately predict the next slots ${\widehat{S}}_t$, the mean squared error (MSE) with the ground truth should be lower compared to competing models. Similarly, the decoder MSE should also be lower, as better dynamics predictions lead to improved reconstructions $\widehat{O_t}$. We calculate the metric at a certain time step t in order to evaluate the quality of predictions at certain horizons.

$$\begin{array}{c}\hfill \begin{array}{cc}\hfill & {\mathrm{MSE}}_t={\displaystyle \frac{1}{N}}\left|\right|S_t-\widehat{S_t}{\left|\right|}^2\hfill \\ \hfill & {\mathrm{MSE}}_t^{\mathrm{decoder}}={\displaystyle \frac{1}{3(H\times W)}}\left|\right|O_t-\widehat{O_t}{\left|\right|}^2\hfill \end{array}\end{array}$$

(9)

To measure consistency, we use the Identity F1 (IDF1) score:

$$\mathrm{IDF}1={\displaystyle \frac{2\mathrm{TP}}{2\mathrm{TP}+\mathrm{FP}+\mathrm{FN}}},$$

(10)

where $\mathrm{TP}$ is the number of correctly matched object identities, $\mathrm{FP}$ is the number of predicted identities that do not correspond to any ground truth identity, and $\mathrm{FN}$ is the number of ground truth identities not recovered by the tracker.

To count identity switches, we first build a pairwise IoU between each predicted slot mask and each ground truth mask. We then apply the Hungarian algorithm to find the optimal matching. Thus, we detect slot permutations and treat them as an identity switch.

To measure trajectory consistency, we treat the predicted mask from time $t-1$ as the ground truth for time t.

For episodic consistency, we fix masks as the ground truth from one reference episode. For every other episode of the same object, we compute the IDF1 between its mask and the ground truth mask from the reference episode.

For the ARI, FG-ARI, and mbo metrics, we follow the setup from [17].

4.3. Training and Evaluation Setup

All transition models are trained and evaluated on pre-computed slots extracted by the COMPAS slot extractor. The slot extractor was trained on 150 epochs with a batch size of 128 for each environment.

More detailed descriptions are available in Appendix A and Appendix B. All configurations and code required for the experiments will be made available.

4.4. Experiments

We introduce OCVP models with simple action conditioning. OCVP (slot act.) uses the processed action vector ${\widehat{a}}_t$ as an additional slot. OCVP (add act.) adds ${\widehat{a}}_t$ to each slot at time step t.

As highlighted in

Table 1. Comparative results for the Causal World (Reach) environment. Mean ± std for 3 seeds. All values are multiplied by 100 for better visual presentation.

	10 Steps		30 Steps		60 Steps
Model	MSE	MSEdecoder	$\mathbf{MSE}$	${\mathbf{MSE}}^{\mathbf{decoder}}$	$\mathbf{MSE}$	${\mathbf{MSE}}^{\mathbf{decoder}}$
Slotformer	$6.48\pm 0.02$	$15.35\pm 0.03$	$7.42\pm 0.12$	$17.95\pm 0.10$	$8.22\pm 0.19$	$19.30\pm 0.19$
OCVP	$6.32\pm 0.03$	$15.10\pm 0.04$	$7.28\pm 0.11$	$17.20\pm 0.13$	$8.08\pm 0.18$	$19.25\pm 0.17$
OCVP (add act.)	$4.66\pm 0.03$	$11.72\pm 0.03$	$5.88\pm 0.09$	$13.58\pm 0.10$	$6.94\pm 0.17$	$15.66\pm 0.17$
OCVP (slot act.)	$4.26\pm 0.02$	$10.86\pm 0.02$	$5.46\pm 0.08$	$12.84\pm 0.09$	$6.62\pm 0.16$	$14.92\pm 0.16$
GNN	$5.18\pm 0.03$	$14.48\pm 0.03$	$6.75\pm 0.10$	$16.18\pm 0.12$	$7.88\pm 0.18$	$18.42\pm 0.18$
COMPAS (ours)	3.18 ± 0.01	4.38 ± 0.01	4.23 ± 0.07	6.82 ± 0.07	5.12 ± 0.15	10.22 ± 0.15

, COMPAS demonstrated superior predictive accuracy in the Causal World environment with the Reach task. This finding underscores the robustness of our approach in handling complex, physically simulated environments.

Furthermore, COMPAS’s performance increased significantly in the more challenging Push task (

Table 2. Comparative results for the Causal World (Push) environment. Mean ± std for 3 seeds. All values are multiplied by 100 for better visual presentation.

	10 Steps		30 Steps		60 Steps
Model	$\mathbf{MSE}$	${\mathbf{MSE}}^{\mathbf{decoder}}$	$\mathbf{MSE}$	${\mathbf{MSE}}^{\mathbf{decoder}}$	$\mathbf{MSE}$	${\mathbf{MSE}}^{\mathbf{decoder}}$
Slotformer	$6.36\pm 0.02$	$15.33\pm 0.04$	$7.38\pm 0.10$	$17.94\pm 0.12$	$8.07\pm 0.18$	$19.08\pm 0.20$
OCVP	$8.65\pm 0.03$	$16.11\pm 0.05$	$9.40\pm 0.12$	$18.98\pm 0.15$	$10.04\pm 0.19$	$21.00\pm 0.18$
OCVP (add act.)	$4.72\pm 0.03$	$12.84\pm 0.03$	$5.98\pm 0.10$	$14.76\pm 0.11$	$7.02\pm 0.17$	$16.92\pm 0.18$
OCVP (slot act.)	$4.38\pm 0.02$	$11.96\pm 0.03$	$5.62\pm 0.09$	$13.88\pm 0.10$	$6.76\pm 0.16$	$15.84\pm 0.17$
GNN	$5.20\pm 0.04$	$14.50\pm 0.03$	$6.80\pm 0.11$	$16.20\pm 0.14$	$7.90\pm 0.17$	$18.40\pm 0.19$
COMPAS (ours)	3.15 ± 0.01	6.40 ± 0.01	4.25 ± 0.08	8.80 ± 0.08	5.10 ± 0.16	10.20 ± 0.16

). In contrast, non-action-dependent models, such as Slotformer and OCVP, exhibited substantial performance degradation under these conditions. This comparison highlights the resilience and adaptability of COMPAS in diverse and complex tasks.

On the RoboSuite dataset, the GNN baseline achieved better results than the other baseline models, likely due to the lower object density and reduced interaction complexity in this environment. Nevertheless, COMPAS still achieved the best overall performance on this dataset, outperforming all baseline methods (

Table 3. Comparative results for the RoboSuite environment. Mean ± std for 3 seeds. All values are multiplied by 100 for better visual presentation.

	10 Steps		30 Steps		60 Steps
Model	$\mathbf{MSE}$	${\mathbf{MSE}}^{\mathbf{decoder}}$	$\mathbf{MSE}$	${\mathbf{MSE}}^{\mathbf{decoder}}$	$\mathbf{MSE}$	${\mathbf{MSE}}^{\mathbf{decoder}}$
Slotformer	$1.70\pm 0.03$	$9.83\pm 0.08$	$1.99\pm 0.11$	$11.52\pm 0.14$	$2.65\pm 0.18$	$12.01\pm 0.20$
OCVP	$1.82\pm 0.05$	$10.55\pm 0.07$	$2.20\pm 0.13$	$12.19\pm 0.15$	$2.58\pm 0.19$	$12.94\pm 0.18$
GNN	$1.00\pm 0.02$	$9.46\pm 0.04$	$1.49\pm 0.12$	$12.08\pm 0.13$	$1.89\pm 0.17$	$12.97\pm 0.19$
COMPAS (ours)	0.50 ± 0.01	8.08 ± 0.03	0.86 ± 0.11	10.53 ± 0.12	1.25 ± 0.16	11.93 ± 0.17

).

On the ViZDoom dataset, COMPAS performed consistently well, achieving better results compared to the alternative methods. The results are presented in

Table 4. Comparative results for the ViZDoom (Defend the Line) environment. Mean ± std for 3 seeds. All values are multiplied by 100 for better visual presentation.

	10 Steps		30 Steps		60 Steps
Model	$\mathbf{MSE}$	${\mathbf{MSE}}^{\mathbf{decoder}}$	$\mathbf{MSE}$	${\mathbf{MSE}}^{\mathbf{decoder}}$	$\mathbf{MSE}$	${\mathbf{MSE}}^{\mathbf{decoder}}$
Slotformer	$14.58\pm 0.05$	$6.51\pm 0.02$	$20.22\pm 0.12$	$14.94\pm 0.10$	$27.39\pm 0.18$	$30.31\pm 0.20$
OCVP	$14.10\pm 0.03$	$6.08\pm 0.03$	$20.41\pm 0.10$	$14.89\pm 0.12$	$26.86\pm 0.20$	$31.04\pm 0.18$
GNN	$8.75\pm 0.04$	$5.20\pm 0.03$	$15.60\pm 0.11$	$12.40\pm 0.11$	$23.10\pm 0.17$	$32.50\pm 0.19$
COMPAS (ours)	4.31 ± 0.01	2.85 ± 0.01	9.43 ± 0.08	10.27 ± 0.08	21.23 ± 0.16	29.17 ± 0.16

.

Moreover, COMPAS demonstrated better object detection quality and slot assignment consistency than the Videosaur model for both short (20 steps) and long (100 steps) trajectories in the Causal World (Push) environment. These results are presented in

Table 5. Performance comparison of COMPAS extractor and Videosaur model over 20 and 100 steps in Causal World (Push) environment. Mean ± std for 3 seeds.

Model	ARI	FG-ARI	mbo	Trajectory IDF1	Episodic IDF1
COMPAS (100 steps)	$\mathbf{0}.\mathbf{43}\pm \mathbf{0}.\mathbf{14}$	$\mathbf{0}.\mathbf{61}\pm \mathbf{0}.\mathbf{05}$	$\mathbf{0}.\mathbf{79}\pm \mathbf{0}.\mathbf{03}$	$\mathbf{1}.\mathbf{00}\pm \mathbf{0}.\mathbf{00}$	$\mathbf{0}.\mathbf{78}\pm \mathbf{0}.\mathbf{05}$
Videosaur (100 steps)	$0.19\pm 0.03$	$0.20\pm 0.02$	$0.19\pm 0.02$	$0.21\pm 0.04$	$0.11\pm 0.02$
COMPAS (20 steps)	$\mathbf{0}.\mathbf{51}\pm \mathbf{0}.\mathbf{06}$	$\mathbf{0}.\mathbf{69}\pm \mathbf{0}.\mathbf{03}$	$\mathbf{0}.\mathbf{82}\pm \mathbf{0}.\mathbf{01}$	$\mathbf{1}.\mathbf{00}\pm \mathbf{0}.\mathbf{00}$	$\mathbf{0}.\mathbf{82}\pm \mathbf{0}.\mathbf{04}$
Videosaur (20 steps)	$0.44\pm 0.05$	$0.51\pm 0.06$	$0.76\pm 0.02$	$0.74\pm 0.04$	$0.12\pm 0.01$

.

The ablation study of the COMPAS extractor, presented in

Table 6. Ablation analysis of COMPAS extractor model. Mean ± std for 3 seeds.

Model	ARI	FG-ARI	mbo	Trajectory IDF1	Episodic IDF1
COMPAS	$\underline{0.43\pm 0.14}$	$\underline{0.61\pm 0.05}$	$\underline{0.79\pm 0.03}$	$\mathbf{1}.\mathbf{00}\pm \mathbf{0}.\mathbf{00}$	$\underline{0.78\pm 0.05}$
No temporal block	$\underline{0.42\pm 0.12}$	$\underline{0.60\pm 0.04}$	$\underline{0.75\pm 0.02}$	$0.89\pm 0.04$	$\underline{0.75\pm 0.04}$
No temporal encoding	$0.18\pm 0.05$	$0.21\pm 0.01$	$0.21\pm 0.03$	$0.41\pm 0.05$	$0.34\pm 0.05$
No temporal block and encoding	$0.16\pm 0.04$	$0.18\pm 0.04$	$0.19\pm 0.02$	$0.32\pm 0.05$	$0.28\pm 0.03$
No slot encoding	$0.00\pm 0.00$	$0.00\pm 0.00$	$0.18\pm 0.01$	$0.21\pm 0.05$	$0.19\pm 0.05$

, demonstrates the effectiveness of the proposed architectural components. While the temporal block has a smaller impact on the segmentation quality compared to temporal encoding, it significantly enhances the model’s episodic temporal consistency.

Similarly, the ablation study of the dynamics model, shown in

Table 7. Ablation analysis of COMPAS dynamics model. Mean ± std for 3 seeds. All values are multiplied by 100 for better visual presentation.

	10 Steps		30 Steps		60 Steps
Model	$\mathbf{MSE}$	${\mathbf{MSE}}^{\mathbf{decoder}}$	$\mathbf{MSE}$	${\mathbf{MSE}}^{\mathbf{decoder}}$	$\mathbf{MSE}$	${\mathbf{MSE}}^{\mathbf{decoder}}$
COMPAS	$\mathbf{3}.\mathbf{15}\pm \mathbf{0}.\mathbf{01}$	$\mathbf{6}.\mathbf{40}\pm \mathbf{0}.\mathbf{01}$	$\mathbf{4}.\mathbf{25}\pm \mathbf{0}.\mathbf{08}$	$\mathbf{8}.\mathbf{80}\pm \mathbf{0}.\mathbf{08}$	$\mathbf{5}.\mathbf{10}\pm \mathbf{0}.\mathbf{16}$	$\mathbf{10}.\mathbf{20}\pm \mathbf{0}.\mathbf{16}$
No interaction block	$4.83\pm 0.02$	$12.17\pm 0.03$	$5.41\pm 0.07$	$13.71\pm 0.09$	$5.96\pm 0.12$	$15.45\pm 0.14$
No temporal block	$4.84\pm 0.03$	$12.02\pm 0.02$	$5.35\pm 0.06$	$13.50\pm 0.10$	$5.82\pm 0.11$	$15.25\pm 0.13$
No action block	$8.65\pm 0.03$	$16.11\pm 0.05$	$9.40\pm 0.12$	$18.98\pm 0.15$	$10.04\pm 0.19$	$21.00\pm 0.18$

, highlights the importance of the action-processing component in improving the prediction quality.

An ablation study of separated and end-to-end training for the COMPAS dynamics model is conducted in the Causal World (Push) environment. In the separated setup, the extractor and dynamics model are trained separately, whereas, in the end-to-end setup, they are trained jointly using different optimizers.

Table 8. Ablation analysis of separated and end-to-end training for the COMPAS dynamics model in the Causal World (Push) environment. Mean ± std for 3 seeds. All values are multiplied by 100 for better visual presentation.

	10 Steps		30 Steps		60 Steps
Model	$\mathbf{MSE}$	${\mathbf{MSE}}^{\mathbf{decoder}}$	$\mathbf{MSE}$	${\mathbf{MSE}}^{\mathbf{decoder}}$	$\mathbf{MSE}$	${\mathbf{MSE}}^{\mathbf{decoder}}$
COMPAS (separated)	3.15 ± 0.01	6.40 ± 0.01	4.25 ± 0.08	8.80 ± 0.08	5.10 ± 0.16	10.20 ± 0.16
COMPAS (end-to-end)	$3.48\pm 0.02$	$6.92\pm 0.02$	$4.76\pm 0.09$	$9.54\pm 0.09$	$5.84\pm 0.17$	$11.18\pm 0.17$

shows that end-to-end training is slightly poorer. A possible reason is that the extractor learns less effective representations under joint optimization, which negatively affects the dynamics model.

We also include an ablation study on the effect of the temporal context size. As shown in

Table 9. Ablation analysis of context size for the COMPAS extractor model in the Causal World (Push) environment. Mean ± std for 3 seeds.

Context Size	ARI	FG-ARI	mbo	Trajectory IDF1	Episodic IDF1
3	$0.40\pm 0.13$	$0.58\pm 0.05$	$0.76\pm 0.03$	$1.00\pm 0.00$	$0.75\pm 0.05$
4	$0.43\pm 0.14$	$0.61\pm 0.05$	$0.79\pm 0.03$	$1.00\pm 0.00$	$0.78\pm 0.05$
10	$0.44\pm 0.13$	$0.60\pm 0.04$	$0.80\pm 0.02$	$1.00\pm 0.00$	$0.79\pm 0.04$

and

Table 10. Ablation analysis of context size for the COMPAS dynamics model in the Causal World (Push) environment. Mean ± std for 3 seeds. All values are multiplied by 100 for better visual presentation.

	10 Steps		30 Steps		60 Steps
Context Size	$\mathbf{MSE}$	${\mathbf{MSE}}^{\mathbf{decoder}}$	$\mathbf{MSE}$	${\mathbf{MSE}}^{\mathbf{decoder}}$	$\mathbf{MSE}$	${\mathbf{MSE}}^{\mathbf{decoder}}$
3	$3.48\pm 0.02$	$6.92\pm 0.02$	$4.76\pm 0.09$	$9.54\pm 0.09$	$5.84\pm 0.17$	$11.18\pm 0.17$
4	$\underline{3.15\pm 0.01}$	$\underline{6.40\pm 0.01}$	$\underline{4.25\pm 0.08}$	$\underline{8.80\pm 0.08}$	$\underline{5.10\pm 0.16}$	$\underline{10.20\pm 0.16}$
10	$\underline{3.16\pm 0.01}$	$\underline{6.42\pm 0.01}$	$\underline{4.24\pm 0.08}$	$\underline{8.82\pm 0.08}$	$\underline{5.11\pm 0.16}$	$\underline{10.19\pm 0.16}$

, reducing the context size leads to a slight drop in performance, indicating that a shorter temporal window provides less information for accurate representation learning and dynamics prediction. In contrast, increasing the context size beyond 4 does not yield noticeable improvements, as the results remain nearly unchanged. This suggests that a moderate context size is already sufficient to capture the relevant temporal information, while a larger context provides little additional benefit.

We also include experiments demonstrating the use of our world model for decision-making tasks—specifically, planning with the DINO-WM [39] baseline—in Appendix D.

Overall, COMPAS proves to be a more robust and adaptable model, particularly in synthetic environments and tasks with both discrete and continuous action spaces, further proving its advantages over object-centric baseline methods.

5. Discussion

The main limitation of our model is the introduction of a new hyperparameter, the context length, which can impact the consistency of the slots. Additionally, the model’s performance is influenced by hyperparameters such as the number of slots and the context time window, both of which can significantly affect the quality of slot representations and dynamics predictions. Furthermore, handling real-world data poses a significant challenge, as object-centric models inherently struggle to process such data efficiently. Future research could address these issues and improve the model’s robustness and adaptability.

6. Conclusions

Our model demonstrates superior consistency over multiple prediction steps compared to analogous models in the field. This confirms the model’s ability to handle different prediction horizons.

Looking ahead, a promising direction is to explore neuro-symbolic perspectives in object-centric approaches. By combining neural networks with symbolic reasoning, models can integrate structured knowledge with learned data, enhancing interpretability, robustness, and generalization. Another key direction is developing algorithms that efficiently leverage object information to drive agent behavior. This requires the careful handling of composable states, as merging multiple object states into a single representation can obscure critical information and reintroduce the binding problem.

In conclusion, our work represents an important step forward in object-centric world models. By integrating action information directly into the model and using an autoregressive transformer model for prediction, we have demonstrated a novel way to encode and predict the dynamics of complex environments. We expect that our work will inspire further research and advances in this exciting area of machine learning.