Grounded Answers for Multi-agent Decision-making Problem through Generative World Model

1. Motivation seems misleading, and results are limited to the SMAC benchmark.

We aim to improve the output of large language models by constructing a differentiable interactive environment and then learning the policy on this interactive environment. The paper focuses on constructing such an environment, including the dynamics and reward models, and demonstrates the potential application of this framework.

Due to the substantial computational resources required to train MARL methods, especially in most cooperative MARL environments that lack image inputs and have extremely low sampling efficiency, we have conducted experiments only on SMAC, which we have mentioned in the limitations section.

Although SMAC cannot represent all possible multi-agent tasks, as a highly complex multi-agent reinforcement learning environment, SMAC involves heterogeneous agents, varying numbers and tactics, and different maps, making the learning process inherently challenging. This complexity can reflect the proposed algorithm's potential applications.

The trick of masking out dead agents aims to simplify the inference difficulty of algorithms, and it seems task-specific. However, it could be a general trick for multi-agent environments with attrition, and it can also be applied to scenarios like MOBA or RTS games. In addition, developing such tricks might be a promising direction in generative models. For example, the outputs of large models are not always reliable, such as generating movements that do not satisfy physical constraints. Introducing human priors into scene prediction results could be an essential way to solve these issues. Methods like PINNs and Lagrangian conservation can impose constraints on certain predictions, thereby enhancing the reliability of the generated results. The tricks specific to the SMAC game in this paper can also be categorized as one of these methods.

2. The generalizability of the proposed method.

The proposed method has a certain generalization ability in multi-agent environments.

• Regarding the dynamic is explicit issue, if we understand correctly, you mean that the input for the reward is the entire trajectory rather than (s, a), as in most previous methods. This is necessary because we aim to achieve temporal utility distribution across the entire trajectory, which cannot be done by relying solely on (s, a), thus failing to address issues like delayed reward. During inference, we first run the dynamics model to obtain the entire reward-free trajectory and then run the reward model to label the rewards so the entire trajectory is known.
  If you mean that dynamics might not be deterministic, such as taking the same action in a given state leading to different following states and rewards, this is not considered in this paper. This is because most multi-agent reinforcement learning environments, including SMAC, Google Research Football, and Multi-agent Particle World, are deterministic. To address this, we may modify the loss function by introducing heteroscedastic aleatoric uncertainty estimation - we can model the predicted state and reward as distributions and estimate the mean and variance of both state and reward.

• Concerning the issue of simulation speed being fast enough for learning and interaction, in our world model, dynamics is an auto-regressive generation process corresponding to interacting with the real environment (e.g., the function of env.step()). This is time-consuming and a common issue in reinforcement learning algorithms. The reward is predicted based on a bidirectional Transformer by inputting the entire trajectory, which is less time-consuming than the auto-regressive dynamics model. Additionally, since the world model is implemented entirely through neural networks, distributed data collection is easier than with the original game engine and does not require communication with the game engine (which is the main reason for the slow interaction in the SMAC environment). When transferred to a real environment, distributed data collection without interacting with the real environment offers higher speed and safety than online learning.

1. The exploration methods (EMC and IIE) for offline data collection

EMC is curiosity-based to boost multi-agent exploration for novel states, i.e., reduce out-of-distribution regions in the offline datasets. IIE belongs to go-explore algorithms that perform imagination without exploration, initialize agents at the end state of the imagination, and then start to explore. This method can reduce the multi-agent exploration space in the time dimension. We chose them to collect data because they have tuned in to the SMAC benchmark and have performed well in super hard maps.

When collecting offline datasets in other MARL tasks, we must choose appropriate exploration methods based on the task properties. If the task involves many decision steps (long-horizon task), go-explore methods can be more effective; if the action space of each agent is huge, we can perform role-based learning (RODE) to decompose the action space; if we want to cover the entire state space, we can use curiosity-based exploration methods (EMC) or hierarchical exploration (MAVEN). Since the joint action and state space are huge in SMAC, we chose EMC to reduce the impact from OOD regions and used IIE to reduce exploration space.

