Learning Interactive World Model for Object-Centric Reinforcement Learning

Thank you for your constructive feedback and review! We truly appreciate your recognition of our methodology. Most of the concerns relate to technical details and clarity. For clarity-related issues, we provide concrete and executable plans that can be incorporated into the camera-ready version. For the technical concerns, we offer detailed clarifications along with additional evaluations in this rebuttal. Please find our responses below.

---

The paper is rich in technical detail, but the overall readability could be enhanced.

R1: Thank you for the constructive feedback. To improve readability, we have devised a concrete revision plan for the camera-ready version, addressing each question below. Our focus is on clarifying the technical content and improving the overall flow, as suggested. Please refer to the tags “[For clarity]” in our responses to see where these revisions are proposed.

---

It would be valuable to include some qualitative visualizations of the disentanglement effect, which could make the results more intuitive and compelling.

R2: Regarding the qualitative visualization of the disentanglement effect, we agree that the illustrations could be useful. We cannot provide it here as we are not allowed to add additional figures during rebuttal. In the camera-ready version, we will make sure to include visualizations of the reconstructed images to provide a more direct qualitative assessment. These will be included in the appendix.

We would also like to emphasize that in the current submission, we primarily evaluated disentanglement through probing results (see Fig. 5(a)), as our goal is to learn disentangled representations directly from pre-trained models such as DINO and R3M, indicating that our method can extract structured object-centric information without additional supervision. In addition, we also reported LPIPS scores (Table 1), which reflect pixel-level reconstruction quality.

---

Figure 1 suggests that a low-level interaction policy is learned during the offline model learning stage. However, the description in Section 3.1 focuses on the world model's components and losses, and the discussion of this policy's learning is deferred to Stage 2. A clarification on whether the low-level policy is indeed trained offline.

R3: The low-level policy $\pi^l$ is trained during the offline stage and then fine-tuned in the online stage. We acknowledge that this setup may have caused some confusion, as the details for learning $\pi^l$ are presented in Section 3.2.

Our original intention was to clearly separate the components: first, explaining how the world model (including dynamics, observation, and reward) is learned, and then detailing how the hierarchical policy is constructed, in order to make the framework easier to follow.

[For clarity] In the revision, we will address this potential misunderstanding by adding a clarifying sentence after Line 166: “Note that the learning of $\pi^l$ is performed in the offline stage and subsequently fine-tuned during the online stage.”

And we plan to add the algorithm to the appendix (here is the brief sketch):

Offline World Model and Online Policy Learning Framework

---

Stage 1: Offline Model Learning
For each $(o_t, a_t, r_t) \in \mathcal{D}$:

• Encode observation: $\hat{o}_t \leftarrow q_\phi(o_t)$
• Infer latent state: $s_t^i \leftarrow q_\phi(s_t^i \mid \hat{o}_t)$
• Factorize latent state: $s_t^i = (d_t^i, c_t^i)$
• Learn reward model: $p_r(r_t \mid s_t, a_t)$
• Learn transition model: $p_s(s_{t+1} \mid s_t, a_t, G_t)$
• Infer interaction graphs: $p_u(G_t \mid s_t)$
• Train interaction policy: $\pi^l(a_t \mid s_t, G_t^g)$

---

Stage 2: Online Policy Learning
For each environment rollout episode:

• Receive current state $s_t$ from world model encoder with input environmental steps
• Select goal graph: $G_t^g \sim \pi^h(G_t^g \mid s_t)$
• Sample action from interaction policy: $a_t \sim \pi^l(a_t \mid s_t, G_t^g)$
• Execute $a_t$ in environment, observe $r_t$, $s_{t+1}$
• Update policies $\pi^h$, $\pi^l$ using collected data

---

In Section 3.2, it is stated that the policy is learned "by leveraging the predicted interactions to learn the inverse mapping from interactions to actions". A more detailed explanation of this inverse mapping process would be beneficial.

R4: The inverse mapping is realized through either MPC or PPO, as described later in Section 3.2 and Appendix C.3. Given a predicted target interaction graph at time $t+k$ (e.g., $G_t^k = G^*$), we treat it as the goal with state s_g and use the learned transition model to infer the target states. For MPC, we optimize the action sequence to minimize the difference between the predicted future state and the inferred target state:$\mathcal{L}\_{\text{MPC}} = \|| s_{t+k} - s_g\||_2^2$.

