Objects matter: object-centric world models improve reinforcement learning in visually complex environments
reinforcement learning;model-based RL;object-centric RL;video object segmentation;Atari;Hollow Knight
Yes
No ethics review needed.
Q1. What is the MBRL training objective? It seems to not be clearly defined. The Preliminaries are melded together with the Related Work, lacking some essential concepts, terminologies, and notations used throughout the paper. 

Q2. How to specify K in more complex environments with a higher number of potential key objects?

Q3. How does the method perform under true partial observability when not all key objects are simultaneously present on the screen?

Q4. Table 2 categorizes Atari games based on whether all key information for the decision can be represented as objects. How is this determined?

Q5. What are potential realistic applicable RL scenarios where OC-STORM method can be useful?

Q6. If evaluating other existing methods is not feasible, why choose Hollow Knight as a benchmark? Wouldn't selecting a more standardized environment allow for clearer comparisons and strengthen the study's validity?

Q7. Does OC-STORM outperform STORM on Hollow Knight only in terms of sample efficiency, or can it achieve comparable overall performance once fully converged? According to the training curves in Appendix D.6, STORM is unable to converge with 100K samples. It would be insightful to compare the methods with an increased number of samples to assess their true performance potential.
S1. The paper is well-written, comprehensible, and offers a profound review of related work within object extraction and MBRL, effectively situating the proposed method within this broader context. The authors adequately explain their motivations for selecting Cutie as the object feature extractor, while outlining certain shortcomings of prior approaches.

S2. The use of a categorical VAE helps ensure more stable and interpretable latent representations, particularly suited for environments where distinguishing between discrete elements is crucial. It enhances the agent’s ability to model structured features effectively.

S3. The spatial-temporal transformer can model both inter-object relationships within a single timestep and the evolution of these objects across time. The spatial block enables understanding interactions between objects at a given moment, while the causal temporal block ensures robust prediction of an object’s trajectory over time, using only past and present information. By combining these, the model can capture both immediate spatial context and long-term temporal dynamics to predict future states. Further incorporating agent actions as a control signal allows the transformer to factor in how an agent’s behavior influences object states. Varying the number of key objects (K∗) adds to the adaptability, enabling effective processing whether using object features, visual input, or both.

S4. The illustrations in the paper are informative. Figure 2 gives a detailed overview of the model's structure, showing how the object module and visual module interact within the architecture. Figure 3 demonstrates the model's ability to capture the state and position of objects, reinforcing the quantitative evidence of its performance. The illustrations in Appendices H and I provide a comparison between the ground truth sample segmentation masks and the ones produced by Cutie.

S5. The method outperforms regular STORM in the Hollow Knight boss fights, and SOTA MBRL baselines in many Atari environments.

S6. The ablation between the module combinations motivates the usage of both vector and visual model components.
The paper introduces OC-STORM, an object-centric MBRL pipeline designed to improve sample efficiency in environments with complex visuals. While traditional MBRL methods with auto-encoding losses can miss small but crucial details in dynamic scenes, OC-STORM addresses this by focusing on key decision-related objects. The pipeline involves annotating key objects, extracting their features using Cutie, a pre-trained vision model, and incorporating these features into the world model for policy training. OC-STORM outperforms it’s predecessor STORM on Hollow Knight and surpasses other SOTA MBRL baselines on many Atari games.
**W1. Motivation**. The authors highlight the theoretical benefits and performance enhancements of their method, such as improved sample efficiency and superior benchmark results. However, there is no discussion on the practical applications of the object-centric MBRL pipeline in potential realistic settings.

**W2. Determining key object annotations.**
1. While using manual annotations to inform agents about key objects is reasonable, determining the optimal number of these objects—represented by **K**—remains unclear. In simple Atari games like *Pong*, setting **K** is heuristically straightforward. However, this task becomes challenging in more complex environments. Users might need to tune **K** for optimal performance, which can be computationally expensive. Given that **K** is integral to the method's effectiveness, it would be beneficial for the authors to assess the method's sensitivity to different **K** values across various settings. This limitation is particularly pronounced in environments where the number of relevant objects is uncertain or fluctuates. An incorrect choice of **K** can lead to under-representation—missing key features—or over-representation—wasting capacity on less important elements.

