# Object-Centric World Models from Few-Shot Annotations for Sample-Efficient Reinforcement Learning

arXiv: https://arxiv.org/abs/2501.16443 · HTML: https://mlsanity.com/paper/2501.16443/

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Placed in Latent predictive world models, one of 11 directions in World models.
- Builds on 47 of its 68 references that are in the map.
- First posted on arXiv 27 Jan 2025.

## Placed in

- World models / [Latent predictive world models](https://mlsanity.com/wm/latent-predictive-world-models/)
- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 30 of 47)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Segment Anything](https://mlsanity.com/paper/2304.02643/) (2023)
- [You Only Look Once: Unified, Real-Time Object Detection](https://mlsanity.com/paper/1506.02640/) (2015)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Grounding DINO: Marrying DINO with Grounded Pre-Training for Open-Set Object Detection](https://mlsanity.com/paper/2303.05499/) (2023)
- [SAM 2: Segment Anything in Images and Videos](https://mlsanity.com/paper/2408.00714/) (2024)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [Sigmoid-Weighted Linear Units for Neural Network Function Approximation in Reinforcement Learning](https://mlsanity.com/paper/1702.03118/) (2017)
- [Object-Centric Learning with Slot Attention](https://mlsanity.com/paper/2006.15055/) (2020)
- [Rainbow: Combining Improvements in Deep Reinforcement Learning](https://mlsanity.com/paper/1710.02298/) (2017)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [Ray: A Distributed Framework for Emerging AI Applications](https://mlsanity.com/paper/1712.05889/) (2017)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [XMem: Long-Term Video Object Segmentation with an Atkinson-Shiffrin Memory Model](https://mlsanity.com/paper/2207.07115/) (2022)
- [Mastering Visual Continuous Control: Improved Data-Augmented Reinforcement Learning](https://mlsanity.com/paper/2107.09645/) (2021)
- [Transformers are Sample-Efficient World Models](https://mlsanity.com/paper/2209.00588/) (2022)
- [Data-Efficient Reinforcement Learning with Self-Predictive Representations](https://mlsanity.com/paper/2007.05929/) (2020)
- [Diffusion for World Modeling: Visual Details Matter in Atari](https://mlsanity.com/paper/2405.12399/) (2024)
- [Unifying Flow, Stereo and Depth Estimation](https://mlsanity.com/paper/2211.05783/) (2022)
- [Personalize Segment Anything Model with One Shot](https://mlsanity.com/paper/2305.03048/) (2023)
- [Conditional Object-Centric Learning from Video](https://mlsanity.com/paper/2111.12594/) (2021)
- [Mastering Atari Games with Limited Data](https://mlsanity.com/paper/2111.00210/) (2021)
- [Track Anything: Segment Anything Meets Videos](https://mlsanity.com/paper/2304.11968/) (2023)
- [Tracking Everything Everywhere All at Once](https://mlsanity.com/paper/2306.05422/) (2023)

## Built on by (most cited first, 10 of 10)

- [Hierarchical Entity-centric Reinforcement Learning with Factored Subgoal Diffusion](https://mlsanity.com/paper/2602.02722/) (2026)
- [AXIOM: Learning to Play Games in Minutes with Expanding Object-Centric Models](https://mlsanity.com/paper/2505.24784/) (2025)
- [Capturing Visual Environment Structure Correlates with Control Performance](https://mlsanity.com/paper/2602.04880/) (2026)
- [Latent State Design for World Models under Sufficiency Constraints](https://mlsanity.com/paper/2605.01694/) (2026)
- [JEDI: Joint Embedding Diffusion World Model for Online Model-Based Reinforcement Learning](https://mlsanity.com/paper/2605.13013/) (2026)
- [Object-Centric World Models Meet Monte Carlo Tree Search](https://mlsanity.com/paper/2601.06604/) (2026)
- [When Object-Centric World Models Meet Policy Learning: From Pixels to Policies, and Where It Breaks](https://mlsanity.com/paper/2511.06136/) (2025)
- [Causal Object-Centric Models for Planning with Monte Carlo Tree Search](https://mlsanity.com/paper/2606.14418/) (2026)
- [Better Slots, Better Worlds: Representation Quality & Robustness in Object-Centric World Models](https://mlsanity.com/paper/2608.12078/) (2026)
- [Lucid Dreaming for World Models: Learning to Doubt Imagination and Decide by Trust](https://mlsanity.com/paper/2609.37156/) (2026)

## Abstract (arXiv)

While deep reinforcement learning (RL) from pixels has achieved remarkable success, its sample inefficiency remains a critical limitation for real-world applications. Model-based RL (MBRL) addresses this by learning a world model to generate simulated experience, but standard approaches that rely on pixel-level reconstruction losses often fail to capture small, task-critical objects in complex, dynamic scenes. We posit that an object-centric (OC) representation can direct model capacity toward semantically meaningful entities, improving dynamics prediction and sample efficiency. In this work, we introduce OC-STORM, an object-centric MBRL framework that enhances a learned world model with object representations extracted by a pretrained segmentation network. By conditioning on a minimal number of annotated frames, OC-STORM learns to track decision-relevant object dynamics and inter-object interactions without extensive labeling or access to privileged information. Empirical results demonstrate that OC-STORM significantly outperforms the STORM baseline on the Atari 100k benchmark and achieves state-of-the-art sample efficiency on challenging boss fights in the visually complex game Hollow Knight. Our findings underscore the potential of integrating OC priors into MBRL for complex visual domains. Project page: https://oc-storm.weipuzhang.com