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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Zhang, Weipu, Jelley, Adam, McInroe, Trevor, Storkey, Amos, Wang, Gang
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866914348658589696
author Zhang, Weipu
Jelley, Adam
McInroe, Trevor
Storkey, Amos
Wang, Gang
author_facet Zhang, Weipu
Jelley, Adam
McInroe, Trevor
Storkey, Amos
Wang, Gang
contents 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
format Preprint
id arxiv_https___arxiv_org_abs_2501_16443
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Object-Centric World Models from Few-Shot Annotations for Sample-Efficient Reinforcement Learning
Zhang, Weipu
Jelley, Adam
McInroe, Trevor
Storkey, Amos
Wang, Gang
Machine Learning
Computer Vision and Pattern Recognition
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
title Object-Centric World Models from Few-Shot Annotations for Sample-Efficient Reinforcement Learning
topic Machine Learning
Computer Vision and Pattern Recognition
url https://arxiv.org/abs/2501.16443