Dyn-O: Building Structured World Models with Object-Centric Representations

Thank you for your thoughtful and constructive comments. We’re glad that you recognize the potential of leveraging large-scale pretrained computer vision models like SAM2 within world model learning, and appreciate our efforts to narrow the gap in this underexplored area. We also thank you for acknowledging our paper’s well-organized presentation as well as its effective illustrations and visualizations. Below, we address your comments and questions.

W1: The experiments rely on trajectories generated by a pretrained PPG agent, raising concerns about applicability in online settings:

• Does the method work in an online RL setting like DreamerV3? Online world models must learn from dynamically changing data distributions, and it is unclear whether the proposed approach would remain stable.
• Could training collapse when learning from scratch? Pretrained policies provide high-quality exploration data, but in online RL, poor initial exploration might lead to degenerate representations.

R1:

• We would like to clarify that our focus is not on online reinforcement learning, but rather on building world models based on object-centric representations. In earlier experiments, we found that training the entire pipeline online was prohibitively resource-intensive, so we decided not to pursue that direction. Notably, many existing works also focus on the offline setting $1-4$ , so we believe our approach still offers meaningful contributions to the field.
  We agree that investigating the stability of the proposed method under dynamically changing data distributions would be insightful. However, doing so would require substantial engineering efforts to significantly optimize our pipeline. We view this as a valuable direction for future work.
• Thank you for the interesting question! First, we would like to clarify that the training data we used were not collected by an expert policy, but rather by different intermediate checkpoints saved during RL training. As a result, the dataset contains a significant portion of suboptimal exploration data.
  To further investigate this, we conducted an ablation experiment in the Bigfish environment, using only data collected from a random policy to train the representation. This was done to assess whether representation collapse would occur. The results are summarized in the table below. We can see using random exploration data leads to slightly worse, but not significantly degraded, performance.

Table 1: Rollout accuracy of Dyn-O trained with different data at 20-th timestep in Procgen Bigfish, measured as mean and standard error

We believe the extent to which representation quality is affected depends heavily on the properties of the environment. In many of the environments we use, most objects remain visible even under random policies, making representation collapse unlikely. However, in particularly challenging environments that require extensive exploration (e.g., Montezuma’s Revenge), such collapse may occur. That said, effective exploration in such environments is itself a broader and open research question, which we consider beyond the scope of our current work.

$1$ Jiang, Jindong, et al. "Object-centric slot diffusion." arXiv preprint arXiv:2303.10834 (2023).

$2$ Wu, Ziyi, et al. "Slotdiffusion: Object-centric generative modeling with diffusion models." Advances in Neural Information Processing Systems 36 (2023): 50932-50958.

$3$ Wu, Ziyi, et al. "Slotformer: Unsupervised visual dynamics simulation with object-centric models." arXiv preprint arXiv:2210.05861 (2022).

$4$ Singh, Gautam, Yi-Fu Wu, and Sungjin Ahn. "Simple unsupervised object-centric learning for complex and naturalistic videos." Advances in Neural Information Processing Systems 35 (2022): 18181-18196.

W2: The paper does not explicitly describe whether reconstruction metrics are computed on training or unseen test trajectories. Is the reported accuracy based on training rollouts? If so, the improvements could stem from overfitting rather than better generalisation.

R2: We would like to point out that, Line 236 of our paper describes that the evaluation is performed on unseen test levels. Therefore, the observed improvements are unlikely to result from overfitting.

Q1: The paper would be significantly more convincing if it included experiments evaluating the agent’s performance, either in a model-based or model-free setting. Currently, the focus on reconstruction accuracy leaves open whether the learned representations actually improve downstream policy learning. I would be willing to accept if this point is covered.

R3: Thank you for the great suggestion and for your willingness to accept our paper if this point is addressed! This is indeed an important aspect that could further strengthen our paper. During the rebuttal period, we conducted additional policy learning experiments as recommended. Specifically, to evaluate whether the learned representations facilitate policy learning, we fixed the representation backbone and trained only the policy head using the PPG algorithm on three Procgen games: Bigfish, Starpilot, and Coinrun. We used the “easy” difficulty setting with an unlimited number of levels.

We compared the performance of object-centric representations against the representations directly output by the Cosmos encoder. Due to NeurIPS policy, we are unable to include figures at this stage, so we provide a table below summarizing the results at different training steps. Each game is trained with three seeds, and some runs are still in progress. We will update the results during the discussion phase once training is complete, and include the full figures in the final version. Based on the results so far, the learned representations lead to clearly improved policy learning performance.

Table 1: Episodic returns in Starpilot

Table 1: Episodic returns in Bigfish

Table 1: Episodic returns in Coinrun

Q2: The paper evaluates the method on 7 Procgen environments but does not justify how these were chosen from the full benchmark suite. A brief explanation (e.g., diversity of difficulty, visual complexity, or relevance to object-centric learning) would clarify the experimental design.

R4: Thank you for pointing this out. We selected 7 Procgen environments based on their relevance to object-centric learning and visual complexity. Specifically, we looked for environments with moving objects, dynamic layouts, and minimal sensitivity to color. Based on these criteria, we excluded maze-like environments such as Maze, Heist, Chaser, and Miner, as well as Plunder, where the ship's color plays a critical role. Among the remaining games, we randomly selected 7 while ensuring diversity across game types, given resource constraints. We would be happy to evaluate additional games if the reviewer deems it necessary, and we will include the above rationale in the next version.

Thank you for your valuable and insightful feedback. We're glad to hear that you found our methodology and representation clearly presented and easy to follow. We're especially encouraged that you consider our use of an object-centric latent space for world modeling to be innovative and a meaningful contribution to structured representation learning. Below, we address your comments and questions.

W1: Experimental Validation: The paper states: “However, the development of object-centric world models has largely been explored in environments with limited visual complexity (such as basic geometries). It remains underexplored whether such models can be effective in more challenging settings.” I fully agree with this assessment; however, the evaluation on the Procgen dataset $A$ , which comprises procedurally generated 2D video game environments, mirrors the limited complexity of prior datasets $D$ . Therefore, this paper's experimental validation does not address whether the proposed methodology is effective in more challenging, real-world environments.

R1: We agree with the reviewer that real-world videos are the best benchmark for evaluating object-centric world models. However, we would like to clarify that, while Procgen consists of 2D video game environments, it offers significantly greater visual and dynamic complexity compared to prior benchmarks such as CATER $D$ and CLEVR $2$ . Specifically, Procgen features:

• diverse shapes and textures - while most objects in prior benchmarks use a single color and basic shape, as shown in Fig 7 and 9, Procgen objects have more complex shapes (like spaceships, trails) and color combinations (like explosion).
• cluttered scenarios - in prior benchmarks, most scenarios have <10 objects, while in Procgen, environments like Starpilot can have ~50 objects.
• complex dynamics - in prior benchmarks, most movements are constrained to linear motion with constant velocity. In contrast, Procgen features complex motions (like jump) and mechanisms (like starship explosion).

Therefore, evaluation on Procgen meaningfully advances the study of object-centric world models toward more challenging settings, and exploring real-world videos remains an exciting avenue for future work.

$1$ Girdhar, R. and D. Ramanan. “CATER: A Diagnostic Dataset for Compositional Actions and Temporal Reasoning.” arXiv:1910.04744, 2019.

$2$ Johnson, Justin, et al. "Clevr: A diagnostic dataset for compositional language and elementary visual reasoning." Proceedings of the IEEE conference on computer vision and pattern recognition. 2017.

W2: More critically, the method is tested solely on Procgen, without comparing its performance against baselines on datasets commonly used in prior work $B, C$ . This omission prevents firm conclusions about whether the proposed methodology advances the field relative to existing approaches. Additionally, the world model’s experimental evaluation is limited beyond the object-centric representation. Could the authors clarify why broader datasets, such as the Arcade Learning Environment (ALE) $E$ , were not included?

R2: Thank you for the helpful feedback! During the rebuttal period, we conducted additional experiments by running our Dyn-O model on the CLEVR dataset. Regarding ALE, we believe Procgen serves as an enhanced benchmark, offering comparable or even greater complexity, with the added advantage of being designed with generalization in mind. As a result, models trained on Procgen are less prone to overfitting. That said, we were able to run experiments on two ALE games, Boxing and Skiing, during the rebuttal period. Detailed results for both the CLEVR and ALE experiments are provided in our response to the next comment.

W3: Unclear Relation to Prior Art: The paper presents two core contributions: object-centric representations and the world model. However, it lacks clarity on how these components compare, conceptually and empirically, to prior work. While the paper references DreamerV3 as a baseline, it reports results only on Procgen, whereas DreamerV3 demonstrates performance across multiple datasets, Procgen, ALE $E$ and several others with greater diversity than 2D games. This limited comparison undermines the paper’s claim of advancing the state of the art.

R3: Thank you for the suggestion! Following your recommendation, we conducted experiments comparing Dyn-O and DreamerV3 on CLEVR and two ALE games. As the results below show, Dyn-O outperforms DreamerV3 in ALE, highlighting the advantages of object-centric world models. The experiments on CLEVR are currently in progress, and we will update the results as soon as they are available.

Table 1: Rollout accuracy at 20-th timestamp in CLEVR, measured as mean and standard error

Table 2: Rollout accuracy at 20-th timestamp in Atari Skiing, measured as mean and standard error

Table 3: Rollout accuracy at 20-th timestamp in Atari Boxing, measured as mean and standard error

Q1: The first sentence in the abstract is: "World models aim to capture the dynamics of the environment". Is this a fair statement? Intuitively, the first thing world model must learn is the static aspect of the world. Not to say learning dynamics isn't important -- it absolutely is -- but this emphasis doesn't sound right to me.

R4: We agree that a world model can capture both the dynamic aspects (transition dynamics) and the static aspects (representation) of an environment. We wrote this statement because most prior work on world models has focused on the dynamic aspect, while the representation side, such as learning object segmentations and image understanding, has been more extensively explored in the broader deep learning and computer vision communities. We are happy to revise this statement based on your suggestions.

Q2: Abstract states: "the dynamics are highly centered on interactions among objects within the environment". Could the rebuttal clarify this? I am confused by this statement.

R5: Thank you for your feedback! By this statement, we mean that most changes in the environment (its dynamics aspect), are either caused by objects or affect them. For example, in the Procgen Coinrun game, changes in the game screen are primarily driven by the agent’s movement or interactions with other objects, such as collecting coins, getting hurt by enemies, or being blocked by chests. We are happy to revise the statement according to your suggestions.

We sincerely appreciate your insightful comments. We're glad to hear that you find our ideas novel and recognize the careful design of the loss functions. Below, we address the questions you raised.

W1: The experimental evaluation of the method seems not very thorough. For example, the comparison is only against DreamerV3. There is no other world model evaluated

R1: Thank you for the suggestion! During the rebuttal period, we further compared our Dyn-O with Dreamweaver $1$ , a recent object-centric world model. The experiments are currently in progress, and we will update the results as soon as they are available.

$1$ Baek, Junyeob, et al. "Dreamweaver: Learning Compositional World Models from Pixels." arXiv preprint arXiv:2501.14174 (2025).

W2: The experimental environments are all about games. More sophisticated setups can be considered such as using real-world videos.

R2: We agree that evaluating on real-world videos would further demonstrate the applicability of our object-centric world model and represents an exciting direction for future work. One avenue we plan to explore is evaluating Dyn-O on robotic datasets such as Bridge-v2 $2$ and RT-1 $3$ to assess its performance in more realistic settings. However, due to time and compute constraints, we were unable to include these experiments during the rebuttal period. We intend to pursue them in future extensions of this work.

$2$ Walke, Homer Rich, et al. "Bridgedata v2: A dataset for robot learning at scale." Conference on Robot Learning. PMLR, 2023.

$3$ Brohan, Anthony, et al. "Rt-1: Robotics transformer for real-world control at scale." arXiv preprint arXiv:2212.06817 (2022).

W3: Some ablation studies can be added. For example, how important is using the state-space model? What if some other models are used?

R3: Thank you for the suggestion! We conducted an ablation study on the backbone model for Dyn-O in the Bigfish environment, replacing the state-space model (SSM) with a transformer. As shown below, the SSM significantly outperforms the transformer.

Table 1: Rollout accuracy at 20-th timestep in Bigfish, measured as mean and standard error

Thank you for your valuable comments and feedback! We're glad to hear that you found our idea interesting. Below, we address your questions.

Q1: My main question is how the findings differ and/or generalize across the different Procgen environments. For instance, Table 1 reports results aggregated for all environments, but it would be useful to assess how Dyn-O performs vs Dreamer or the non-object-centric version across all environments. Similarly, evaluating the object-centric representations is only done on BigFish and the static/dynamic disentanglements only on CoinRun – how do these findings generalize to the other environments?

R1: Due to space limit, the results for each Procgen environment are deferred to the supplementary materials. Specifically,

• Dynamics rollout accuracy: As shown in Table 6 (Page 15), Dyn-O outperforms both Dreamer and the non-object-centric baseline in most environments. Qualitative results are provided in Fig. 8 - 10.
• Object-centric representation quality: As shown in Table 5 (Page 14), Dyn-O outperforms SOLV in 3 out of 4 environments. Corresponding qualitative results can be found in Fig. 7.
• Static/dynamic disentanglement: As shown in Tables 7 - 10, Dyn-O successfully disentangles static and dynamic features. In particular, for static privileged properties, static features achieve significantly higher probing accuracy than dynamic features, and vice versa.

We will add clarification in the main text to help reduce potential confusion.