2. In all tasks from Atari and *Hollow Knight*, the number of annotated objects ranges between 2-4. However, there is no detailed discussion on how these values were determined beyond noting that they are based on heuristics, and it is unclear whether other values were explored. More complex environments may have a greater number of essential objects requiring annotation, raising the question of whether the method can effectively handle such cases. Reliance on a low number of annotated objects (low **K**) may limit the method's applicability, as it could struggle to incorporate and process a higher number of annotated objects crucial for success in more complex settings.

3. The method appears to be limited to scenarios where all annotated objects appear in the same frame—a property common to all Atari games. Similarly, most boss fights in *Hollow Knight* feature a fixed background apart from some minor visual disparities. Only with Pure Vessel and Mage Lord, the screen will slightly pan left or right with the player's horizontal movement. Regardless of this panning, the key objects necessary for the agent's success can always appear in the frame. As shown in Figure 15, the annotations typically include the agent, the boss, and optionally any projectiles the boss might fire. It is straightforward to cherry-pick and annotate the most useful frames where both the boss and harmful projectiles are visible.  
However, this approach becomes unfeasible in environments that are truly partially observable, such as navigating 3D landscapes from an egocentric perspective. Not only does the agent not appear in all sections of the environment, but it also generally has a limited field of view—roughly 90 degrees—occluding most information. This means it is common not to obtain even a single frame where all key objects are visible. This limitation reduces access to informative annotation masks. Moreover, since Cutie attempts to mask the specified number of key objects, it will be incapable of doing so when not all objects are present, resulting in incorrect masks. This raises the question of the applicability of OC-Storm in broader contexts.
        
**W3. Choice of Hollow Knight**. The authors' choice of *Hollow Knight* boss fights as benchmarking tasks is commendable due to the visual and dynamic complexity compared to Atari, along with the game's challenge for both AI and human players. Understandably, since *Hollow Knight* is not an established benchmark, adapting prior MBRL methods to this setting is non-trivial because of significant differences that the authors outline, like sample step limits, resolution, environment wrapping, and reward functions. The authors acknowledge that these differences make direct comparisons with existing methods impractical, which then raises the question of whether *Hollow Knight* is the most appropriate choice for demonstrating the superiority of their approach.
    
While the authors argue that *Hollow Knight*'s visual complexity, dynamic elements, rare duplicates, and less critical backgrounds make it a suitable testbed for object-centric learning, more commonly used 2D platformer environments like **CoinRun** [1] or benchmarks like **ProcGen** [2] could serve the same purpose. Equally, in the 3D realm, game-based platforms such as **ViZDoom** [3], **DMLab** [4] or more realistic environment suites like **Habitat** [5], and **AI2Thor** [6] offer environments with egocentric perception. These environments feature objects and entities that agents need to avoid, interact with, or navigate toward, making them suitable for object-centric RL approaches. Moreover, they have been widely studied in the MBRL domain and come with existing implementations and baseline evaluations, facilitating direct comparisons with prior work.
    
The Mage Lord and Pure Vessel boss fights show the lowest performance and are the only ones where the screen pans as the player moves left and right. The other boss fights maintain a fixed perspective and consistent background, similar to most Atari games. The lower performance in these panning scenarios suggests that the method may require static backgrounds with minimal variations for effective object-centric learning. To demonstrate the robustness of OC-Storm, it would be valuable to assess the method in environments with more dynamic backgrounds and highly disentangled consecutive frames, such as those found in embodied perception environments.
    

[1] Cobbe, Karl, et al. "Quantifying generalization in reinforcement learning." *International conference on machine learning*. PMLR, 2019.

[2] Cobbe, Karl, et al. "Leveraging procedural generation to benchmark reinforcement learning." *International conference on machine learning*. PMLR, 2020.

