Better Slots, Better Worlds

Representation quality and robustness in object-centric world models.

Shukrullo Nazirjonov1, Sai Prasanna1,2†, Anna Manasyan1,2†, Georg Martius1,2

1University of Tübingen  2Max Planck Institute for Intelligent Systems

†Equal contribution

Abstract

TL;DR. We study how object-centric representation quality affects visual planning. Using a fixed world model and planner, we vary the slot encoder, training objective, and robot-state input. Planning success tracks standard slot-quality metrics, with the gains saturating at high quality. Given a well-bound slot encoder, the masking objective and proprioceptive input that prior methods relied on do not improve planning. And under visual distribution shifts, world models built on frozen pretrained features stay more robust than an end-to-end baseline, whether those features are object-centric or patch-based.

Learning world models from offline trajectories enables agents to accomplish different tasks through planning. Object-centric (OC) representations, which decompose a scene into a set of slots that bind to its objects, have been proposed as an inductive bias for world models that are more sample-efficient and generalize better. Yet prior object-centric world models (OCWMs) take the slot encoder as given and evaluate only in-distribution, leaving open whether the object-centric bias actually delivers for planning and what within the OCWM drives it. We conduct a controlled study of OCWMs for visual model-predictive control along two axes: object-centric representation quality and generalization under distribution shift relative to scene-centric models. We find that (i) planning success correlates positively with unsupervised slot-quality metrics (FG-ARI, mBO), though the gains saturate at high slot quality; (ii) with well-bound slots, the auxiliary proprioception inputs and masking inductive bias that prior methods relied on become unnecessary; and (iii) an OCWM with well-bound slots plans more robustly under unseen distribution shifts than the end-to-end trained scene-centric LeWM, while DINO-WM, built on similar frozen pretrained features, remains comparably robust — suggesting that pretrained visual representations are an important contributor to robustness.

Setup

We build on a recently proposed object-centric world model, Causal JEPA (C-JEPA from here on) [3]: a frozen encoder turns each frame into a set of object slots [6], and an action-conditioned bidirectional transformer predicts their next state. In addition to predicting future latent slots, C-JEPA introduces a history masking objective that forces the model to recover a masked object’s dynamics by relying on its interactions with other objects in the scene. Given such an object-centric world model, we evaluate it by planning in slot space with model-predictive control and the cross-entropy method [11]. Unless stated otherwise, each world model is trained with a single training seed, and every reported planning success rate is averaged over three planning seeds.

C-JEPA world model
The C-JEPA world model and planner.
MPC + CEM planning loop
At each step, the planner evaluates candidate action sequences in slot space and replans after executing the first action.

We evaluate on 2D PushT and 3D OGBench Single Cube [10]. Baselines include DINO-WM [1], LeWM [2], and C-JEPA [3], all with the same planning budget.

Q1 · Does slot mask quality drive planning?

If objects are the right abstraction, how cleanly a model separates them should matter, and it does, up to a point. We take intermediate checkpoints of the state-of-the-art object-centric encoder SlotContrast [4], which vary in slot-mask quality as measured by standard metrics (FG-ARI, mBO), and train the same world model on each. FG-ARI and mBO measure how well objects are separated into slot masks and how sharp those masks are, respectively. As the slot masks improve, downstream planning performance follows. The trend is clean on PushT and weaker on OGBench-Cube, where an aggregate overlap score can be dominated by one large object while the small, task-relevant one stays poorly bound.

Slot masks over time on PushT
Slot masks over an episode on PushT. SlotContrast separates the agent and the T-block into distinct, temporally consistent slots; VideoSAUR entangles them.
Slot masks over time on OGBench-Cube
Slot masks over an episode on OGBench-Cube. SlotContrast binds the robot arm, gripper, and cube in 3D, holding them consistent even through contact.
Slot quality vs planning success on PushT
Planning success versus slot quality on PushT. The trend is clean along both FG-ARI and mBO.
Slot quality vs planning success on OGBench-Cube
Planning success versus slot quality on OGBench-Cube. The same ordering holds but the trend is weaker.

Takeaway 1. Planning success is positively correlated with slot quality, but the gains saturate: once slots are well-bound, better-bound slots add little.

Q2 · Do auxiliary objectives and privileged inputs help?

C-JEPA adds a slot-history masking objective and feeds proprioception to the world model as a separate token. If the agent is already observable from the visual (slot) input alone, as in PushT, a natural question is whether the world model needs proprioception at all. To test this, we sweep the number of masked slots, with and without the proprioception token, for a strong encoder (SlotContrast [4]) and a weaker one (VideoSAUR [5]). With the proprioception token present, which directly encodes task-critical information such as agent velocity and position, C-JEPA’s masking objective helps the weaker VideoSAUR-based encoder. For the stronger SlotContrast encoder, masking is not beneficial even with proprioception (blue lines). When the proprioception token is removed, the benefit of masking disappears entirely: for both VideoSAUR- and SlotContrast-based world models, planning success falls as the number of masked slots increases.

Masking and proprioception sweep
Planning success with zero, one, or two masked slots, with and without proprioception.

Takeaway 2. A sufficiently object-centric representation makes slot-history masking and proprioception unnecessary; they mainly compensate for weak representations.

Q3 · What makes planning robust to distribution shift?