At each iteration, we sample a population of action sequences from a Gaussian distribution and update the mean and covariance using gradient descent, enabling the policy to reach the desired interaction outcome. For PPO, we use the same reward formulation: $-\|| s_{t+k} - s_g\||_2^2$, and train a goal-conditioned policy to maximize expected return using standard PPO updates. [For clarity] We will add these to the second paragraph of Sec 3.2.

---

The use of Slot Attention is a critical component for learning object-centric representations. Given that the number of objects can vary across different environments, it would be helpful to know how the number of slots is determined. Conducting ablation studies on the effect of the number of slots on the model's performance in different environments would be better.

R5: Thank you for the insightful comment. In our experiments, we heuristically set the number of slots to be slightly larger (typically +2 or +3) than the true number of objects present in the environment.

This allows Slot Attention to stably learn object-specific representations, given its inherent limitations in dynamically adapting to varying object counts. And we think this heuristic is also practical in real-world settings where a rough estimate of object count is often available, either through prior domain knowledge or by incorporating auxiliary tools such as segmentation-based object proposal methods (e.g., SAM).

While slot count selection is orthogonal to the core contributions of our method, we agree that it is helpful to include an ablation. Below, we report performance on the Gym-Fetch (Compositional generalization, 5 objects) task by varying the number of slots:

From the results, we see that using +2 to +4 slots yields stable and comparable performance. Setting the number of slots to exactly the number of objects can limit flexibility, making it harder to handle occlusions or background clutter. Too many slots may lead to redundancy or object fragmentation. Slight over-provisioning offers a good balance between expressiveness and stability.

---

The paper leverages pre-trained vision models like DINO-v2 and R3M to extract high-quality features. A clarification would be helpful as to whether these pre-trained models are fine-tuned during the world model learning stage.

R6: We do not fine-tune the pre-trained vision models (DINO-v2 and R3M) during world model learning, which makes the framework efficient and flexible. Our goal is to build task-specific world models efficiently by leveraging the rich representations already captured by these models. [For clarity] We will clarify this point in the first paragraph of Section 3 in the camera-ready version.

The reconstruction loss is applied in the object-centric embedding space rather than the raw image space. The reasoning behind this choice could be further explained. While this approach is likely computationally efficient, it appears to heavily rely on the quality and generalization of the pre-trained model's features, which might lead to suboptimal learned representations if the pre-trained model's features are not sufficiently expressive for a given complex task.

R7: Yes, we use reconstruction loss in the object-centric embedding space to reduce computational cost while focusing learning on semantically meaningful features. R3M, being pre-trained on robotic data, provides features well-suited for control tasks. DINO-v2 captures strong semantic representations of visual scenes (also validated in DINO-WM (Zhou et al., 2025)). We believe these pre-trained embeddings offer a sufficiently rich foundation, which our disentanglement module further refines to support downstream dynamics learning and decision-making.

---

We hope these responses address your concerns. If so, we would be truly grateful if you would consider adjusting your rating accordingly. Thank you again for your time and thoughtful review.

We thank you for the detailed and constructive review, and we appreciate your recognition of our model contributions and writing. Most of the questions pertain to technical details, which we address through clarifications and additional experiments provided below. In the camera-ready version, we will incorporate these clarifications and include visualizations at the pixel level. We hope these responses address your concerns.

---

There is no visualization of the slot mappings to the environment observations.

R1: We agree that visualizing slot-to-object mappings would help clarify the effectiveness of our object-centric representations. Since in the rebuttal we cannot add figures anymore, we will not be able to provide them here, but in the camera-ready version, we will include such visualizations (e.g., reconstructed images and slot masks).

While our current evaluation focuses on probing results (Fig. 5(a)) to show disentanglement of factors like position, velocity, and color, these fine-grained results suggest that object-level separation is also achieved. This is further supported by the use of the Videosaur backbone, which has shown strong object-wise disentanglement in prior work. That said, we acknowledge that a direct qualitative assessment would strengthen this point, and we will ensure that these visualizations are included in the final version.

[1] Zadaianchuk, Andrii, Maximilian Seitzer, and Georg Martius. "Object-centric learning for real-world videos by predicting temporal feature similarities." NeurIPS 2023.