[3] Kempka, Michał, et al. "Vizdoom: A doom-based ai research platform for visual reinforcement learning." *2016 IEEE conference on computational intelligence and games (CIG)*. IEEE, 2016.

[4] Beattie, Charles, et al. "Deepmind lab." *arXiv preprint arXiv:1612.03801* (2016).

[5] Szot, Andrew, et al. "Habitat 2.0: Training home assistants to rearrange their habitat." *Advances in neural information processing systems* 34 (2021): 251-266.

[6] Kolve, Eric, et al. "Ai2-thor: An interactive 3d environment for visual ai." *arXiv preprint arXiv:1712.05474* (2017).
Yes
No ethics review needed.
- How the mask for 6-shot and 1-shot be for the experiments in the Figure 11?

- What is the additional computational cost compared to the base MBRL model for calculating the object extractor and the added features?
- The manuscript is well-written and easy to read, making it straightforward to understand the four key steps of the proposed method.

- By training STORM using both object features obtained through the object extractor and resized observations, it shows improved performance over the original MBRL algorithm, STORM.

- In addition to quantitative evaluations, appropriate qualitative results are also provided.
In this work, the authors propose a object-centric model-based reinforcement learning pipeline, which can capture small or dynamic elements than can be critical for decision making. The authors point out the limitations of the previous MBRL methods, which often rely on auto-encoding with an L2 loss, which could be vulnerable to capture the moving objects, or small parts in the scenes. The proposed method, OC-STORM, provides a object-centric MBRL pipeline done by four steps: (1) annotate key objects related to rewards and goals, (2) extract object features using a pre-trained foundation model, (3) incorporate object features with the raw observations and predict dynamics, and (4) train the policy. As a result, the proposed method shows improved performance on diverse Atari Benchmarks.
- Rather than comparing with Dreamer/STORM/DIAMOND, the proposed method need to be compared with algorithms where additional representation techniques are applied. For instance, MWM [1] is a very similar work, which added a masked autoencoder to Dreamer and showed it also could capture small moving blocks in the scenes. Algorithms like MWM would actually be more suitable baselines.

- Additionally, although not MBRL, there are numerous studies that consider the dynamics of the environment  in representation part [2]. Rather than comparing with these representation learning methods, this paper mainly focuses on which aspects it improves over normal MBRL. As a result, it lacks persuasive arguments for why the proposed method is better than other environment dynamics-aware representation learning methods.

- For these reasons, the description in the related works section is somewhat lacking.

- Appendix E.2 includes an ablation study on the number of annotation masks, where the authors claim that "a greater number of masks contributes to more robust performance." However, even in static environments like "Atari Boxing," where at least two masks (for the player and the opponent) seem necessary, using only one mask ultimately achieves a similar convergence value to using six masks. Therefore, I would like to ask the authors to share results with a wider range of mask counts and in more diverse environments.

[1] Seo, Younggyo, et al. "Masked world models for visual control." *Conference on Robot Learning*. PMLR, 2023.