To probe the robustness of learned world models, we test zero-shot planning while changing colors, backgrounds, object scale, and object shape at test time. SlotContrast-WM (C-JEPA with no masking objective) and DINO-WM keep more of their performance than the end-to-end LeWM baseline. The two share a frozen pretrained encoder but differ in whether their features are object-centric, which points to pretraining, not object-centricity by itself, as the source of robustness.

OOD robustness on PushT
Planning success under appearance, scale, and shape shifts on PushT.
OOD robustness on Cube
Planning success under appearance, scale, and shape shifts on OGBench-Cube.

Takeaway 3. World models planning over frozen pretrained features — object-centric or not — tolerate appearance and frame-level shifts far better than the end-to-end trained LeWM.

Additional ablations

How should the predictor be trained for rollout?

Planning rolls the world model out for several steps, so the way the predictor is trained for multi-step prediction matters. C-JEPA pairs a bidirectional transformer with single-step teacher forcing, which works well in that setting. We ask a complementary question: which predictor holds up when trained autoregressively, on its own multi-step rollouts? Alongside C-JEPA we look at SlotFormer [7], an earlier object-centric predictor built on a causal transformer. Under teacher forcing the two are comparable. The difference appears under autoregressive training, where the model is fed its own predictions: the causal predictor takes to it naturally and stays more accurate over the rollout (85 → 90 on PushT), while the bidirectional transformer, built for single-step teacher forcing, does not transfer as well to this regime (84 → 66). The same trend appears on OGBench-Cube. For multi-step autoregressive training, then, a causal transformer is a natural fit.

Teacher forcing vs autoregressive training on PushT
Teacher forcing vs autoregressive training on OGBench-Cube
Planning success with one-step teacher forcing (rollout = 1) versus five-step autoregressive training (rollout = 5), on PushT (left) and OGBench-Cube (right). Unlike the other experiments, results here are averaged over 4 training seeds.

DINOv2 versus DINOv3

We also check whether a more recent version of DINO, namely DINOv3, yields a more robust world model. To this end, we replace DINO-WM’s original frozen DINOv2 [8] encoder with two DINOv3 [9] variants. On OGBench-Cube, DINOv3 improves neither in-distribution nor OOD planning.

DINOv2 vs DINOv3 backbone for DINO-WM
DINO-WM planning success with DINOv2 and two DINOv3 encoders on OGBench-Cube.

Imagined rollouts

For interpretability, we train a slot-to-pixel decoder, use it to decode the world model’s predicted slots back to pixels, and place them beside the ground truth.

PushT decoded rollout
Cube decoded rollout
Predicted vs. ground-truth rollouts, decoded for inspection.

Limitations

We test these findings on only two control tasks, which we note as a limitation. Future work should verify these claims more rigorously in more complex, multi-object environments.

References

  1. G. Zhou, H. Pan, Y. LeCun, L. Pinto. DINO-WM: World Models on Pre-trained Visual Features Enable Zero-shot Planning. ICML 2025.
  2. L. Maes, Q. Le Lidec, D. Scieur, Y. LeCun, R. Balestriero. LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels. arXiv:2603.19312, 2026.
  3. H. Nam, Q. Le Lidec, L. Maes, Y. LeCun, R. Balestriero. Causal-JEPA: Learning World Models through Object-Level Latent Interventions. arXiv:2602.11389, 2026.
  4. A. Manasyan, M. Seitzer, F. Radovic, G. Martius, A. Zadaianchuk. Temporally Consistent Object-Centric Learning by Contrasting Slots. CVPR 2025.
  5. A. Zadaianchuk, M. Seitzer, G. Martius. Object-Centric Learning for Real-World Videos by Predicting Temporal Feature Similarities. NeurIPS 2023.
  6. F. Locatello, D. Weissenborn, T. Unterthiner, A. Mahendran, G. Heigold, J. Uszkoreit, A. Dosovitskiy, T. Kipf. Object-Centric Learning with Slot Attention. NeurIPS 2020.
  7. Z. Wu, N. Dvornik, K. Greff, T. Kipf, A. Garg. SlotFormer: Unsupervised Visual Dynamics Simulation with Object-Centric Models. ICLR 2023.
  8. M. Oquab, T. Darcet, T. Moutakanni, et al. DINOv2: Learning Robust Visual Features without Supervision. TMLR 2024.
  9. O. Siméoni, H. V. Vo, M. Seitzer, et al. DINOv3. arXiv:2508.10104, 2025.
  10. S. Park, K. Frans, B. Eysenbach, S. Levine. OGBench: Benchmarking Offline Goal-Conditioned RL. ICLR 2025.
  11. R. Y. Rubinstein. The Cross-Entropy Method for Combinatorial and Continuous Optimization. Methodology and Computing in Applied Probability, 1999.

BibTeX

If you find this work useful, please use the following BibTeX entry to cite:

@inproceedings{nazirjonov2026betterslots,
  author    = {Shukrullo Nazirjonov and Sai Prasanna and
               Anna Manasyan and Georg Martius},
  title     = {Better Slots, Better Worlds: Representation Quality and
               Robustness in Object-Centric World Models},
  year      = {2026},
  eprint       = {2608.12078},
  archivePrefix = {arXiv},
  primaryClass = {cs.CV}
}