---

Some essential details are missing from the paper. The paper notes that either MPC or PPO could be used for low-level policies in the second step, but does not indicate which was used in the experiments.

R2: Thank you for pointing this out. We apply MPC for the Gym-Fetch and Franka-Kitchen environments, and PPO for LIBERO and iGibson tasks. We will add this clarification after Line 220 in the revision.

Additionally, we conducted an ablation study to compare the use of PPO and MPC across all environments. The results, which show comparable final success rates, confirm that both approaches are feasible. We will include this ablation table in Appendix D.

---

The architecture is difficult to understand. Including schematic diagrams of the dynamics model and the attributes VAE would greatly improve clarity and help to better understand the structure of the world model.

R3: Thank you for the suggestion. The dynamics model and attribute VAE are currently shown as modules in Fig. 3. To improve clarity, we will include two additional schematic diagrams illustrating their input/output flow and network structures in Appendix D of the camera-ready.

For reference, the detailed loss functions are provided in Section 3.1, and the corresponding hyperparameters are listed in Appendix E.1.

---

In Appendix E.3, it appears that the Videosaur training time is not included. For the sake of fairness, it would be important to account for the slots attention training time as well, particularly since it needs to be trained separately for each environment.

R4: Thank you for the suggestion. You are right that for a fair comparison, the training time of the Slot Attention (Videosaur) module should be reported. We have now included the training times for each environment (all measured on an A100 GPU with 40GB memory):

• Sprites-World: 1 hour
• Fetch: 4 hours
• iGibson: 5 hours
• Libero: 5 hours
• Franka-Kitchen: 4.5 hours

We will add this information to Appendix E.3 in the camera-ready version.

---

It seems that Table A4 is missing the hidden dimensions for the MLP decoder used in Videosaur.

R5: Thank you for catching this. The hidden dimensions for the MLP decoder in Videosaur were omitted from Table A4. We use:

• iGibson, Libero, Franka-Kitchen: 4-layer MLP with 1024 hidden units

• Fetch: 3-layer MLP with 512 hidden units

• Sprites-World: 2-layer MLP with 256 hidden units

These settings are chosen based on the complexity of each environment. We will update Table A4 accordingly in the camera-ready version.

---

The paper mentions using both PPO and MPC frameworks for low-level policies. Are there any recommendations or insights on when to prefer PPO over MPC, or vice versa? If time permits, it would be valuable to see ablation exploring this comparison.

R6: We choose MPC for tasks with relatively simple semantics and shorter horizons (e.g., Gym-Fetch, Franka-Kitchen), and PPO for more complex environments involving long-horizon goals and rich semantics (e.g., LIBERO, iGibson). This decision balances efficiency and learning flexibility. As mentioned previously, we also include an ablation comparing PPO and MPC across environments, which shows that both are feasible with similar final performance. We will add these insights and the corresponding table in Appendix D.

---

The paper claims that the method can handle long-horizon tasks. How long exactly? Would it be feasible for environments requiring thousands or even hundreds of thousands of steps? Additionally, does the performance degrade as the horizon length increases?

R7: By “long-horizon,” we refer to the method’s ability to solve tasks with complex temporal structure by decomposing them into meaningful sub-tasks using learned interaction structures and disentangled features. The actual episode length depends on when the agent completes the task, rather than being fixed. While we have not explicitly tested environments requiring hundreds of thousands of steps, our hierarchical framework is designed to handle increasing task complexity through temporal abstraction. To support this, we evaluate variants of the Franka-Kitchen environment with different numbers of sub-tasks. For example:

• Task 1: Open the microwave → Move the kettle → Turn on the stove → Turn on the light (4 sub-tasks)
• Task 2: Open the microwave → Turn on the light → Slide the cabinet → Open the cabinet → Turn on the top burner (5-subtasks)
• Task 2: Open the microwave → Move the kettle → Turn on the light → Turn on the stove → Open the cabinet → Turn on the top burner (6-subtasks)

These results show that while both methods perform comparably on moderate tasks (4 sub-tasks), FIOC outperforms TD-MPC2 (best-performing baseline for this task) as task complexity increases, showing better scalability and robustness for long-horizon tasks. This supports our claim that the hierarchical design and object-centric representations in FIOC enable more effective temporal abstraction.

