Electronics, Vol. 14, Pages 4884: MoE-World: A Mixture-of-Experts Architecture for Multi-Task World Models

Electronics doi: 10.3390/electronics14244884

Authors:
Cong Tang
Yuang Liu
Yueling Wu
Wence Han
Qian Yin
Xin Zheng
Wenyi Zeng
Qiuli Zhang

World models are currently a mainstream approach in model-based deep reinforcement learning. Given the widespread use of Transformers in sequence modeling, they have provided substantial support for world models. However, world models often face the challenge of the seesaw phenomenon during training, as predicting transitions, rewards, and terminations is fundamentally a form of multi-task learning. To address this issue, we propose a Mixture-of-Experts-based world model (MoE-World), a novel architecture designed for multi-task learning in world models. The framework integrates Transformer blocks organized as mixture-of-experts (MoE) layers, with gating mechanisms implemented using multilayer perceptrons. Experiments on standard benchmarks demonstrate that it can significantly mitigate the seesaw phenomenon and achieve competitive performance on the world model’s reward metrics. Further analysis confirms that the proposed architecture enhances both the accuracy and efficiency of multi-task learning.

Source link

Cong Tang www.mdpi.com