2. $L_t$ in the dynamics model

Thanks for pointing it out. The input only contains one target sequence. In general, $L_t$ should reflect the task and the environment, which can be time-varying. The language description for the trajectory from $s_t$ and the terminated (goal) state $s_T$, e.g., return-to-go). However, in practice, we simplify this language description to be time-invariant, which is always the language description of terminated (goal) state $s_T$. We have corrected this for better understanding.

3. Hindsight-relabelling in the reward model

We use the same text label for the reward and dynamics models to specify the trajectory and the reward generation. We use hindsight learning to provide a specific condition for reward generation to avoid the influence of suboptimal data in the offline dataset. Without this condition, the reward model would learn to approximate all possible trajectories in the dataset. Due to the presence of suboptimal trajectories, this would negatively impact the reward model's learning. We train the reward model using conditional generation to learn the reward functions corresponding to different outcomes. During inference, we set a condition, such as asking the agents to completely eliminate the enemy, to extract the reward function corresponding to a successful trajectory from the reward model.

In policy models, this type of condition is usually set as return-to-go (RTG) to achieve desired effects. In contrast, we use language descriptions to set the conditions, which allows for the generation of an interactive simulator and enables the model to have multimodal processing capabilities.

4. Only focus on the SMAC benchmark

Due to the substantial computational resources required to train MARL methods, especially in most cooperative MARL environments that lack image inputs and have extremely low sampling efficiency, we have conducted experiments only on SMAC, which we have mentioned in the limitations section.

Although SMAC cannot represent all possible multi-agent tasks, as a highly complex multi-agent reinforcement learning environment, SMAC involves heterogeneous agents, varying numbers and tactics, and different maps, making the learning process inherently challenging. This complexity can reflect the proposed algorithm's potential applications.

5. Why do some tables report win rates and some rewards?

In SMAC, a slight difference in return (19 vs. 20) can result in a significant difference in win rate (0% vs. 90%). If we only show the win rate, the results for most current offline methods can always be 0, making comparisons less meaningful. Using cumulative rewards as a metric is also reasonable because reinforcement learning inherently optimizes cumulative rewards.

6. Consider the number of time steps in the data collection.

First, the pre-trained world model is an interactive simulator, and we compare it with the real environment, not the online methods themselves. This experiment aims to demonstrate that training a randomly initialized joint policy on the learned world model is faster than training it in the real environment - the learned world model (interactive simulator), specifically the learned transition and reward functions, is accurate. So, we do not shift the LBI curve to the right by the number of time steps in the data collection.

7. Clarifications for equation (1).

$D$: the demonstrations provided by an expert policy.

The index $i$ of $\tau$ in the $f$ term: Thanks for pointing this out. We perform this regularization on the predicted rewards from the expert demonstrations and the data collected by the MA-SAC policy. Namely, in the $f$ term, $\tau\sim \\{\boldsymbol{\pi}^\theta,D\\}$, $i$ and $t$ should be added in the agent and the timestep dimensions.

Over all possible actions: This is an implenmentation trick. In practice, the reward model $RM$ takes the action-free trajectory and outputs the rewards for all actions of each agent. Namely, the individual reward head outputs $|U|$ rewards for each agent, where $|U|$ is the action space. Then, we choose the reward corresponding to the action using the torch.gather function. We add the L2 penalization for these $|U|$-dim predicted rewards to avoid overestimation on out-of-distribution actions. For example, if the offline dataset does not contain the action $u_t^i$, then the predicted reward of $u_t^i$ does not have a target, introducing potential risks for policy learning.

1. It requires significant computational resources

We listed the hyperparameters in the interactive simulator in Tables 8 and 9, which requires computational resources. However, we believe this framework is meaningful. We utilize the generative model as an interactive environment for RL to enhance the quality of large models' responses to decision-making problems. RL requires significant interaction with the environment. Since the world model is implemented entirely through neural networks, distributed data collection is easier than with the original game engine and does not require communication with the game engine (which is the main reason for the slow interaction in the SMAC environment). When transferred to a real-world task, distributed data collection without interacting with the real-world environment offers higher speed and safety than online learning.

