NeurIPS 2025

M3W: Learning and Planning Multiple Tasks via an MoE-based World Model

M3W · MoE-based Multi-Task World Model

Zijie Zhao^1,2, Zhongyue Zhao², Kaixuan Xu², Yuqian Fu^1,2, Jiajun Chai³, Yuanheng Zhu^2,1,†, Dongbin Zhao^2,1,†

¹ School of Artificial Intelligence, University of Chinese Academy of Sciences

² Institute of Automation, Chinese Academy of Sciences ³ Meituan

^† Corresponding authors.

Code PDF OpenReview

Abstract

Learn shared task dynamics. Plan task-specific actions.

Multi-task reinforcement learning aims to develop a single model capable of solving a diverse set of tasks. Existing methods often struggle with substantial variation in optimal policies across tasks. In contrast, many tasks exhibit bounded similarity in their underlying dynamics: high similarity within related task groups, but substantial differences between unrelated groups. We present M3W, which applies mixture-of-experts to a world model for both learning and planning. A SoftMoE dynamics model and a SparseMoE reward predictor support knowledge reuse across similar tasks while reducing conflicts across dissimilar tasks. An MPPI planner evaluates and optimizes actions through predicted rollouts. Experiments on 14 Bi-DexHands tasks and 24 MuJoCo tasks demonstrate improved performance, sample efficiency, and multi-task adaptability over policy-centric baselines.

M3W framework connecting multiple tasks, an MoE world model, and planning — **M3W at a glance.** A modular world model captures task dynamics and rewards, supporting planning across diverse tasks.

Method

One world model, multiple tasks.

M3W combines modular prediction with model predictive path integral (MPPI) planning.

Learn dynamics
SoftMoE models latent dynamics, reusing knowledge within groups of related tasks.
Predict rewards
SparseMoE predicts task rewards with selective expert activation to reduce interference.
Plan together
An MPPI planner optimizes joint actions using imagined world-model rollouts.

Detailed M3W architecture with SoftMoE dynamics, SparseMoE reward prediction, and planning — **Learning and planning.** Task-conditioned predictions support action selection across tasks with different dynamics and objectives.

Results

Multi-task performance across manipulation and locomotion.

M3W is evaluated on Bi-DexHands and MuJoCo against policy-centric baselines.

Performance comparison between M3W and reinforcement learning baselines on Bi-DexHands and MuJoCo — **Multi-task performance.** Learning curves for representative tasks. See the paper for metrics, evaluation protocols, and full experimental results.

Task demonstrations

Behaviors learned by M3W.

Representative rollouts from dexterous manipulation and mujoco locomotion tasks.

Bi-DexHands

Catch Abreast task rollout — Catch Abreast

Catch Underarm task rollout — Catch Underarm

Lift Underarm task rollout — Lift Underarm

Open Outward task rollout — Open Outward

MuJoCo

Cheetah Run Front task rollout — Cheetah Run Front

Cheetah Run Backward task rollout — Cheetah Run Backward

Walker Stand task rollout — Walker Stand

Citation

BibTeX

@inproceedings{zhao2025m3w,
  title     = {Learning and Planning Multi-Agent Tasks via an 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 = {Advances in Neural Information Processing Systems},
  year      = {2025},
  url       = {https://openreview.net/forum?id=fi24ry0BX5}
}