COMBO: Compositional World Models for Embodied Multi-Agent Cooperation

## Abstract

Abstract In this paper, we investigate the problem of embodied multi-agent cooperation, where decentralized agents must cooperate given only partial egocentric views of the world. To effectively plan in this setting, in contrast to learning world dynamics in a single-agent scenario, we must simulate world dynamics conditioned on an arbitrary number of agents’ actions given only partial egocentric visual observations of the world. To address this issue of partial observability, we first train generative models to estimate the overall world state given partial egocentric observations. To enable accurate simulation of multiple sets of actions on this world state, we then propose to learn a compositional world model for multi-agent cooperation by factorizing the naturally composable joint actions of multiple agents and compositionally generating the video.
By leveraging this compositional world model, in combination with Vision Language Models to infer the actions of other agents, we can use a tree search procedure to integrate these modules and facilitate online cooperative planning.
To evaluate the efficacy of our methods, we create two challenging embodied multi-agent long-horizon cooperation tasks using the ThreeDWorld simulator and conduct experiments with 2-4 agents. The results show our compositional world model is effective and the framework enables the embodied agents to cooperate efficiently with different agents across various tasks and an arbitrary number of agents, showing the promising future of our proposed framework. More videos can be found at https://vis-www.cs.umass.edu/combo/ .

###### Keywords:

## 1 Introduction

Building cooperative embodied agents that can engage in and help humans in multi-agent tasks is a valuable yet challenging endeavor. To cooperatively plan in a multi-agent scenario, in contrast to learning world dynamics in a single-agent scenario, there is an additional challenge to simulate world dynamics conditioned on an arbitrary number of agents’ actions given only partial egocentric observations of the world.

Large generative models have brought remarkable advances to various domains, including language understanding and generation [, , , ], image understanding and generation[, , , , , , ], and even video generation [, , , , , , , , ]. Many have explored how to leverage these powerful foundation models for embodied AI, [, , , , , ] leverage Large Language Models for decision-making, [] incorporates LLMs to build modular embodied agents for cooperation and communication, [, , ] uses Multimodal Language Models to build capable vision agents, [, ] use Video Models to generate visual plans for robots, [, , , ] explores modeling world dynamics with video models to improve single-agent planning. However, how to leverage Multimodal Language Models and Video Models to help build embodied agents capable of planning under a visual cooperation task is under-explored where it’s important to accurately simulate the world dynamics conditioned on an arbitrary number of agents’ actions given only partial egocentric observations at each step for efficient cooperation.

We propose to learn a compositional world model for multi-agent cooperation by factorizing the naturally composable joint actions of an arbitrary number of agents and compositionally generating the video to enable accurate simulation of multiple sets of actions on the world state. To address the partial observability issue, we leverage a learned generative model to estimate the overall world state given partial egocentric observations. By leveraging the compositional world model, in combination with Vision Language Models, we propose COMBO, a novel Compositional wOrld Model-based emBOdied multi-agent planning framework to facilitate online cooperative planning. COMBO first estimates the overall world state from partial egocentric observations, then utilizes Vision Language Models to act as an Action Proposer to propose possible actions, an Intent Tracker to infer other agents’ intents, and an Outcome Evaluator to evaluate the different possible outcomes. In combination with the compositional world model to simulate the effect of joint actions on the world state, COMBO uses a tree search procedure to integrate these modules and empowers embodied agents to imagine how different actions may affect the world with other agents in the long run and plan more cooperatively.

We build two challenging embodied multi-agent long-horizon cooperation tasks on ThreeDWorld []: TDW-Game and TDW-Cook where 2-4 decentralized agents must cooperate to finish several puzzles according to the visual clue or dishes according to the recipe given only partial egocentric views of the world as shown in Fig. 1, requiring extensive visual cooperation. The agents need to estimate the overall world state given partial egocentric observations, infer other agents’ intents, and accurately simulate how the world state may be affected by multiple sets of actions to make efficient cooperation in the long run. Our extensive experiments show our learned compositional world model delivers accurate video synthesis conditioned on multiple sets of actions from an arbitrary number of agents and COMBO enables the embodied agents to cooperate efficiently with different agents across various tasks and an arbitrary number of agents. In sum, our contribution includes:

- •
We propose to learn a compositional world model for multi-agent cooperation by factorizing joint actions of an arbitrary number of agents and compositionally generating the video to enable accurate simulation of multiple sets of actions on the world state.
- •
We introduce COMBO, a Compositional wOrld Model-based emBOdied multi-agent planning framework to empower the agents to imagine how different actions may affect the world with other agents in the long run and plan more cooperatively.
- •
We create two challenging embodied multi-agent planning tasks with 2-4 agents, conduct experiments, and show our framework enables the embodied agents to cooperate efficiently with different agents across various tasks and an arbitrary number of agents.