2. The scope might be too narrow.

Due to the substantial computational resources required to train MARL methods, especially in most cooperative MARL environments that lack image inputs and have extremely low sampling efficiency, we have conducted experiments only on SMAC, which we have mentioned in the limitations section.

Although SMAC cannot represent all possible multi-agent tasks, as a highly complex multi-agent reinforcement learning environment, SMAC involves heterogeneous agents, varying numbers and tactics, and different maps, making the learning process inherently challenging. This complexity can reflect the proposed algorithm's potential applications.

3. Unfair comparisons.

First, the pre-trained world model is an interactive simulator, and we compare it with the real environment, not the online methods themselves. This experiment aims to demonstrate that training a joint policy on the learned world model is faster than training it in the real environment - the learned world model (interactive simulator), specifically the learned transition and reward functions, is accurate.

Second, this paper mainly compares offline Q-learning and inverse learning methods, demonstrating their generalization ability to unseen tasks.

4. The effectiveness of image reference.

Using images as a reference can improve performance (LBI and LBI-wo-IR), although less significant than the residual term (LBI and LBI-wo-RT). However, we demonstrated that the error and return when using real images (LBI-GTI) are similar to LBI, indicating that the benefit of image reference is limited rather than suggesting that the model is not optimized. The visual results show that the generated images are similar to the real ones.

It is important to note that the model uses images and text as initial inputs and then performs autoregressive prediction, making images necessary for this task.

5. The world model architecture.