---

As stated in the original Videosaur paper, the quality of slot extraction tends to degrade with longer video lengths. Has this limitation affected the performance of FIOC-WM in your experiments?

R8: In our experiments, we did not observe significant degradation in slot extraction quality with longer video lengths. This is likely because the environments we consider are not as semantically complex or visually diverse as the long-form natural videos (e.g., YouTube-VIS) used in the original Videosaur paper. Most robotic control and embodied AI environments involve more structured and repetitive dynamics. That said, we acknowledge this known limitation of Videosaur and will add a note in Appendix E.1 stating: “While slot extraction quality may degrade over very long or highly diverse videos, we did not observe such issues in the structured environments used in our experiments.”

---

Thank you again for your positive feedback and insightful questions! Please let us know if you have any further questions!

Thank you so much for your response and your positive feedback! We are glad that we were able to address your concerns, and we truly appreciate your recognition of the paper’s contribution to the object-centric learning community.

We believe this work can also provide useful insights for broader areas such as (object-centric) representation learning, world models, and RL/control. Also, as mentioned, we will make sure to incorporate all the promised revisions in the camera-ready version.

Thank you for your careful review. We appreciate the detailed comments and your recognition of our model. Most of the concerns are centered around technical details, and we provide clarifications and additional evaluation results for each point below. We hope these responses address your concerns. For points requiring clarification, we outline concrete plans for how they will be incorporated into the camera-ready version. The additional evaluations further support the soundness of our technical design. Please let us know if you have any remaining questions or concerns.

---

No qualitative visualizations comparing long-term predictions between FIOC-WM and baseline methods.

R1: We agree that qualitative visualizations of pixel observations would be useful. Since in the rebuttal we cannot add figures anymore, we will not be able to provide them here, but in the camera-ready version, we will include side-by-side image rollouts in the appendix to compare long-term predictions made by FIOC-WM and baseline methods.

---

Role of [ $u_t$] and clarification of learning process

R2: The $\mathbf{u}_t$ variable represents the latent parameters of the distribution from which the interaction graph is sampled at each timestep, which allows us to model the fact that the interaction graphs (and their distributions) are time-varying.

Specifically, for each object pair $(i, j)$ at time $t$, we encode their latent states $s_t^i$ and $s_t^j$ using an RNN-based encoder to obtain a pairwise embedding:

$$\mathbf{u}_t^{ij} = f\_{\text{enc}, \phi_u}(s_t^i, s_t^j)$$

(as described in Equation (A2)).

We then use these embeddings to infer the graph structure $G_t$ in one of two ways:

1. Sampling edges using a Gumbel-Softmax distribution over $\mathbf{u}_t$ (Equation A3), or

2. Quantizing $\mathbf{u}_t$ into a discrete codebook of interaction types (Equation A4), followed by decoding into a graph via:$G_t \sim g\_{\text{dec}}(e_z)$, where $e_z$ is the prototype vector. (as shown in Equation (A5)). In the camera-ready version, we will put the above explanation into the main paper as we have one extra page available.

---

prior distribution P(G) over graph structures defined? What learning procedure is used to infer the graph structures?

R3: The prior distribution $P(G)$ is defined as an independent Bernoulli distribution over each edge: $G_{ij} \sim \text{Bernoulli}(p),$ where $p \in (0, 1)$ controls the sparsity level. We set $p \ll 0.5$ (e.g., $p = 0.1$) to encourage sparse and interpretable interaction graphs by biasing toward no-edge configurations.

In the camera-ready, we will clarify this in Section 3.1 (another paragraph after Line 151) by explicitly describing the sparsity prior and summarizing the graph inference procedure from state encoding to discrete graph generation.

---

The motivation for adopting two specific approaches—(i) variational mask learning and (ii) conditional independence testing—for learning the state transition distribution is not sufficiently justified.

R4: We adopt variational masking and CIT because both are well-established structure learning strategies. Variational methods benefit from scalability and can be efficiently trained via likelihood matching, while CIT offers statistical grounding for inferring dependency structures. We also note that, due to the flexibility of our framework, it can accommodate any off-the-shelf structure learning method. To further support our design choice, we additionally evaluated Dynotears [1], a strong baseline for causal structure learning [2]. We use it on the Sprites-World environment with 3, 5, and 9 objects. The results are shown below:

The results show that DynoTears performs similarly to CIT, but worse than variational methods across all object counts. This supports our use of variational masking and CIT as effective and principled choices for object-centric interaction modeling. To clarify our choice, we will add more results for DynoTears for other settings in the Appendix of the revised paper.

[1] Pamfil, Roxana, et al. "Dynotears: Structure learning from time-series data."AISTATS 2020.

[2] Vowels, Matthew J., Necati Cihan Camgoz, and Richard Bowden. "D’ya like dags? a survey on structure learning and causal discovery." ACM Computing Surveys 55.4 (2022): 1-36.

---

Pre-trained models for Dreamer and TD-MPC

R5: Thank you for the suggestion. We originally did not include this comparison because Dreamer-V3 is designed to reconstruct pixel-level observations, while TD-MPC2 predicts actions, terminal flags, and rewards directly from raw pixels. However, to better isolate the contribution of our proposed factorized world model from the benefits of pre-trained visual features, we implemented a fairer comparison.

Specifically, we adapted both baselines to use DINO embeddings as input:

• For Dreamer-V3, we replaced the pixel encoder with a frozen DINO encoder and used DINO-WM's decoder to reconstruct the observations.
• For TD-MPC2, we used DINO features to predict actions, terminal values, and rewards via its original decoding heads.

The updated results are shown below:

Using pre-trained DINO features improves both Dreamer-V3 and TD-MPC2 performance in some cases, but they still underperform compared to FIOC. This suggests that while strong visual representations help, our proposed factorized world model and hierarchical policy structure are key contributors to the performance gain.

---

Is the same offline training setup used for baseline models such as DreamerV3 and TD-MPC2?

R6: Yes, all are trained on the same dataset as we mentioned in Table A6.

---

You mention using both MPC and PPO for training the low-level policy. Under what conditions is each method used?

R7: We use variational masks for the results reported in Table 2. Specifically, for all methods (excluding those evaluated under the nSHD metric), we adopt variational masks to infer the interaction structures. For downstream control, we apply MPC for the Gym-Fetch and Franka-Kitchen environments, and use PPO for LIBERO and iGibson tasks. We also conducted an ablation study to evaluate the impact of using MPC or PPO on all environments. The results of this ablation are shown below. The results indicate that using either PPO or MPC is feasible with similar final success rates. We will add this Table in Appendix D.

---

When PPO is applied, is it trained using imagined rollouts from the learned world model (i.e., model-based PPO)?

R8: Yes, we used both the imagined rollouts and the rollouts from the replay buffer to learn the PPO policy.

---

How is the action sampling process implemented when using MPC with the Cross Entropy Method (CEM)? What are the hyperparameters (e.g., number of gradient descent steps, population size)?

R9: We implement MPC using the Cross Entropy Method (CEM) with the following hyperparameters:

• CEM iterations: 5 for Fetch, Franka, and i-Gibson; 8 for Libero
• Population size: 500 for Fetch, 1000 for others
• Planning horizon: 10 for Fetch, 20 for Franka, i-Gibson, and Libero We will add these in the camera-ready Appendix E.2 as an individual table and include an algorithmic pseudocode as well in the appendix.

---

Eq6 defines the reward prediction loss, but it is unclear how this learned reward model is utilized during policy training or planning. Is the predicted reward used in MPC rollout evaluation or in PPO optimization?

R10: The reward prediction model is used during Stage 1 for learning the world model. Specifically, we supervise it using ground-truth rewards from the offline dataset, following the same setup as in Dreamer-v3 and TD-MPC2.

The model assumes a fixed number of object slots. It is unclear how it generalizes to environments where the number of objects changes, or where previously unseen objects appear.

R11: To assess the model's ability to generalize to a greater number of objects, we train on environments with 3 objects and evaluate on tasks involving 6 and 8 objects with the fixed world model. As shown in the table below, FIOC achieves strong generalization performance, comparable to or better than baselines like EIT, and outperforms Dreamer-V3 and TD-MPC2 under distribution shifts in object count.

---

We hope these can address the concerns. Please let us know if you have any further questions. Thank you again for your time and efforts!