[2] Cui, Zichen Jeff, et al. "DynaMo: In-Domain Dynamics Pretraining for Visuo-Motor Control." *arXiv preprint arXiv:2409.12192* (2024).
Yes
No ethics review needed.
- Why do you base your method on STORM, but not DreamerV3? Given that DreamerV3 has better overall performance on multiple domains. Will OC-DreamerV3 perform better than DreamerV3?
- How applicable is OC-STORM when faced with a moving camera, or where not all objects appear in a single frame? This situation is common in open-world environments like Minecraft.
- Can you demonstrate results in domains other than video games? For example on the robotics domain?
- Why are the carrots in gopher not specified as key objects? Is it because they are stationary?
- The method is a straightforward extension of an existing model-based RL algorithm and should be easily replicable.
- Experiments results across two video game benchmarks show that OC-STORM effectively improves the performance of the base method STORM.
- Modifications to the Hollow Knight benchmark can help the agent learn a more meaningful policy and should be regarded as a valuable contribution to the research community.
This paper introduces OC-STORM, which incorporates object-centric state-based representations obtained from a pre-trained visual model into the MBRL method STORM. The OC-STORM extractor requires a few hand-annotated segmentation masks to identify key objects before extracting their features. Experimental results indicate that OC-STORM outperforms STORM on both Atari 100K and Hollow Knight benchmarks.
- The core idea is relatively simple. Only subsection 3.1 in the method section contain novel contributions (i.e. introducing object features as part of the input), while the remaining sections elaborate on DreamerV3's (and STORM's) model architecture and training methodology. 
- The generality of the object-centric MBRL pipeline is unclear. It is uncertain how the effectiveness of OC-STORM extends beyond the specific STORM architecture. Further experiments combining this pipeline with other MBRL methods, such as DreamerV3, should be conducted.
-  The experimental results are insufficient to support the paper's claim. While the paper mentions "visually complex environments" in its title, experiments are limited to two video game benchmarks and exclude robotics control benchmarks, where scenarios with visual distractions are prevalent, as seen in Natural Background DMC[1] and DMC Remastered[2]. Atari, one of the benchmarks utilized, may not be regarded as a visually complex environment. Within the video game domain, Minecraft benchmarks such as Minedojo[3] feature visually complex elements closer to real-world scenarios and should more suitable for inclusion in this paper. Additional experiments on diverse benchmarks are necessary to support the paper's argument.
- The scope of applicability is unclear. The extractor OC-STORM uses relies on hand-annotated segmentation masks. While these mask are easy to obtain in Atari and Hollow Knight, it may may pose challenges in open-world tasks due to complex object interactions. Beyond acquisition difficulties, the impact of mask quality on the extractor's accuracy remains an unresolved issue. Authors should address these challenges or acknowledge them as part of the limitations. 

[1] Zhang et al. Natural environment benchmarks for reinforcement learning.

[2] Grigsby et al. Measuring visual generalization in continuous control from pixels.

[3] Fan et al. Minedojo: Building open-ended embodied agents with internet-scale knowledge.
Yes
No ethics review needed.
1. The selection of Cutie as the object detector is well-justified, given its robustness in generalizing across diverse environments. However, evaluating additional object detection models or conducting an ablation study would provide valuable insights into the flexibility and performance of different detectors within OC-STORM. Could the authors share experimental results that compare the effectiveness of various object detection models in this context?

2. Integrating object detection and feature extraction may introduce computational overhead. A comprehensive analysis of the trade-offs between computational cost and performance gains in complex environments would enable potential users to better understand OC-STORM's scalability and efficiency. Could the authors provide experimental results detailing this trade-off?
1. The paper introduces a unique method of integrating object-based features with reinforcement learning, allowing the agent to prioritize decision-relevant elements within complex visual environments. This approach addresses a significant gap in traditional pixel-based MBRL methods by enabling more focused training.

