Language-Guided Object-Centric World Models for Predictive Control

Thanks for your valuable feedback. We are currently running additional experiments to address reviewers’ concerns, and below are our responses to the specific concerns if answerable without additional experimental results. We will fully cover all concerns during the extended review timeline, thanks for your understanding!

[Weakness 1] Despite the paper being relatively clearly written, I would highly recommend using ICLR 2025’s extended page limit to increase the size and quality of your visualizations. For instance, for Figure 3, it’s very hard to see where the cube and moon are in the scene. I likewise cannot see any empirical quality differences between your approach’s generations and SuSiE’s.

Thanks for pointing out the difficulty in parsing the qualitative results in the figures and visualizations, we improved figure visibility to make the results clearer and more interpretable. To aid understanding, we have added a video related to Figure 3 in the supplementary material. The video includes four segments: original.gif (the original video), ours.gif (the video decoded by our method), and seer_s.gif and seer_f.gif (videos generated by the Seer baselines). While Seer is restricted to generating 10 frames, our method, being autoregressive, has no such limitation. Additionally, as per your suggestion, we have increased the size of Figures 3 and 4 to improve readability.

However, both Seer and our method appear to lose shape details when decoding into the pixel space in Figure 3. In our case, this is likely due to the limitations of the image decoder, as well as the fixed resolution, which constrains the potential for quality improvement. We believe this limitation is acceptable, as the primary contribution of our method is not in producing high-quality frame decodings but in accurately predicting future states within the latent space.

[Weakness 2] The approach is not tested on a wide range of tasks – only the simulated LanguageTable benchmark. It’s not at all clear to me that it would generalize to the real world. Given that SuSiE was evaluated both in sim (CALVIN) and on real-world robots on Bridge-like tasks and showed good performance (compared to strong real-world baselines like RT-2), it is unclear if the present paper’s approach would similarly scale to such more complex tasks.

To show our method in a more challenging environment, we conducted additional experiments on more complex simulated environments (Libero) using our method with SAVi and with state-of-the-art video slot encoder model (SOLV). Please refer to the general comment for additional experiments.

[Weakness 3] Similarly, the authors claim: However, the major drawback of language-guided video-generation models is the requirement of large-scale labeled language-video datasets and the corresponding high computational cost. Therefore, latent predictive models, which abstract video to predict forward in compact latent state spaces, can serve as an alternative from a computational efficiency perspective.

• If this is true, it seems more sensible to evaluate in the real world, where data is more limited than in sim.
• Additionally, SuSiE does show that an off-the-shelf image generation model pre-trained on general image-language data can be fine-tuned to work well with robot data just on existing relatively limited robot datasets. If that’s the case, it seems highly unclear that sample efficiency is a problem.

We currently have no access to real world demonstration data and physical manipulation devices. Also, complex domain cases like real-world are limited due to our default setting’s (SAVi with simple CNN) expressiveness capability, which is addressed in additional experiments. Extending and testing in real-world scenarios is a priority for future work once the necessary resources become available.

We considered internet scaled pre-training also as a part of the dataset and computation burden, which can be resolved leveraging an open-sourced model. However, the scale of limited fine-tuning data in their work is not very trivial compared to pre-training data scale. Also, our main contribution is computational efficiency during training and inference.

Thanks for your valuable feedback. We are currently running additional experiments to address reviewers’ concerns, and below are our responses to the specific concerns if answerable without additional experimental results. We will fully cover all concerns during the extended review timeline, thanks for your understanding!

[Weakness 1] The authors do not have a SlotFormer baseline, which does not use any language conditioning. Given that one of the key claims of the paper is that language conditioned object centric world models help downstream tasks, checking the importance of being language centric is critical. Adding that baseline would be helpful.

We conducted an additional experiment for the suggested SlotFormer baseline replacing LSlotFormer to SlotFormer with identical experiment settings only without language instruction. Slotformer prediction successfully reconstructed in-domain images, but semantically fails to perform block-to-block tasks. Action decoder trained with SlotFormer successed 2/200 (1%) of environment evaluation tasks.

This is because in a language table environment, the agent does not have any information other than language instruction to perform given tasks. One can claim that even without instruction, when successfully learned, the model could be able to generate random A-to-B trajectory resulting in a success rate of 2/(n(n-1)); 2 for B-to-A execution leading to success. However, the predicted trajectories do not show a long-term task solving behavior due to its task ambiguity and inconsistency.

[Weakness 2] For the evaluation of this approach, the authors have used the language table simulation environment, which involves some objects to be manipulated on a table top setting. This makes sense since there is a clear distinction between foreground (the objects) and background, which favors object centric approaches over general video generative models. However, showcasing some other scenarios or evaluation setups where maybe lets say intuitively a video generative model would have an edge, would have been interesting and more convincing to see.

We expect video generative models, especially when pretrained, would have an edge in more complex image domains. To show our method in a more challenging environment, we conducted additional experiments on more complex simulated environments (Libero) using our method with SAVi and with state-of-the-art video slot encoder model (SOLV). Please refer to the general comment for additional experiments.

[Weakness 3] Minor: The qualitatives in figure 3 are not the easiest to parse, if the author’s method works well, but having a video to see the predictions would make the difference much clearer, I couldn’t find anything in the suppmat.

Thanks for pointing out the difficulty in parsing the qualitative results in the figures, we added video prediction visualizations in supplementary materials to make the results clearer and more interpretable.

Thanks for your valuable feedback. We are currently running additional experiments to address reviewers’ concerns, and below are our responses to the specific concerns if answerable without additional experimental results. We will fully cover all concerns during the extended review timeline, thanks for your understanding!

[Weakness 1] Overall, the proposed method is a simple modification of the SlotFormer adding language goal conditioned predictions and trained on a large dataset of the demonstrations. On it its own it is not a big problem, if the proposed methods would be studied on diverse and challenging environments and compared with other methods that are state-of-the-art world models

[Weakness 3] It is not clear how the methods compare to the standard baselines on this task: while outperforming diffusion models for video prediction, it is not clear if usage of world model with object-centric representations are comparable or not with state-of-the-art algorithms using the same data for training.

To show our method in a more challenging environment, we conducted additional experiments on more complex simulated environments (Libero) using our method with SAVi and with state-of-the-art video slot encoder model (SOLV). Please refer to the general comment for additional experiments.

For the baselines, we considered language conditioned diffusion video/image prediction models as world model baselines. Our motivation is to substitute diffusion world models with more efficient object-centric world models. Although direct comparison between these methods can be difficult due to their difference in design, our approach shows better performance over diffusion models, both when trained from scratch on the target dataset and when fine-tuned from a pre-trained model, while also offering greater computational efficiency.

The primary motivation behind our research stems from studies such as UniSim (Yang et al., 2023) and UniPi (Du et al., 2023), which employ video diffusion models as world models but require substantial time and computational resources. For instance, UniPi utilizes 256 TPU-v4s to train video diffusion models for 2 million steps with a batch size of 2048. Similarly, UniSim uses 512 TPU-v3s to train video diffusion models for 1 million steps over 20 days, processing nearly 821 million data examples. Furthermore, these diffusion models require several sampling steps during inference, which can also pose significant challenges in practical applications.

To address these challenges, we aimed to implement a world model using a latent predictive approach. Therefore, we compare our method against SOTA diffusion-based generative world models as baselines. While we greatly appreciate your suggestion to compare with other types of world models, we believe that doing so could divert the focus from the primary objectives of our work.

Additionally, as shown in Table 2, we provide a comparison of the data usage between diffusion-based generative world models and our approach. Notably, we outperform Seer-S, a version of Seer trained from scratch using the same amount of data as our method. This result further underscores the efficiency and effectiveness of our proposed approach.

[Weakness 2] While the improved performance on the synthetic dataset is encouraging, it is still not clear how the method would perform in more realistic scenarios where both object-centric models and corresponding agents can struggle. As mentioned by authors, recently it was shown that object-centric methods are able to decompose much challenging images or videos (e.g. see DINOSAUR (Seitzer et al. (2023)) for images or VideoSAUR [5] /SOLV (Aydemir et al., 2024) for videos). Thus, it would be important to test how object-centric world models perform in more realistic environments with visual more complex scenarios, e.g. by training LSlotFormer on VideoSAUR or SOLV slots on environments like ManiSkill2).

To show our method in a more challenging environment, we conducted additional experiments on more complex simulated environments (Libero) using our method with SAVi and with state-of-the-art video slot encoder model (SOLV). Please refer to the general comment for additional experiments.