## 2 Related Work

### 2.1 Multi-Agent Planning

Multi-agent planning has a long-standing history [], with various tasks and methods have been introduced [, , , , , , , , , , , , ]. For embodied intelligence, [, ] explored the social perception of the agents during householding tasks, [] studied communication and cooperation of two agents in a multi-room house. However, these works didn’t dig into the explicit challenge of modeling the world dynamics conditioned on an arbitrary number of agents’ actions given partial egocentric observations of the world, which is essential for the agents to cooperate efficiently in long-horizon cooperation-extensive visual tasks. In contrast, we learn a compositional world model for multi-agent cooperation and propose a compositional world model-based embodied multi-agent planning framework to empower the embodied agents to imagine and plan more cooperatively with only partial egocentric observations.

### 2.2 Large Generative Models for Embodied AI

With the recent advance of large generative models, numerous works have explored how they can help build better agents [, , , ], especially in embodied environments [, , , , , , , , , ]. Specifically, [, , , , , , , , , , , , ] leverage Large Language Models to help decision-making. [, , , , ] use multimodal language models to facilitate embodied agents to plan end-to-end in visual worlds. [, ] use diffusion models for decision-making. [, , , , ] use video models to help the robot make visual plans, however, they assume the observation is static and non-partial and neglect the challenge of learning world dynamics conditioned on multiple sets of actions of an arbitrary number of agents. We learn a compositional world model for multi-agent cooperation to enable accurate simulation of multiple sets of actions on the world and leverage a learned generative model to estimate the overall world state given partial egocentric observations to address the partial observability issue.

## 3 Problem Statement

Embodied multi-agent cooperation problem can be formalized as a decentralized partially observable Markov decision process (DEC-POMDP) [, , ], defined by $(n,S,\{A_{i}\},\{O_{i}\},T,G,h)$, where:

- •
$n$ denotes the number of agents;
- •
$S$ is a finite set of states;
- •
$A_{i}$ is the action set for agent $i$;
- •
$O_{i}$ is the observation set for agent $i$, including partial egocentric visual observation the agent receives through its sensors;
- •
$T(s,a,s^{\prime})=p(s^{\prime}|s,a)$ is the joint transition model which defines the probability that after taking joint action $a\in A=A_{1}\times\cdots\times A_{n}$ in $s\in S$, the new state $s^{\prime}\in S$ is achieved;
- •
$G$ is the final goal state;
- •
$h$ is the planning horizon.

Starting from an initial state $S_{0}\in S$, $n$ decentralized agents need to plan their actions $a_{i}\in A_{i}$ given partial egocentric RGBD observations $o_{i}\in O_{i}$ each step to achieve the goal state $G$ using a minimal number of steps.

## 4 Compositional World Model

Figure: Figure 2: Compositional World Model. Given a text condition of a joint action of multiple agents $a$, we first factorize it into several components corresponding to each agent $a_{i}$, then the compositional world model generates multiple scores conditioned on the text components, then composed to generate the final output video.

To cooperatively plan in a multi-agent scenario, in contrast to learning world dynamics in a single-agent scenario, there is an additional challenge to simulate world dynamics conditioned on an arbitrary number of agents’ actions. We propose to learn a compositional world model for multi-agent cooperation as shown in Fig. 2 by factorizing the naturally composable joint actions of an arbitrary number of agents and compositionally generating the video to enable accurate simulation of multiple sets of actions on the world state. We first discuss how we learn the composable video diffusion model to act as a compositional world model in Sec. 4.1, then introduce the Agent-Dependent Loss Scaling we designed for the challenging situation where we need to accurately model multiple agents’ manipulating multiple objects to improve the video synthesis performance in Sec. 4.2.

### 4.1 Composable Video Diffusion Models

Text-guided video diffusion models learn to model the distribution of videos conditioned on text prompts $P_{\theta}(x|a)$, where the text prompts $a$ in our scenario corresponds to the joint actions of $n$ agents $a=(a_{1},\cdots,a_{n})\in A=A_{1}\times\cdots\times A_{n}$, which can be naturally factorized into $n$ components $a_{1},\cdots,a_{n}$ corresponding to the action of each agent. We propose to learn the text-guided video diffusion model as a composition of individual video diffusion models conditioned on each text component $a_{i}$:

$$ $\displaystyle P_{\theta}(x|a)=P(x|a_{1},\cdots,a_{n})\propto P_{\theta}(x) \prod_{i=1}^{n}\frac{P_{\theta}(x|a_{i})}{P_{\theta}(x)}$ (1) $$

Leveraging the interpretation that diffusion models are functionally similar to Energy-Based Models [, ], we can train the model to learn a set of score functions $\epsilon(x_{t},t|c_{i})$ and compose them as in Equation 2, which corresponds to the production of the probability densities in Equation 1.

