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 improves with standard slot-quality metrics, but the benefit saturates once objects are represented well. At that point, masked-slot prediction and proprioceptive input do not help. Additionally, we find that under visual distribution shifts, world models built on frozen pretrained features are 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.

The testbed

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 testbed
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.

Result 1. Planning improves with slot quality, then plateaus.

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.

Result 2. With well-bound slots, masking and proprioception do not improve planning.

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.

Result 3. Frozen pretrained features improve robustness to visual shifts; object-centric structure alone does not explain the difference.

Additional ablations

How do predictor choice and training objective affect performance?

C-JEPA [3] predicts masked slots with a bidirectional transformer trained by teacher forcing. We ask whether that design is necessary, or whether a plain causal next-slot predictor, the older SlotFormer [7] architecture, does as well. The two respond to training in opposite ways. The bidirectional C-JEPA transformer works only with single-step teacher forcing; trained on its own five-step rollouts (autoregressive), it degrades sharply on PushT (84 → 66). The causal predictor does the reverse: autoregressive training helps it (85 → 90). A similar trend appears in the OGBench-Cube environment. A bidirectional predictor may not be needed, and to benefit from multi-step autoregressive training, a causal transformer is the better choice.

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}
}