Thank you for your thoughtful review and feedback. We sincerely appreciate your suggestions and your recognition of our methodology and performance.

Most of the concerns relate to writing clarity and a few technical details. For clarity, we propose concrete and executable revisions that can be addressed through targeted clarifications in the camera-ready version. For the technical concerns, we provide essential clarifications and additional evaluations. We hope our responses adequately address your questions, and please let us know if you have any further concerns.

---

what G means exactly and how it's learned

R1: We define $\mathcal{G} =$ { $G_1, \ldots, G_T$ } as the sequence of interaction graphs that represent the relational structure among objects over time (Lines 72–73 and 137–138). Each $G_t$ captures the pairwise interactions between objects at time step $t$, where each edge indicates whether an interaction exists between a pair of objects. Concretely, this can be represented as a binary adjacency matrix of size $N \times N$, where $N$ is the number of objects in the environment.

To learn $G_t$, we introduce a surrogate latent variable $\mathbf{u}_t$ that parameterizes the distribution over interaction graphs. This modeling choice allows us to capture the fact that interactions and their underlying distributions may vary over time.

Specifically, for each object pair $(i, j)$ at time $t$, we encode their latent states $s_t^i$ and $s_t^j$ using an RNN-based encoder to obtain a pairwise embedding:

$$\mathbf{u}_t^{ij} = f\_{\text{enc}, \phi_u}(s_t^i, s_t^j)$$

(as described in Equation (A2)).

We then use these embeddings to infer the graph structure $G_t$ in one of two ways:

1. Sampling edges using a Gumbel-Softmax distribution over $\mathbf{u}_t$ (Equation A3), or

2. Quantizing $\mathbf{u}_t$ into a discrete codebook of interaction types (Equation A4), followed by decoding into a graph via:$G_t \sim g\_{\text{dec}}(e_z)$, where $e_z$ is the prototype vector. (as shown in Equation (A5)).

These mechanisms are designed to flexibly infer time-varying, discrete graph structures from continuous object states. More details can be found in Lines 146–151 and Appendix C.2. [For clarity] In the camera-ready version, we will move the above explanation into the main paper and expand it with additional details from Appendix C.2, as we have one extra page available.

---

Typos and color suggestion for Fig.6.

R2: [For clarity] We will correct all identified typos in the camera-ready version. For Fig. 6, we will use more distinguishable colors. For setrence at 112, we will make it as “We encode the observations using object-centric representation learning built on top of pre-trained models such as DINO-v2 [23] and R3M [24], which have been empirically shown to provide high-quality image understanding capabilities [26–28, 18, 29] and facilitate robotic manipulation tasks [19].”

---

How do you get the ground truth interaction structures G to get results shown Figure 6? Did you only do that for SpriteWorld, and that ground truth is not required unless you want to measure SHD? Also, how exactly is G represented?

R3: You are correct that ground-truth sequence of interaction graphs $G$ are only used in the SpriteWorld environment for evaluation purposes to compute the nSHD and we do not use it during training. During world model learning, we infer the estimated sequence of interaction graphs $\hat{G}$ in a self-supervised manner by maximizing the overall likelihood and optimizing the loss functions described in Section 3.1.

As for obtaining the ground-truth $G$ in SpriteWorld, we derive it from visual inspection of the collected videos: if two objects are sufficiently close and exhibit abrupt changes in their trajectories, we label them as interacting. [For clarity] We will clarify this process and add the above explanation after Line 268 in the revised version.

For representation, $G$ is a discrete binary adjacency matrix of size $N \times$, where $N$ is the number of objects. While the model may produce a continuous-valued matrix during optimization, we threshold it to obtain a binary matrix for evaluation. [For clarity] We will add this sentence after Line 73.

---

In the online stage, do you also update the world model with the new data you collect?

R4: For computational efficiency, and because our experiments show that the offline-trained world model is already robust given high-quality offline data, we do not update the world model during the online stage in our default setting. Please note that this is just an empirical choice, and it is feasible to incorporate online updates. To further investigate this, we conducted an ablation study comparing performance with and without online world model updates. The results are shown below:

Results averaged over 5 seeds indicate that there is no significant difference, though we observe slight improvements in most environments for both single-task and generalization settings. [For clarity] We will include this comparison as an additional ablation study in the appendix, and clarify the experimental setting with the following description: “We also consider the case where online interaction data is used to further update the world model components during policy learning or planning. The results are given in Appendix D.” Additionally, we will emphasize in the main paper that “we fix the learned world model during policy learning and planning”, and will add this clarification after Line 150 to avoid confusion about whether the model is updated online by default.

---

Do you train the low-level policy fully inside the learned world model or using online observations?

R5: Thank you for the question. We train the low-level policy during the offline training stage, and then fine-tune it during the online learning stage using the newly collected interaction data. [For clarity] We will explicitly state at Line 167: “We learn the low-level policy during world model learning (Stage 1), and then fine-tune it with online data during Stage 2, where the policy is updated each time new interaction data becomes available.”

We will provide the algorithm in the appendix (here is a sketch):

Offline World Model and Online Policy Learning Framework

---

Stage 1: Offline Model Learning
For each $(o_t, a_t, r_t) \in \mathcal{D}$:

• Encode observation: $\hat{o}_t \leftarrow p\_\text{pre-trained}(o_t)$
• Infer latent state: $s_t^i \leftarrow q_\phi(s_t^i \mid \hat{o}_t)$
• Factorize latent state: $s_t^i = (d_t^i, c_t^i)$
• Learn reward model: $p_r(r_t \mid s_t, a_t)$
• Learn transition model: $p_s(s_{t+1} \mid s_t, a_t, G_t)$
• Infer interaction graphs: $p_u(G_t \mid s_t)$
• Train interaction policy: $\pi^l(a_t \mid s_t, G_t^g)$

---

Stage 2: Online Policy Learning
For each environment rollout episode:

• Receive current state $s_t$ from world model encoder with input environmental steps
• Select goal graph: $G_t^g \sim \pi^h(G_t^g \mid s_t)$
• Sample action from interaction policy: $a_t \sim \pi^l(a_t \mid s_t, G_t^g)$
• Execute $a_t$ in environment, observe $r_t$, $o_{t+1}$
• Update policies $\pi^h$, $\pi^l$ using collected data

---

In line 173, you say you only focus on subsets of objects with size of less than 2. Is that correct? Do you only focus on sets of single objects?

R6: Yes, at each step, we focus on a small subset of objects, typically one or two, to induce diverse interactions for them. This does not mean we only consider single-object interactions overall. Instead, we adopt an iterative strategy, where over multiple steps, we cover all objects by sequentially targeting different subsets. This is an empirical design choice that balances interaction diversity and sample efficiency.

---

FIOC considers two options for low-level policy: MPC and RL, and two options for graph inference: Variational and CIT. What version do you use in Table 2 results?

R7: We use variational masks for the results reported in Table 2. Specifically, for all methods (excluding those evaluated under nSHD), we adopt variational masks to infer the interaction structures. For downstream control, we apply MPC for the Gym-Fetch and Franka-Kitchen environments, and use PPO for LIBERO and iGibson tasks. We also conducted an ablation study to evaluate the impact of using MPC or PPO on all environments. The results of this ablation are shown below. The results indicate that using either PPO or MPC is feasible with similar final success rates. [For clarity] We will add this Table in Appendix D.

---

We hope these clarifications address your concerns. If so, we would be truly grateful if you would consider adjusting your rating accordingly. Thank you again for your time and efforts!

Also, is there any way to look into what the slots and interactions represent, or is it purely a black-box?

Yes, in the camera-ready version, we will include a figure visualizing the reconstructed slots. (Due to file upload constraints, we cannot include the image here, but it has already been prepared and will be added.) For interactions, we can find the list of object pairs $(i, j)$ that are activated at each high-level step. This will clearly indicate which objects are interacting or being focused on during each subgoal transition.

Got it, please clearly state what version of your method you're using in each experiment in the paper.

We will insert the following clarification after Line 233 in the camera-ready version: "We use variational masks for the results in Table 2. Specifically, for all methods (excluding those evaluated under nSHD), we adopt variational masks to infer the interaction structures. For downstream control, we apply MPC for the Gym-Fetch and Franka-Kitchen environments, and use PPO for LIBERO and iGibson tasks." We will put further ablations (given in the rebuttal) in Appendix D.

---