$$ $\displaystyle\hat{\epsilon}(x_{t},t|a)=\epsilon_{\theta}(x_{t},t)+\sum_{i=1}^{ n}\epsilon_{\theta}(x_{t},t|a_{i})-\epsilon_{\theta}(x_{t},t)$ (2) $$

Then the sampling process changes to

$$ $\displaystyle x_{t-1}\sim\mathcal{N}(x_{t}-\hat{\epsilon}(x_{t},t|a),\sigma_{t }^{2}I)$ (3) $$

We train this composable video diffusion model with two stages. In stage one, we learn to model the distribution of $P_{\theta}(x|a_{i})$ by training with only conditions corresponding to a single agent’s action component using the standard denoising diffusion training objective

$$ $\displaystyle L_{MSE}=\sum_{i=1}^{n}\|\epsilon_{\theta}(x,t|a_{i})-\epsilon\|^ {2}$ (4) $$

Then in stage two, we fine-tune the model to specifically learn compositional generation to model $P_{\theta}(x|a)$ by training with conditions containing joint actions of multiple agents using the composed score function loss

$$ $\displaystyle L_{Composed}=\|\overline{\epsilon}(x_{t},t|a)-\epsilon\|^{2}= \left\|\frac{1}{n}\sum_{i=1}^{n}\epsilon_{\theta}(x_{t},t|a_{i})-\epsilon \right\|^{2}$ (5) $$

After the two-stage training, we can sample from the composed distribution with the composed score function at inference time given conditions of joint action $a=(a_{1},\cdots,a_{n})$ as

$$ $\displaystyle\hat{\epsilon}(x_{t},t|a)=\epsilon_{\theta}(x_{t},t)+\sum_{i=1}^{ n}\omega(\epsilon_{\theta}(x_{t},t|a_{i})-\epsilon_{\theta}(x_{t},t))$ (6) $$

where $\omega$ is the guidance weight controlling the temperature scaling on the conditions. When the condition corresponds to a single action component, the above sampling procedure reduces to the classifier-free guidance[].

### 4.2 Agent-Dependent Loss Scaling

The performance of stage one training impacts the final compositional generation performance a lot given the accurate modeling of $P(x|a_{i})$ is the basis for modeling $P(x|a)$. We have employed a technique named Agent-Dependent Loss Scaling to assist in stage one training.

Formally, we define an agent-dependent loss coefficient matrix $C\in\mathbb{R}^{n\times H\times W}$ for $n$ agents and $H\times W$ pixel, and change the Equation 4 to

$$ $\displaystyle L_{MSE}=\sum_{i=1}^{n}C_{i}\cdot\|\epsilon_{\theta}(x,t|a_{i})- \epsilon\|^{2}$ (7) $$

We simply set the loss coefficient matrix based on each agent’s reachable region in the image to supervise the model to focus more on the related pixels.
We have observed that even with this straightforward loss coefficient approach, there is a significant improvement in the model’s accuracy of modeling $P_{\theta}(x|a_{i})$.

Figure: Figure 3: (a) Given partial egocentric RGBD observations, COMBO first reconstructs and inpaints the top-down orthographic image as the overall world state estimation. (b) COMBO then leverage the planning sub-modules built with Vision Language Models to propose actions, infer other agents’ intents, and evaluate the outcomes simulated with the compositional world model to plan online with a tree search procedure to cooperate in the long run.

## 5 Compositional World Model for Multi-Agent Planning

To plan efficiently in the multi-agent cooperation problem, we still need to address the challenge of partial egocentric observation and model complex world dynamics in the long run. We propose COMBO, a novel Compositional wOrld Model-based emBOdied multi-agent planning framework, shown in Fig. 3. After receiving the egocentric observations during the last execution step, COMBO first estimates the overall world state to plan on, as discussed in Sec. 5.1. COMBO then utilizes Vision Language Models to act as an Action Proposer to propose possible actions, an Intent Tracker to infer other agents’ intents, and an Outcome Evaluator to evaluate the different possible outcomes, as detailed in Sec. 5.2. In combination with the compositional world model to simulate the effect of joint actions on the world state $s_{i+1}=CWM(s_{i},a)$ introduced in Sec. 4, we discuss the tree search procedure to integrate these planning sub-modules in Sec. 5.3. COMBO empowers embodied agents to imagine how different plans may affect the world with other agents in the long run and plan cooperatively.

### 5.1 World State Estimation with Partial Egocentric Views

Directly planning based on partial egocentric views presents a considerable challenge. To address this, we initially reconstruct partial point clouds from multiple egocentric RGBD views $o_{i}$ captured from different perspectives. These point clouds are then transformed into a unified top-down orthographic image representation, serving as the world state. This representation is constructed by overlaying the views in chronological order. It is important to note that our world is inherently dynamic, with other agents actively interacting, resulting in a top-down orthographic image representation that is both noisy and incomplete, as depicted in Fig. 3 (a). To refine this representation, we employ a diffusion model to inpaint the partial and noisy orthographic image, thereby enhancing the estimation of the world state $s_{i}$ for subsequent planning. This approach allows us to effectively represent and rectify the world state, enabling more accurate planning in multi-agent environments.

### 5.2 Planning Sub-modules with Vision Language Models

Action Proposer Given the estimateed overall world state $s_{i}$ and the long horizon goal $G$, the Action Proposer first searches over the potential action spaces $A_{i}$, and then proposes multiple possible actions $a_{i,1...p}=AP(s_{i},G)$ to explore on. We implement this module by querying a VLM to generate the possible actions in the text given the task goal and encoded image world state as context. We finetune LLaVA on randomly collected rollouts with possible actions labeled to construct this module.

Intent Tracker Inferring what other agents may do given observation history is important for effective multi-agent cooperation. We implement this module $a_{-i}=IT(s_{i},G)$ by keeping the estimated image world state from the last $k$ steps and then feed into the VLM together with the task goal and to query for the possible actions of other agents $a_{-i}$. We construct this module by finetuning LLaVA on collected short rollouts with other agents’ actions labeled.

Outcome Evaluator It’s vital to have a way to assess the value of the achieved state from different plans so the search can be effectively deepened utilizing pruning. We implement an Outcome Evaluator to fulfill this functionality $v=OE(s,G)$ by generating a heuristic score $v$ for each image state $s$ given the task goal $G$. To construct this module, we finetune LLaVA to describe the state of each object in the image and the corresponding heuristic score representing steps left to achieve the task goal considering all objects.

### 5.3 Planning Procedure with Tree Search

Figure: Algorithm 1 COMBO Planning Procedure for Agent $i$.

With powerful generative sub-modules implemented so far, we can use an effective tree search algorithm to integrate them into achieving the best planning performance. Formally, we need to search for a sequence of actions that compose the plan to achieve the task goal with a minimal number of steps. To sample the long-horizon plans, we can chain the sub-modules to expand future states from current state $s_{i}$ with $s_{i+1}=CWM(s_{i},AP(s_{i},G),IT(s_{i},G))$. By recursively deploying these chained modules we can expand the plans until achieving the goal state. However this is impractical due to the large search space, we implement a limited tree search procedure instead by always keeping the $B$ best-scored plan beams and exploring $P$ action proposals at each plan step for a maximum of $D$ rollout steps. We show the complete COMBO planning procedure in Algorithm 1.

Due to the uncertainty of other agents, COMBO replans at every step to reflect the fast-changing world state. As similarly observed in [], when searching for the best plan, the Outcome Evaluator may leverage the irregular artifacts of the Compositional World Model to get artificially high scores, such as where key objects are teleported to desired positions. To mitigate this artificial exploitation issue, the Outcome Evaluator will use the default score from the last state if the new state generated from the Compositional World Model suspiciously increases the score above a threshold.

## 6 Experiments

### 6.1 Experimental Setup

We build two challenging embodied multi-agent planning tasks TDW-Game and TDW-Cook on ThreeDWorld [], where 2-4 decentralized agents must cooperate to finish some long-horizon goals, as illustrated in Fig. 1. All the agents have an observation space of egocentric $336\times 336$ RGBD images and the corresponding camera matrix. Both tasks have a maximum planning horizon of 60 steps. In TDW-Game, 3-4 agents cooperate to pass and place 6-8 puzzle pieces scattered randomly on the table into the correct puzzle box according to visual clues such as the shape. The location of the puzzle box and pieces are randomly initialized across different episodes. The action space includes

- •
wait
- •
pick up <obj>
- •
place <obj> onto <loc>

In TDW-Cook, 2 agents cooperate to pass, cut, and place 6-8 out of 10 possible food items randomly scattered on the table to make some dishes according to the recipe. The location of the food items and the recipe are randomly initialized across different episodes. Specifically, there is only one cutting board where agents can pass or cut objects, making cooperation vital for efficient plans. The action space includes

- •
wait
- •
pick up <obj>
- •
cut <obj>
- •
place <obj> onto <loc>

Metrics We evaluate the agent’s cooperation ability by task success rate and average steps of the successful episodes when cooperating with different agents across 10 episodes. An episode is done when the goal state is reached and considered successful or the maximum planning horizon is reached and considered failed. Specifically, we implement two different agent policies as cooperators on both tasks to reliably evaluate the agents. In TDW-Game, Agent 1 takes the policy of always passing the unwanted puzzle pieces in a clockwise manner, while Agent 2 takes the policy of always passing the unwanted puzzle pieces in a counter-clockwise manner. In TDW-Cook, Agent 1 takes a “selfish” policy of always operating objects in its own recipe first, while Agent 2 takes an “altruism” policy of always prioritizing operating objects in the cooperator’s recipe.

Baselines
We compare our methods to several multi-agent planning methods:

- •
Recurrent World Models[], using VAE to encode the visual observations and a generative recurrent neural network to learn the world model. An evolution-trained simple policy is used as a controller.
- •
MAPPO[], a MARL baseline where agents with shared weight Actor and Critic are trained jointly with PPO.
- •
CoELA (Oracle Vision)[], an LLM-based agent with access to the Oracle vision perception of the state re-implemented due to the lack of ground truth segmentation needed to train the Perception Module of CoELA.
- •
LLaVA[], LLaVA-v1.5-7B fine-tuned with the same collected rollouts to train the Action Proposer in our COMBO, directly learn a mapping from the visual observation and task goal to action.
- •
Oracle Cooperator (Oracle Vision) has access to the Oracle state of the environment and other agents’ policies, acting as a strong baseline.

**Table 1: Main results on TDW-Game and TDW-Cook. We report the mean success rate and average steps of the successful episodes over 10 episodes here. COMBO (w/o IT) denotes COMBO without the Intent Tracker module. Oracle Cooperator has access to the Oracle state of the world and other agents’ policies. Agent 1 and Agent 2 are cooperators with different policies, as detailed in Sec. 6.1.**
|  |  | TDW-Game | TDW-Cook |  |  |
| --- | --- | --- | --- | --- | --- |
|  | Cooperating with | Agent 1 | Agent 2 | Agent 1 | Agent 2 |
| Recurrent World Models[] | Success Rate | 0 | 0 | 0 | 0 |
| Average Steps | / | / | / | / |  |
| MAPPO[] | Success Rate | 0 | 0 | 0 | 0 |
| Average Steps | / | / | / | / |  |
| CoELA (Oracle Vision)[] | Success Rate | 1.0 | 0.8 | 0.5 | 0.2 |
| Average Steps | 17.8 | 26.1 | 38.2 | 29.5 |  |
| LLaVA[] | Success Rate | 0.9 | 1.0 | 1.0 | 0.9 |
| Average Steps | 26.0 | 31.1 | 31.5 | 36.7 |  |
| COMBO (w/o IT) | Success Rate | 0.8 | 0.7 | 0.9 | 0.9 |
| Average Steps | 15.6 | 17.4 | 24.8 | 21.0 |  |
| COMBO (Ours) | Success Rate | 1.0 | 1.0 | 1.0 | 1.0 |
| Average Steps | 16.8 | 18.1 | 21.3 | 21.7 |  |
| Oracle Cooperator | Success Rate | 1.0 | 1.0 | 1.0 | 1.0 |
| (Oracle Vision) | Average Steps | 14.9 | 16.8 | 23.4 | 19.3 |

Implementation details The video diffusion model of the compositional world model is built upon AVDC[] codebase, with the architectural modification of introducing a cross attention layer to the text condition in the ResNet block and replacing the Perceiver with an MLP to enhance the text conditioning. We use the base resolution $128\times 128$ for the video diffusion model and train a super-resolution diffusion model to get the final $336\times 336$ images. We use the T5-XXL encoder to pre-process all the text conditions for a better contextual representation. All vision language models used in the main experiments are LLaVA-v1.5-7B. The planning parameters are set to Action Proposes P=3, Planning Beams B=3, and Rollout Depths D=3 unless other specified. More details are in Appendix A.

Figure: Figure 4: Joint action-conditioned video generation. Our Compositional World Model can simulate world dynamics conditioned on the joint action of multiple agents accurately while AVDC struggles with simulating which agents should act, and COMBO w.o ADLS may simulate actions incorrectly.

**Table 2: Human evaluation of generated videos from different World Models and the corresponding plan performance. Single denotes the accuracy of synthesized videos conditioned on a single action of an agent among $20$ samples. Multiple denotes the accuracy of synthesized videos conditioned on multiple actions of multiple agents among $20$ samples. Plan denotes the average steps (success rate) across cooperating with both Agent 1 and Agent 2 over 10 episodes. Best performances are in bold.**
|  | TDW-Game | TDW-Cook |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | Single | Multiple | Plan | Single | Multiple | Plan |
| AVDC | $65\%$ | $20\%$ | 29.7(80%) | $85\%$ | $25\%$ | 34.5(90%) |
| COMBO w.o ADLS | $70\%$ | $55\%$ | 26.9(100%) | $85\%$ | $70\%$ | 28.3(100%) |
| COMBO | 95% | 75% | 17.5(100%) | 100% | 100% | 21.5(100%) |

### 6.2 Results

COMBO can cooperate efficiently with different cooperators As shown in Tab. 1, both agents of different policies achieve the best performance when cooperating with COMBO on both tasks, finishing all the tasks within the limited horizon and cost the least steps. Both Recurrent World Models and MAPPO perform poorly due to the difficulty of the multi-agent long-horizon task with only egocentric visual observations. Compared to LLM Agent CoELA and VLM Agent LLaVA, COMBO with a compositional world model to explicitly model the world dynamics and plan accordingly achieves more efficient cooperation.

Intent Tracker contributes to efficient cooperation Comparing the results of COMBO (w.o IT) and COMBO in Tab. 1, we can see the Intent Tracker module is an important module that empowers the agent to infer other agents’ intents and consider them during planning for better cooperation in the long run.

Compositional World Model is crucial to make cooperative plans We compared the generated video quality with AVDC[], a video diffusion model designed for robotics tasks, and ablate on the effect of the Agent-Dependent Scaling Loss (COMBO w.o. ADSL) in Tab. 2 and Fig. 4. Compared to AVDC, our compositional world model can accurately simulate which agents are acting and what interactions are described in the text condition, leading to a large performance gain in the accuracy of modeling agent actions on the world state. Removing the Agent-Dependent Scaling Loss leads to considerable accuracy degradation, showing the effectiveness of the introduced technique. Moreover, with the compositional world model, the overall plan performance boosts to 17.5 steps for TDW-Game compared to 29.7 steps with AVDC.

More computation budgets lead to better Plan We study the effect of computation budgets on the long-horizon plan quality in Tab. 4, where we can see more computation budgets with more action proposals, larger plan beams, and deeper rollout depth leads to higher success rate and less average steps. A qualitative example of generated plans with different computation budgets is shown in Fig. 5, where more searches can obtain better plans in the long run.

**Table 3: Plan performance improves with more compute budgets. We report the mean success rate and average steps of the successful episodes cooperating with Agent 1 and Agent 2 over 5 episodes on TDW-Game here.**
| Action | Rollout | Plan | Success | Average |
| --- | --- | --- | --- | --- |
| Proposals $P$ | Depth $D$ | Beam $B$ | Rate | Steps |
| 3 | 1 | 1 | 50% | 27.8 |
| 2 | 2 | 2 | 70% | 18.1 |
| 3 | 3 | 3 | 100% | 17.5 |

Figure: Figure 5: More Computation budgets leads to better plan. With increased computation budgets (lower part), COMBO can search for a better plan where Alice first clears the common region with David so that he can pass the next puzzle piece to Alice instead of having to wait, leading to a better world state after same number of steps.

COMBO can generalize to an arbitrary number of agents COMBO trained with only data of four agents playing TDW-Game can surprisingly generalize to three agents version well as shown in Tab. 4, showing the strong generalization and promising future of the compositional world model-based modularized planning framework for multi-agent cooperation.

## 7 Conclusion

In this work, we learn a compositional world model for embodied multi-agent cooperation by factorizing the naturally composable joint actions of an arbitrary number of agents and compositionally generating the video to enable accurate simulation of multi-agent world dynamics. We then propose COMBO, a novel Compositional World Model-based embodied multi-agent planning framework to empower the agents to infer other agents’ intents and imagine different plan outcomes in the long run. Our experiments on two challenging embodied multi-agent cooperation tasks show COMBO can enable the embodied agents to cooperate efficiently with different agents across various tasks and an arbitrary number of agents.

Future Work Our method leverages tree search combined with Large Generative Models to devise long-horizon plans. The necessity for multiple inferences using large models leads to a relatively slow inference speed, restricting the applicability of our method in scenarios demanding rapid response. Exploring the development of more efficient models could potentially mitigate this drawback and enhance the practicality of our approach in time-sensitive environments.

## Appendix 0.A Additional Experiment Details

### 0.A.1 Task

Figure: Figure 6: Objects of TDW-Cook and TDW-Game. Left: 13 food items that may occur on the table while 10 of them (except for the third row) may occur in the recipe. Right: 4 puzzles containing 3 to 4 pieces each with different colors and shapes.

We build two challenging embodied multi-agent planning tasks TDW-Game and TDW-Cook on ThreeDWorld [], where 2-4 decentralized agents must cooperate to finish some long-horizon goals. The object assets of the two tasks are illustrated in Fig. 6.

### 0.A.2 Compositional World Models

##### Dataset Collection

We collect random rollouts with a scripted planner and generated 107k videos of TDW-Game and 50k videos of TDW-Cook.

##### Model Parameters

The parameters of our inpainting diffusion model, super-resolution diffusion model, and video diffusion model are shown in Tab. 5.

**Table 5: Model Parameters.**
|  | Inpainting | Super-Resolution | Video Diffusion Model |
| --- | --- | --- | --- |
| num_parameters | 84M | 38M | 198M |
| input_resolution | $336\times 336$ | $128\times 128$ | $128\times 128$ |
| output_resolution | $336\times 336$ | $336\times 336$ | $128\times 128$ |
| base_channels | 64 | 64 | 128 |
| num_res_block | 2 | 2 | 2 |
| attention_resolutions | (16,) | (16,) | (8, 16) |
| channel_mult | (1, 2, 4, 6, 8) | (1, 2, 3, 4, 5) | (1, 2, 3, 4, 5) |

##### Agent-dependent Loss Scale

We use a simple method to set the loss scale. If agent $i$ is located in the upper half of the image, then we set $C_{i,1:H/2,1:W}=2$ and others are $1$. Similarly, the same approach is applied to other situations or conditions.

##### Compute

We train the world model for 50k steps in the first stage with 384 batch-size on 192 V-100 GPUs in 1 day. And then fine-tune the model for 25k steps in the second stage with 120 batch-size on 120 V-100 GPUs in 1 day. Both the inpainting model and the super-resolution model are trained for 60k steps with 288 batch-size on 24 V-100 GPUs in 1 day.

##### Samping

We use DDIM sampling across the experiments with guidance weight 5 for the text-guided video diffusion model.

### 0.A.3 Baselines

#### 0.A.3.1 MAPPO

Our task is multi-agent reinforcement learning with decentralized observation and control, as different agents have disjoint observations to produce their actions. We followed the training procedure in []. Each agent is an actor-critic network where the actor and critic share a common convolutional backbone. We pre-define a set of all possible actions for each environment and let the PPO agent choose among them. We design rewards to reflect the steps left to finish the episode in a similar spirit to the Outcome Evaluator’s heuristics score design.

However, in our setting, the embodied multi-agent simulation samples are relatively slow, creating a higher demand for sample efficiency of the method. We conducted a grid search on hyper-parameters of learning rate in $\{1e^{-6},5e^{-5},1e^{-3}\}$, batch size in $\{2,64\}$. However, the results are poor and the behavior either falls to near-random or collapses to a single action, which verifies the challenge of cooperative planning on ego-centric RGBD observations.

#### 0.A.3.2 Recurrent World Models

We followed the training procedure in [], firstly training a variational autoencoder (VAE) on the same data collected for training COMBO to encode the egocentric observations in the TDW-Game and TDW-Cook environments into a 32-dimensional latent vector $z\in\mathbb{R}^{32}$. Then, we train a mixture density network combined with a recurrent neural network (MDN-RNN) for predicting the future latent vector given a textual action which is encoded by BERT[], the current egocentric observation and the current hidden state of the RNN, which is modeling $P(z_{t+1}|a_{t},z_{t},h_{t})$. Finally, we train two simple controllers using the covariance-matrix adaptation evolution strategy (CMA-ES) [] to maximize the expected cumulative reward of a rollout by interacting with the actual environments.

The recurrent world model failed to model world dynamics accurately conditioned on a single action, as shown in Fig. 7. Specifically, we predict the latent code $z_{t+1}$ from MDN-RNN and use the decoder of the VAE to generate the image of the next time step. Additionally, the simple controller doesn’t work well on a task that requires complex planning as it keeps collapsing to single actions in our experiments.

Figure: Figure 7: Recurrent world models failed to model world dynamics conditioned on single actions accurately.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - 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?
- As mentioned in the weaknesses, I have concerns about potential scalability. Can this be generalized to continuous action spaces, such as in robotics tasks? If so, it would be beneficial to include experiments on this. Additionally, 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.
soundness: 3
strengths: **Novelty**. Suggested framework, named COMBO, offers a unique solution to the multi-agent planning problem by utilizing compositional world modeling for accurate simulation.

**Clear Framework Explanation.** The framework is presented clearly and is easy to follow. Each setting and procedure is understandable through Figure 3 and Algorithm 1. The roles of each module are well-explained in the text and formulation.

**Well-designed Experiments And Clarified Implications of Results.** The experimental setup and baseline choices effectively test the COMBO approach's capabilities. The results demonstrate impressive performance and highlight the necessity of the design choices in their architecture.
summary: This paper proposes a compositional world model for multi-agent cooperation, leveraging large generative models and compositional diffusion models to build accurate simulations with an arbitrary number of agents. 

The authors present a novel framework called COMBO, which involves the following procedures:

1. **Estimate the Current World State**: Use a diffusion model to infer the state from partial egocentric observations.
2. **Action Proposal and Evaluation**: Employ a pretrained Vision-Language Model (VLM) to suggest actions, predict other agents’ intentions.
3. **Future Frame Generation**: Produce future frames based on the current world state image and a text prompt, representing joint actions of multiple agents. This is performed by the Compositional World Model, leveraging a compositional diffusion architecture.
4. **Evaluate and Plan**: Assess the predicted outcomes from the compositional world model using the VLM, and plan with a tree search algorithm based on these evaluations.

To assess their model, the authors test their approach on three datasets—TDW-Game, TDW-Cook, and 2D-FetchQ—and compare it to several baselines, including a VAE-based world model, MAPPO, CoELA, and LLaVA. The results demonstrate that the COMBO framework significantly outperforms these baselines, especially in planning capabilities.
weaknesses: **Lack of Figure Clarity and Interpretability.** The figures in the paper are unclear and difficult to interpret. Illustrations should enhance understanding, but these require reading the text to decipher them. For instance, Figure 1-(b) displays a random assortment of images without any labels. I believe this can be resolved by adding explicit labels for sequential processes that each frame means. Similarly, Figure 4 presents consecutive frames without explanations. Adding annotations, e.g. state, instruction, prediction t=1~3, would make the figures more informative.

**Limited Scalability and Impact.** 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 and operating without access to ground truth environmental labels 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.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: - 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?  

- Potentially more details could be provided on modification of 2D FetchQ environment, as I didn't find any in the appendix. 

- Why Co-Gaild baseline is used in this task but not used on the original environments? 

- Is usage or fine-tuning of VLM really necessary for these tasks? If so, does this fine-tuning leads to generalizable action generation /agents intends? For example, how robust agents would follow recipes if some ingredients are not seen in the training? Without studying and showing such generalizations, it is not clear why LLM are needed and if more simple supervised models would do the same job.  

- Some formulas could be better connected. E.g in section 4 X is used for world model state, while in 5 s is used.
soundness: 2
strengths: - The problem of join cooperation of several agents is interesting, and the approach for world modeling by merging observations from different agents seems plausible. 

- Decomposition of the full state to regions that can be affected by each agent (while not correct in general e.g. turning on light would change the whole image, however it is a reasonable assumption that the overall scene in effected by agents mostly independently)
summary: The paper propose to use compositional video diffusion model as "world model" for an environment with several agents cooperatively solving the task. In addition to learning scores for each agent independently, the diffusion model is trained only on regions reachable by an agent. Such model was used in COMBO agent that is effectively planning by fine-tuning VLM to propose reasonable action given current states while also fine-tuning VLM for intend prediction and for the evaluation of the current state.
weaknesses: - 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. Also, what about regions what are not reachable by any agent? In current formulation, it is not clear if those regions are modeled or not in the world model.  

 - 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). 

- In both of the TDW-Game and TDW-Cook performance is saturated learning to agents that perform near (and in case of the TDW-Cook somehow better 22.8 vs 24.0) than Oracle Cooperator. Thus, it would be great if authors can increase the difficulty of those tasks and show possible failure cases of their approach. For example, currently world model assumes that visual information and actions is enough to predict the next state. However, for in closer to real world scenarios some parameters are hidden (e.g. objects mass) and thus effective world model would need to deal successfully with estimation of such parameters and using them for the determining of the optimal actions.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - 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.
soundness: 3
strengths: - They designed explicit compositional world model. It is one of the most distinguishable designs in their modeling, which supports the look ahead planning (tree search planning), and they showed it is beneficial through the empirical evaluation results. As a part of this, the world state estimation is good to build the world model for multi-agent setting.
- It is a well written paper. I can easily follow their discussions without unnecessary questions.
summary: The authors of this paper explored a compositional workload model for embodied multi-agent cooperation. Unlike previous approaches to embodied multi-agent cooperation, they introduced an explicit world model that simulates the next state by incorporating the estimated actions from the Intent Tracker along with the agent's actions. This model stands out from other single-agent, model-based reinforcement learning approaches by explicitly accounting for multi-agent cooperation dynamics. Utilizing this world model, they implemented a tree search planning strategy similar to Monte-Carlo Tree Search [1]. Their approach demonstrated superior performance in embodied multi-agent cooperation tasks and, due to the compositional nature of their model, exhibited strong generalization capabilities.

[1] Silver, David, et al. "Mastering the game of Go with deep neural networks and tree search." nature 529.7587 (2016): 484-489.
weaknesses: - The room of the evaluated benchmarks is too small to show the effectiveness of their proposed modeling. 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.
- 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? If yes, it should be mentioned in the caption of the figure or the paragraph. If no, then we think this empirical evidence is weak to show the generalizationality of the COMBO.

Meta-review: COMBO presents a significant contribution to embodied multi-agent cooperation through its novel compositional world model approach, demonstrating superior performance across multiple benchmarks. While reviewers initially raised concerns about scalability and generalization capabilities, the authors provided comprehensive responses including new real-world human-robot experiments, quantitative results showing strong generalization to different numbers of agents, and thorough ablation studies demonstrating the necessity of their design choices. The extensive validation across both simulation and real-world scenarios, combined with clear practical impact for multi-agent robotics, makes this a valuable contribution worthy of publication. I recommend accepting this paper.

Decision: Accept (Poster)