The official implementation of the paper Learning and Planning Multi-Agent Tasks via an MoE-based World Model, which published on NeurIPS 2025
M3W is a Mixture-of-Experts world model framework for multi-task multi-agent reinforcement learning.
M3W leverages the idea of bounded similarity in task dynamics, combining SoftMoE dynamics learning and SparseMoE reward prediction within a world model. By planning directly on predicted rollouts with a multi-agent MPPI planner, it enables efficient knowledge reuse, conflict avoidance, and scalable multi-task adaptability, surpassing policy-centric approaches.
conda create -n m3w python=3.8 -y
conda activate m3w
pip install -r requirements.txtIf you encounter issues when installing MA-MuJoCo, Isaac Gym or Bi-DexHands, please refer to MA-MuJoCo, Isaac Gym and Bi-DexHands.
For training, please run:
python examples/train.py --load_config configs/mujoco/installtest/m3w/config.jsonYou can modify the configuration in configs/mujoco/installtest/m3w/config.json to customize the training process.
We extend Bi-DexHands to multi-task settings using multiprocessing, which makes the sampling process significantly slower. To accelerate training, you can refer to the settings in shadow_hand_meta settings and enable GPU-based parallel multi-task execution.
M3W achieves superior performance on both Bi-DexHands and MA-MuJoCo, showing superior sample efficiency and multi-task adaptability compared to policy-centric baselines.
Here we showcase several representative task rollouts demonstrating the diverse cooperative behaviors learned by M3W.
Our code is built upon HARL, TDMPC2 and Light-SoftMoE. We thank all these authors for their nicely open sourced code and their great contributions to the community.
If you find our research helpful and would like to reference it in your work, please consider the following citations:
@inproceedings{zhao2025m3w,
title = {Learning and Planning Multi-Agent Tasks via a MoE-based World Model},
author = {Zhao, Zijie and Zhao, Zhongyue and Xu, Kaixuan and Fu, Yuqian and Chai, Jiajun and Zhu, Yuanheng and Zhao, Dongbin},
booktitle = {},
year = {2025}
}