We would like to once again express our sincere gratitude to the reviewer for the thoughtful review. Your suggestions have helped strengthen the paper. We hope the above responses have addressed your further questions, and please do not hesitate to reach out if there are any remaining concerns.

Thank you very much for your feedback and follow-up questions. We address them point by point below.

What is the dimension of $u$? Is it $u \in \mathbb{R}^{N \times N}$? Please add this to the paper if it's not there, that would help in understanding the method. And same for other variables you introduce. Having a clear understanding of variable dimensions and types would help in understanding what represents what.

For the two cases, we used different ways to sample $G_t$:

1. Sampling edges using a Gumbel-Softmax distribution over $\mathbf{u}_t \in \mathbb{R}^{N \times N}$ (Equation A3), which models the interaction logits between all object pairs. Here, the dimensionality of $\mathbf{u}_t$ is $N \times N$

2. Quantizing $\mathbf{u}\_t$ into a discrete codebook of interaction types (Equation A4), followed by decoding into a graph via:$G_t \sim g_{\text{dec}}(e_z),$ where $e_z$ is the selected prototype vector. In our implementation, we set $\mathbf{u}_t$ equal to the quantized embedding. The 16-dimensional vector $e_z$ is then decoded into the parameters of Bernoulli distributions, from which we sample the final interaction graph $G_t$. Hence, here, the dimensionality of $\mathbf{u}_t$ is the same as $e_z$: 16.

We will incorporate the dimensionality details, along with the description provided in our response to the first question in the rebuttal, into the main paper (after Line 151), while keeping the pointer to Appendix C.2 for further reference. To make it clear, we will also add a table for notations in Appendix C.

This line is still confusing to me, less than 2 means 1, so the phrase itself is confusing. And I'm still confused as to what the high-level model is doing. My current understanding is that for a given state $S_t$ ,there's some interaction graph you can extract from it $G^t$. Then, the high-level policy $\pi_h$ should select the target interaction graph you want to reach next $G^g$. If we consider the Franka Kitchen environment, and let's say we're trying to move the kettle from the counter to the stove top, the sequence of intermediate subgoal graphs we'd consider is 1) graph where the arm and the kettle are interacting (arm is holding the kettle) and then 2) graph where kettle and the stovetop are interacting. And by saying that you focus on a subset of objects with size of less than 2, you're saying that the delta between $G^t$ and $G^g$ should not be too big, meaning only one interaction between objects i, j, $G^g_{i, j} \neq G^t_{i,j}$, with the rest of the graph remaining unchanged.

Yes, we should fix the main paper "less than 2" to "1 or 2". Thank for the pointer.

For the rest of understanding, yes, you are correct. However, one clarification we would like to make is regarding the number of objects mentioned. For each subgoal interaction (i.e., the pair $(i, j)$), we first select one or two objects as candidates for $i$, and then sample $j$ conditioned on the chosen $i$. This hierarchical selection is designed for efficiency, as directly sampling from all $N \times N$ possible interactions would be computationally expensive.

To clarify, we will add the following after Line 176 in camera-ready (we really like your example so we will also use it here, many thanks!!):

---

To clarify the high-level planning process, we follow the perspective that each state $\mathbf{s}\_t$ is associated with an interaction graph $G^t$, and the final task corresponds to reaching a desired target graph $G^g$. $\pi_h$ selects a sequence of intermediate subgoal graphs that gradually transform $G^t$ into $G^g$ ($G^g_{i,j} \neq G^t_{i,j}$), where each subgoal graph differs from the previous one in only a single interaction. For example, in a task such as moving a kettle from the counter to the stove, the graph transitions might involve first enabling an interaction between the arm and the kettle, followed by an interaction between the kettle and the stovetop.

To make the subgoal selection both tractable and structured, we do not sample directly from the full space of possible object interactions. Instead, at each decision point, we first identify a small subset of objects (typically 1–2) as primary candidates for initiating interaction changes. These candidates define the anchor object(s) $i$, and we then select a target object $j$ conditioned on $i$ to form the proposed subgoal interaction $(i, j)$. This scheme reduces the combinatorial action space and leads to more localized graph transitions. Note that the selected subset does not constrain the interaction to only occur between these objects; rather, it defines a focused region of the graph for subgoal exploration.

---