2. The experiments cover both Atari games and Hollow Knight, demonstrating the versatility and practical relevance of OC-STORM. This shows that the method is not limited to simplified or controlled environments but can extend to visually and structurally complex settings.
The paper proposes OC-STORM, an object-centric model-based reinforcement learning (MBRL) framework designed to improve sample efficiency and performance in visually complex environments. Unlike conventional MBRL methods that rely on pixel-based auto-encoding, OC-STORM leverages recent advances in object detection to focus on key, decision-relevant elements within scenes. The approach uses an object-centric pipeline that includes object annotation, feature extraction, and training with simulated trajectories, enhancing reinforcement learning in environments such as Atari games and Hollow Knight. The paper demonstrates OC-STORM's performance improvement over existing MBRL algorithms, particularly in settings where object information is essential for decision-making.
See question.
Yes
No ethics review needed.
- As mentioned in the weaknesses, how scalable is this to visual domains with continuous action spaces and possibly to real-world tasks?
- Good discussion on the thought process in selecting the model for extracting object representations and other hyperparameter/design choices
- Paper is well-motivated, reasoning for using an object-centric WM is justified
- Good related work section discussing literature in MBRL and object-centric representations for RL
- Results are convincing, good comparisons to relevant baseline methods
- Nice qualitative analysis highlighting what objects are being segmented by the object model in the evaluation environments
- Good discussion of potential limitations of the work
This paper proposes an object-centric world model to handle environments that contain small, dynamic objects that are likely important for decision-making. Their approach involves annotating key objects using segmentation masks, extracting the object features using a pretrained model, training a world model with these features and finally training a policy using MBRL. They claim to be the first to adopt object-centric learning on Atari and the visually complex game of Hollow Knight. They outperform the prior MBRL work on 18 of 26 tasks in the Atari 100k benchmark.
- Is the primary novelty in the paper the use of the pretrained features from Cutie in a MBRL framework? It is not clear apart from this what other contributions are presented in this work.
- Environments are carefully selected such that an object-centric approach has an advantage. I wonder whether this approach could be applicable in continuous control tasks and potentially real-world robot tasks. It would be nice to see such results in a future study.
- It would be nice to see a model-free baseline there for comparison as well. It is not clear to me what the benefit is for MBRL in these environments where the simulator is readily available. If it is about sample-efficiency, it would be good to see how this stacks against your standard SAC or PPO.
Yes
No ethics review needed.
1. Additional Overhead for Inference Time: In OC-STORM, the segmentation model is used as a pre-processing step to extract object features. This likely introduces extra overhead during both training and inference, particularly given the upscaled observations and the large backbone of the foundation model. Can you clarify how significant this overhead is, especially during test time?

2. Analysis of Segmentation Model Quality and Errors: Since segmentation foundation models are used, some segmentation errors are inevitable. The paper addresses a few aspects of this: failure cases of segmentation models are presented qualitatively (Fig. 12), and the effect of annotation count is discussed (Fig. 11). However, a more detailed analysis of segmentation model quality in relation to downstream control performance would strengthen the work. For example, examining the correlation between segmentation quality and control performance could provide insights into whether advanced segmentation would lead to better results. Additionally, analyzing segmentation model failures during test time could help evaluate the method’s robustness to segmentation errors, which is crucial for safe deployment.

3. Other Segmentation Models: Related to the second aspect, it would be beneficial to explore how OC-STORM performs when combined with other segmentation models, such as SAM2.

4. Other MBRL Models: Since the main contribution is the incorporation of additional input into the world model, I believe this approach could be applied to other model-based reinforcement learning methods, such as Dreamer. While some modifications may be necessary due to the non-symmetric nature of non-Transformer latent dynamics models, demonstrating that the inclusion of object features can consistently enhance performance across various base MBRL models would strengthen the argument.

5. Other Domains: While the proposed method performs well in video game environments, it may prove even more advantageous in other, more object-centric domains, such as object manipulation, where decision-relevant objects are more unique.

