COMBO: Compositional World Models for Embodied Multi-Agent Cooperation

We sincerely thank the reviewer for the time to read our paper, and we appreciate your positive and constructive comments! We address your questions in detail below and updated our paper according to your suggestions.

Q1: In Table 1, COMBO outperformed previous works. Although, when comparing with LLaVA, it shows comparable success rate except TDW-Cook Cooperator 1 setting. Their efficiency on solving the tasks is clearly better than LLaVA, but we felt it is not good enough to evaluate the effectiveness of their modeling.

Thanks for pointing this out! We agree that the current success rate metric does not provide enough detailed information for comparing the baselines. To address this, we have introduced an additional metric success rate@30, which evaluates the success rate with a shorter planning horizon of $h=30$. The results, shown in the table below, highlight that our method COMBO clearly outperforms other baselines, including LLaVA, achieving an average success rate of 0.98 compared to 0.5 for the LLaVA baseline.

• The higher success rate of LLaVA with a longer planning horizon can be attributed to the randomness in its predictions, which occasionally allows it to escape from dead loops when given sufficient steps.

Q2: The generalization performance evaluation is too weak to show that in lines 521-524 and Table 5. They trained the model with 4 agents and applying it to 3 agent only, then what happens if we test with 2 or much more agents such as 10? Additionally, COMBO is more efficient than LLaVA, but LLaVA also showed good success rates for 3 agents setting. The LLaVA is trained on the 3 agents setting?

Thanks for your valuable question. Following your suggestions, We additionally tested our method and the LLaVA baseline on the 2-agent version of TDW-Game without further training. The success rate@30 and average steps when cooperating with two different cooperators over 10 episodes are reported in the table below. The results demonstrate that our method COMBO clearly outperforms LLaVA across all settings of different numbers of agents, successfully completing all tasks within 30 steps without requiring additional training. These results have also been incorporated into Table 5 of the manuscript.

• Although the LLaVA baseline was also not further trained, it benefits from being pretrained on internet-scale data and only needs to propose an action for itself, which makes it inherently easier to transfer to new scenarios. On the other hand, COMBO includes a world model trained entirely from scratch using only the 4-agent TDW-Game trajectories. It is therefore remarkable that COMBO demonstrates strong generalization to the 3-player TDW-Game, despite the significantly different agent dynamics in this setting, due to its compositional structure.

Q3: In Figure 1, what means the surrounding figures? Maybe we guess it is to visualize the look ahead planning, but I am not sure it is a good visualization for showing that.

We're sorry for the confusion. Yes, it's meant to visualize the look ahead planning. We have updated a better version with fewer frames and more clear annotations to improve the interpretability.

We sincerely appreciate your constructive comments. Please feel free to let us know if you have further questions.

Best,

Authors

First of all, thank authors for preparing concrete rebuttal with additional experiments in the limited time.

Their rebuttals addressed my concerns well, but I have one more question.

• Could you explain why Success Rate@30 is good measurement for checking the performance? It looks showing the efficiency of the plannings.

We sincerely thank the reviewer for the time to read our paper, and we appreciate your positive and constructive comments! We address your questions in detail below and updated our paper according to your suggestions.

Q1: Scaling loss with the reachability assumes that reachability is provided externally, in real world agents "reachability" should be additionally estimated/discovered from the exploration data. It would be great if the authors would cover better how to discover the reachability regions if they are not provided.

Thanks for the valuable question!

• In our specific setting, reachability is defined by the task rules, where agents are restricted to accessing only one side of the table and must collaborate closely to complete the task. Therefore, we utilize the agents' reachable regions to define the coefficient matrix.
• That said, this is not the only approach to setting the loss coefficient matrix based on reachability. When such information is unavailable, advanced object detection and segmentation models like GroundingDINO [1] and SAM2 [2] can be employed to preprocess the data to label the semantic regions and set the coefficient matrix accordingly.

Q2: what about regions that are not reachable by any agent? In the current formulation, it is not clear if those regions are modeled or not in the world model.

Thanks for bringing up this issue!

• In the first stage of training the CWM, agent-dependent loss scaling is employed to better capture $P(X|x_0,a_i)$. As detailed in Appendix A.3, the coefficient matrix $C$ is set to 2 for reachable regions and 1 for other regions. This encourages the model to focus more on relevant pixels while avoiding the complete exclusion of unreachable regions.
• In the second stage, during compositional training, all regions in the image are treated equally. This includes regions that are unreachable by any agent, which are also explicitly modeled in the CWM.

Q3: Fine-tuning of VLM on the "collected short rollouts" potentially lead to agent that memories how other agents behave. While this approach can work, it would always require fine-tuning "intend prediction" if the agent is changing policy. It would be great if the authors can show that VLM can use some of the initial observations of the agent behaviors to adapt its predictions on the fly (e.g. with in context learning).

Thank you for the thoughtful feedback! We agree that it is desirable for the intent tracker module to generalize to unseen cooperators and adapt its predictions dynamically based on historical observations.

• Our approach employs a shared intent tracker module across all tasks and cooperators. The test-time cooperator's policy is agnostic to our agent during training, meaning the intent tracker must rely on observation history to adapt its predictions and track the intents of diverse cooperators effectively. As demonstrated in Table 1, our agent successfully collaborates with cooperators following different policies.
• To address this point more explicitly, we add a new qualitative result in Appendix C.2 with Figure 13 which illustrates how the intent tracker module adapts its predictions based on historical observations. For example, given the same current state, the intent tracker accurately predicts the next action of Agent Bob by leveraging its observation history, which reveals a preference for preparing his recipe first or helping Agent Alice first. With accurate intent tracking, it becomes possible for Agent Alice to select the best action to take. This adaptation showcases the model's ability to generalize and predict dynamically without relying solely on memorizing.

Q4: Please describe in more details on what data intend and outcome evaluation was fine-tune on? How realistic collection of such data for real-world agents?

Thanks for bringing up this issue! We provide more details below and also added them in Appendix C.1.

• We generated 4k training environments by randomly sampling a recipe in TDW-Cook and a puzzle box in TDW-Game for each agent. All objects were initialized at random positions on the table. Scripted planners with randomness were then used to play out the episodes and collect the state and action histories.
• For the Intent Tracker, we collected 40k short rollouts consisting of three images of consecutive observations and a textual description of the next actions of all the agents, converted by a template given the action history.
• For the Outcome Evaluator, we collected 138k data consisting of one image of the observation and a textual description of the state of each object in the image and the heuristic score, converted by a template given each object's location.
• This dataset collection pipeline is transferable to real-world robotics tasks. For instance, in practical scenarios, robots can follow a similar approach by utilizing predefined task rules and leveraging planners to simulate and annotate task executions in diverse settings.

Thank you for your continued valuable feedback and insightful suggestions to strengthen our work. We are glad you found our responses and updates useful. Based on your suggestions, we have incorporated computational requirements into Appendix C and provide further clarifications below to address your remaining concerns.

Reachability and Agent-Dependent Loss Scaling

Thank you for raising the point regarding the assumption of reachability. We provide additional clarification and discuss potential approaches:

• Reachability is not strictly required to set the loss coefficient matrix: The key idea of this technique is to encourage the model to focus more on the pixels related to the action prompt. There are also other possible ways to leverage this technique. One promising way is to use advanced vision models to identify the action prompt-relevant regions in an image (e.g., "the lettuce" in the action prompt "Bob pick up the lettuce") and adjust the loss scale for those regions.
• Alternative approaches: We agree with your suggestion of discovering reachability during exploration. One promising avenue is bootstrapping from zero-shot vision-language models (VLMs), where an agent could learn reachability by interacting with the environment.
• Effectiveness without agent-dependent loss scaling: Even without this agent-dependent loss scaling, our proposed method still significantly improves the previous methods by leveraging the composable generation. As shown in Table 3, we observe improvements in success rates from 20%–25% to 55%-70% in the challenging setting of simulating the effects of multiple sets of actions on the world state.
• We've also incorporated this discussion into Appendix A.3 in the manuscript.

Adaptability of the Intent Tracker to Unseen Agent Behaviors

We appreciate your clarification on the question about adapting to agents with previously unseen behaviors. Below, we outline how our method addresses this challenge:

• Intent tracker adaptability: Our intent tracker leverages observation history to adapt its predictions dynamically, even for cooperator behaviors not encountered during training. While the tracker is trained on a limited set of scripted planners (e.g., the 4-agent TDW-Game setting), it generalizes reasonably well to new behaviors due to its reliance on cooperator-independent observation patterns.
• Generalization across settings: As demonstrated in Table 5, the same intent tracker was applied without retraining in 3-agent and 2-agent versions of TDW-Game, where cooperators followed policies unseen during the training phase. The intent tracker module demonstrated adaptability, achieving correct intent prediction rates of 75.8% in the 2-agent cooperation scenario and 64.6% in the 3-agent cooperation scenario. Among the wrong predictions, 20% attributes to the intent prediction of "wait" while it shouldn't be. These results underscore the intent tracker module’s ability to adjust to new behaviors.
• Future improvements: We acknowledge that further generalization to more diverse and complex cooperator behaviors remains a challenge and is a promising avenue for future research. Enhancements such as Population-Based Training [1] and Behavior Diversity Play [2] could increase the diversity of cooperators during training, thus improving robustness. Additionally, incorporating more sophisticated modeling techniques, such as the Bayesian Theory of Mind [3], could enable better inference and adaptability in multi-agent collaborations.
• We've also incorporated this new result and discussion into Appendix C.2 in the manuscript.

[1] Human-level performance in 3D multiplayer games with population-based reinforcement learning. Science,2019

[2] Adaptive Coordination in Social Embodied Rearrangement. ICML,2023

[3] Too Many Cooks: Bayesian Inference for Coordinating Multi-Agent Collaboration. Topics in Cognitive Science,2021

We sincerely appreciate your insightful comments and detailed suggestions. Please feel free to let us know if you have further questions or require additional clarification.

Best,

Authors

We appreciate the positive and constructive comments from you! We address your questions in detail below and have updated our paper according to your suggestions.

Q1: Lack of Figure Clarity and Interpretability.

Thanks for the detailed and actionable suggestions to improve our work. We've reorganized Figure 1-(b) with fewer images and added explicit labels, and Figure 4 with added annotations to improve the interpretability of the figures.

Q2: Despite its impressive performance, the proposed model may struggle to scale in more realistic scenarios, such as handling low-level controls with continuous action spaces.

Thanks for bringing up this issue!

• Our compositional framework is designed to improve multi-agent task planning capabilities and can easily integrate various low-level control policies, including reinforcement learning and trajectory optimization for robotics tasks.
• We add a new real-world experiment following your suggestion, and the results show our method is capable in real-world scenarios involving low-level control and non-trivial generalization. Two figures showing our task setup and results and more details can be found in Appendix B. We also upload the demo video in the updated supplementary materials.
  • We instantiate the TDW-Game benchmark in the real world with two agents. A human and an XArm sit on the upper and right sides of the table, each can only reach the objects on his side of the table and must cooperate to pass and place 3-5 puzzle pieces scattered randomly on the table into the correct puzzle box according to visual clues such as the shape. In each step, the human or the robot can only pick or place one object within his reachable region. With 5 trials of experiments, our method succeeded in 4 trials with an average step of 8.4.

Q3: I wonder if the CWM training setup will be applicable in realistic scenarios. In my opinion, we might only have access to egocentric observations and action history during training. I believe that those compositional world modeling abilities seem reliant on structured inputs through prompts, which aren't typically available in realistic multi-agent cooperation setups.

Thanks for the valuable question.

• Our environments are built on top of ThreeDWrold[1], which is a physically realistic simulation platform. Therefore our training pipeline could also be employed in the real world.
• The only data we need to collect to train the CWM are ego-centric observations, action history, and orthographic views, which could all be collected in the real world with RGBD cameras.
• The input to CWM is just converting the actions of agents into text through a template, which is easy to get in even real world.

Q4: Does the action prompt to CWM use a cropped image token along with text, as depicted in Figure 4, or is that just an illustration?

We're sorry for the confusion.

• The cropped image in Figure 4 is just an illustration for better interpretability.
• The action prompt to CWM is pure text describing the joint actions, such as "Alice pick up the lettuce. Bob pick up the pickle slice."

[1] ThreeDWorld: A platform for interactive multi-modal physical simulation. NeurIPS21

We sincerely appreciate your detailed and constructive comments. Please feel free to let us know if you have further questions.

Best,

Authors

• Following reviewer Dcsa's suggestion, we additionally tested our method and the LLaVA baseline on the 2-agent version of TDW-Game without further training. The success rate@30 and average steps when cooperating with two different cooperators over 10 episodes are reported in the table below. The results demonstrate that our method COMBO clearly outperforms LLaVA across all settings of different numbers of agents, completing all tasks within 30 steps without requiring additional training. These results have also been incorporated into Table 5 of the manuscript.

• We add a new qualitative result in Appendix C.2 with a Figure.13 which illustrates how the intent tracker module adapts its predictions based on historical observations as Reviewer Rt2j suggested.
• We add an additional section in Appendix C.3 to discuss the alternative implementation choices for the planning submodules and conduct experiments to show VLMs not only offer better generalizability but also achieve superior performance compared to a simpler supervised model as Reviewer Rt2j suggested. This is also reported in Table 7 in the manuscript.

• We add failure case analysis in Appendix A.5, providing Figure 9 and Figure 10 to illustrate common errors such as Image Quality and Misunderstanding and discuss the future extension with physics-aware video generation approaches as Reviewer Rt2j suggested.
• We update the Figure 1(b) and Figure 4 for better interpretability as Reviewer yTaz and Dcsa suggested.
• We add more details about the 2d-fetchQ challenge in Appendix A.1.2 as Reviewer Rt2j suggested.
• We add the details of training data collection for the planning submodules in Appendix C.1 as Reviewer Rt2j suggested.

We hope our detailed responses below convincingly address all reviewers’ questions.