[Weakness 4] The analysis suggesting that future state prediction alone suffices for action decoding is questionable; the low accuracy for "instruction + 0 future steps" (2.5%) compared to near-zero performance for Seer implies that baseline results may lack rigor, potentially outperforming when future states are not predicted.

Your observation that it might be possible for baseline methods to achieve success using only the current frame and instruction, without future-predicted slots, is valid, especially since our proposed method demonstrated around a 2.5% success rate under these conditions. To test this hypothesis, we trained an action decoder (policy) using features extracted by inputting the current RGB frame into a VAE, along with the instruction. We evaluated this approach in the same language table setting described in the paper across 200 episodes. The resulting success rate was only 0.5%, supporting our claim that control tasks are challenging without future prediction, relying solely on the current state and instruction.

Your implication that the Seer results with future prediction may lack rigor is also reasonable. However, as presented in the main text and supplementary material (Figure 3, Figure 6-7), this is likely due to Seer’s consistent failure in future prediction.

[Weakness 5] The dataset used is overly simplistic, limiting the scope of validation for the world model. Testing across multiple, varied environments would better demonstrate the model’s general applicability and robustness.

To show our method in a more challenging environment, we conducted additional experiments on more complex simulated environments (Libero) using our method with SAVi and with state-of-the-art video slot encoder model (SOLV). Please refer to the general comment for additional experiments.

Finally, we hope our response could address the concerns, and we thank the reviewer again for the helpful comments. We are glad to discuss further comments and suggestions.

Thanks for your valuable feedback. We are currently running additional experiments to address reviewers’ concerns, and below are our responses to the specific concerns if answerable without additional experimental results. We will fully cover all concerns during the extended review timeline, thanks for your understanding!

[Weakness 1] The contribution in terms of "object-centric" design feels limited, as it primarily substitute the encoder for SAVi without introducing distinct object-centric innovations.

While most object-centric literatures are limited to image reconstruction and segmentation, integration to more practical downstream tasks can enrich the usability of object-centric representation. Therefore, rather than proposing a novel method of object-centric representation, we focus on properly utilizing object-centric representation which is done simply and efficiently using transformers with permutation invariance.

[Weakness 2] The lack of an experiment comparing your proposed model with a VAE-based variant (ours + VAE in Tab.1) makes it difficult to conclusively justify the benefits of slot attention.

The reason we did not explicitly include “ours + VAE” in the main text is that the primary purpose of our world model is not to decode into the RGB space but to predict states in the latent space of slots. We considered it somewhat unnatural to decode into the RGB space and then re-encode it using a VAE. However, your observation is valid, and we are currently conducting additional experiments to address this question. We will provide a response as soon as we have the results.

[Weakness 3] Comparison against video diffusion models would be more appropriate than models like instructpix2pix, as diffusion models are more aligned with the proposed model's multi-frame prediction capability.

One of the baseline models, Seer, is a language-conditioned video diffusion model conditioned on 6 frames to generate 10 future frames like our method.

The motivation behind our research stems from studies such as UniSim (Yang et al., 2023) and UniPi (Du et al., 2023), which employ video diffusion models as world models but require substantial time and computational resources. For instance, UniPi utilizes 256 TPU-v4s to train video diffusion models for 2 million steps with a batch size of 2048. Similarly, UniSim uses 512 TPU-v3s to train video diffusion models for 1 million steps over 20 days, processing nearly 821 million data examples. Furthermore, these diffusion models require several sampling steps during inference, which can also become a significant obstacle in practical applications.

To address these challenges, we aimed to implement a world model using a latent predictive approach. To evaluate the time efficiency of our method, we selected Seer (Gu et al., 2023), a latent video diffusion model known for its strength in time efficiency, as a baseline. As you pointed out, SuSiE (Black et al., 2023) may not align perfectly in terms of multi-frame prediction capability. However, if we consider its functionality solely as a world model, it can be seen as more efficient than video diffusion models like Seer due to its ability to generate only the essential frames. Additionally, because it generates frames autoregressively, it tends to produce coarser videos. From this perspective, we concluded that SuSiE, based on instruct2pix2pix, is a suitable comparison point for time efficiency as a generative world model. Its strong time efficiency aligns well with our goal of benchmarking the proposed model against generative methods in this aspect.