The world model consists of a VQ-VAE model for the image tokenizer, a casual transformer model for dynamics prediction, and a bi-directional transformer model for reward prediction. We build the VQ-VAE and the transformer model implementation based on (https://github.com/Project-MONAI/GenerativeModels/blob/main/generative/networks/nets/vqvae.py) and Decision Transformer (which borrows the code from minGPT https://github.com/karpathy/minGPT), the publicly available re-implementation of VQ-VAE and GPT2, respectively. The dynamics and reward models share the same transformer architecture, where the reward model does not use causal masks. We mentioned it in Appendix D.3.

6. Traning and inference.

The model is trained in two phases following a standard autoregressive video generation pipeline: We train the image tokenizer first, which is used for the dynamics and reward models. We then train the dynamics model and reward model separately.

VQ-VAE: (1) sample a batch of images; (2) calculate the reconstruction and the quantization error; (3) optimize the VQ-VAE model.

Dynamics model: (1) sample a batch of trajectories from the offline buffer; (2) calculate the output for each token; (3) optimize the dynamics model by minimizing the cross-entropy loss for actions and mean-squared error for others.

Reward model:

• Randomly initialize a MA-SAC policy network

• While True do:

  • Use MA-SAC to collect a batch of trajectories $B$ from the ground-truth environment

  • Sample a batch of trajectories $D$ from the offline buffer

  • Mask out rewards of the trajectories $B$ and $D$

  • Generate the reward via the reward model

  • Update the reward model by maximizing the predicted reward for $D$, and minimizing the predicted rewards for $B$;

  • Update MA-SAC policy network for k-step (k=5).

Inference:

• Randomly initialize a policy network

  • While True do:

    • Given an initialized image and a task description, calculate the image embedding via the VQ-VAE model and the language embedding via a pre-trained language tokenizer named bert-base-uncased

    • While all agents or all enemies are alive do

      • Generate state via dynamics model

      • Generate joint observations via dynamics model

      • Generate actions (only for the behavior-regularization) via the dynamics model

      • Obtain joint actions (for dynamics prediction) via current policy network (using $\epsilon$-greedy)

      • Generate the next image embedding

      • Concatenate the language embedding

    • Add masks into the reward-free trajectory for the reward prediction

    • The reward model takes the entire trajectory as input and generates corresponding rewards

    • Store the whole trajectory with actions for the behavior regularization in replay buffer

    • If the trajectories in replay buffer > batch size then

      • Sample a random batch of trajectories

      • Calculate the TD loss with the behavior regularization for the policy network

      • Update policy network

1&2. Trivial framework and lack of investigation of related multi-agent world model literature.

If I understand correctly, the distinct characteristic of multi-agent scenarios you mentioned should be decentralized execution, which is a common way to simplify computational demands during the execution and thus enhance the scalability in practical applications. However, centralized training is generally required during the training process because agents learn the joint policy in a non-stationary environment. For example, QMIX [1] uses monotonic value decomposition to learn the joint Q-value function, and HAPPO [2] proposes sequential learning and multi-agent advantage decomposition. Otherwise, using individual Q-learning could lead the algorithm into local optima (unless we use some independent learning methods like lenient learning or design individual reward functions).

The most significant difference is that our aim is to pre-train a world model and utilize it to train the policy model rather than directly using it as a part of the policy model.

• [3] is an image-based single-agent world model, while ours is a multi-modal multi-agent world model. This dynamics model structure is general and is also used in GAIA, DT, TT, and Genie, which we mentioned in the related work on model-based and Transformer methods.

• [4][5][6] are online learning methods that use the world model to generate imagined future trajectories at a given state for planning or better decision-making. Our goal is to learn an interactive simulator from offline datasets.

• [3][4][5][6] suppose the reward is given, and [4][5][6] enable communication during execution to add necessary inputs for the agents. We suppose the reward is unknown and use inverse RL to learn the reward, and one can make use of different MARL methods (e.g., model-based or model-free, value-based or policy-based) to learn the policy on the proposed world model.

For the world model, since we do not use it in the policy execution but treat it as a simulator – the learned policy can still be decentralized. To provide enough information for the input, we train directly using joint behaviors and observations.

Thanks for pointing out these works. We would like to add them to the World Model and Multi-agent Reinforcement Learning subsection.

3. Unfair comparisons.

First, the pre-trained world model is an interactive simulator, and we compare it with the real environment, not the online methods themselves. This experiment aims to demonstrate that training a joint policy on the learned world model is faster than training it in the real environment - the learned world model (interactive simulator), specifically the learned transition and reward functions, is accurate. Designing a multi-agent environment, precisely the reward function, is complex and requires much prior. Our framework shows that utilizing the world model (a causal transformer for dynamics and a bidirectional transformer for reward) can successfully perform inverse reinforcement learning and learn meaningful reward functions.

Second, this paper mainly compares offline Q-learning and inverse learning methods, demonstrating their generalization ability to unseen tasks. The comparison with online MARL is only briefly mentioned in the main text, with detailed results provided in the appendix.

4. Heavily handcrafted processing - the agent death masks.

This design aims to simplify the inference difficulty of algorithms, and it seems task-specific. However, it could be a general trick for multi-agent environments with attrition, and it can also be applied to scenarios like MOBA or RTS games. For example, MAPPO [7] also uses such death masking for better value estimation.

In addition, developing such tricks might be a promising direction in generative models. For example, the outputs of large models are not always reliable, such as generating movements that do not satisfy physical constraints. Introducing human priors into scene prediction results could be an essential way to solve these issues. Methods like PINNs and Lagrangian conservation can impose constraints on certain predictions, thereby enhancing the reliability of the generated results. The tricks specific to the SMAC game in this paper can also be categorized as one of these methods.

[1] Tabish Rashid, et al. Qmix: Monotonic value function factorisation for deep multi-agent reinforcement learning.

[2] Jakub Grudzien Kuba, et al. Trust region policy optimisation in multi-agent reinforcement learning.

[3] Micheli, Vincent, Eloi Alonso, and François Fleuret. Transformers are Sample-Efficient World Models.

[4] Egorov, Vladimir, and Alexei Shpilman. Scalable Multi-Agent Model-Based Reinforcement Learning.

[5] Zhang, Yang, et al. Decentralized Transformers with Centralized Aggregation are Sample-Efficient Multi-Agent World Models.

[6] Liu, Qihan, et al. Efficient Multi-agent Reinforcement Learning by Planning.

[7] Chao Yu, et al. The surprising effectiveness of ppo in cooperative multi-agent games.