6. Missing Information in Table 1 and Table 2: Without standard deviation or error, it’s hard to assess the reliability and consistency of the results presented.
1. The authors present an effective approach to overcoming challenges in learning control policies for visually complex, object-centric tasks by leveraging recent advances in computer vision. While foundation models for RL/robotics are still emerging and there is no established standard for learning useful representations for control agents, this paper shows that inheriting knowledge from vision foundation models can be a promising direction. The authors seamlessly integrate the segmentation model with RL, requiring a little bit of task-specific prior knowledge and few annotations.
2. The authors conduct comprehensive experiments across a range of video game tasks in two different domains. While their method does not consistently outperform others, they demonstrate that OC-STORM performs effectively in object-centric games by categorizing games into two groups, as shown in Table 2.
3. The paper is well-written and easy to follow, with prior works thoroughly and up-to-date in their presentation.
The authors present an object-centric MBRL pipeline that enables agents to focus on key, decision-relevant elements. They leverage recent advances in computer vision by extracting object features using a pre-trained segmentation foundation model with few annotations (six annotations in their experiments). These extracted features are incorporated into the world model as an additional modality along with raw observations. The authors demonstrate that their method (OC-STORM) is more effective than the baseline (STORM) for object-centric tasks in Atari games and Hollow Knight.
1. Although prior works are well presented, comparisons with them are missing in the experiments. While Table 1 compares the method with other general-purpose MBRL approaches, it would be valuable to see how OC-STORM performs against methods specifically designed for object-centric policy learning, such as [1], [2], [3], and [4]. Without these comparisons, it is challenging to fully assess the method’s effectiveness relative to similar approaches.
[1] Stefano Ferraro, Pietro Mazzaglia, Tim Verbelen, and Bart Dhoedt. FOCUS: Object-Centric World Models for Robotic Manipulation. In Intrinsically-Motivated and Open-Ended Learning Work- shop @ NeurIPS2023, November 2023
[2] Jaesik Yoon, Yi-Fu Wu, Heechul Bae, and Sungjin Ahn. An Investigation into Pre-training Object- centric Representations for Reinforcement Learning. In International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pp. 40147–40174. PMLR, 2023. URL https://proceedings. mlr.press/v202/yoon23c.html.
[3] Akihiro Nakano, Masahiro Suzuki, and Yutaka Matsuo. Learning Compositional Latents and Behaviors from Object-Centric Latent Imagination. In The 38th Annual Conference of the Japanese Society for Artificial Intelligence, JSAI 2024, 2024.
[4] Younggyo Seo, et al. Masked world models for visual control. Conference on Robot Learning. PMLR, 2023.
2. Although the motivation and design of the architecture and pipeline are sound, the experimental results are not particularly impressive or convincing. Given the advantages of prior knowledge on decision-related objects, few annotations, and the use of foundation models, one would expect a more significant performance boost over methods that lack these benefits, particularly in terms of final performance and sample efficiency (see Table 1 and Fig. 4).
We propose an object-centric model-based RL pipeline, which integrates recent advances in computer vision to allow agents to focus on key decision-related elements.
@inproceedings{
anonymous2024objects,
title={Objects matter: object-centric world models improve reinforcement learning in visually complex environments},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=Q2hkp8WIDS},
note={under review}
}
Deep reinforcement learning has achieved remarkable success in learning control policies from pixels across a wide range of tasks, yet its application remains hindered by low sample efficiency, requiring significantly more environment interactions than humans to reach comparable performance.
Model-based reinforcement learning (MBRL) offers a solution by leveraging learnt world models to generate simulated experience, thereby improving sample efficiency.
However, in visually complex environments, small or dynamic elements can be critical for decision-making.
Yet, traditional MBRL methods in pixel-based environments typically rely on auto-encoding with an $L_2$ loss, which is dominated by large areas and often fails to capture decision-relevant details.
To address these limitations, we propose an **object-centric MBRL pipeline**, which integrates recent advances in computer vision to allow agents to focus on key decision-related elements.
Our approach consists of four main steps: (1) annotating key objects related to rewards and goals with segmentation masks, (2) extracting object features using a pre-trained, frozen foundation vision model, (3) incorporating these object features with the raw observations to predict environmental dynamics, and (4) training the policy using imagined trajectories generated by this object-centric world model.
Building on the efficient MBRL algorithm STORM, we call this pipeline **OC-STORM**.
We demonstrate OC-STORM's practical value in overcoming the limitations of conventional MBRL approaches on both Atari games and the visually complex game Hollow Knight.
Code and videos are available in the supplementary materials.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
reinforcement learning model-based RL object-centric RL video object segmentation Atari Hollow Knight
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/b09b7fcb6935453d3a8bfb4797f4b1c1fff45419.pdf
reinforcement learning
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
/attachment/d4fca91e4e349fde1591174c4c1898d5f95e5eb9.zip
Objects matter: object-centric world models improve reinforcement learning in visually complex environments
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission