Mixture-of-World Models: Scaling Multi-Task Reinforcement Learning with Modular Latent Dynamics

TL;DR

MoW addresses the challenge of sample-efficient multi-task reinforcement learning in high-dimensional visual environments by introducing a modular world-model framework that decouples perception and dynamics via task-specific VAEs and a mixture of Transformer experts guided by task embeddings. A gradient-based warmup clustering and specialized losses (task prediction and expert balance) enable effective parameter sharing without collapsing expert usage, while harmonious loss stabilizes multi-task optimization. Empirically, MoW delivers state-of-the-art performance on Meta-World and competitive Atari 100k results with roughly half the parameters of comparable baselines, demonstrating scalable, parameter-efficient generalist world models for visual MTRL. This architecture paves the way for scalable multi-task agents that can learn rich latent dynamics across diverse tasks with improved reconstruction fidelity and sample efficiency.

Abstract

A fundamental challenge in multi-task reinforcement learning (MTRL) is achieving sample efficiency in visual domains where tasks exhibit substantial heterogeneity in both observations and dynamics. Model-based reinforcement learning offers a promising path to improved sample efficiency through world models, but standard monolithic architectures struggle to capture diverse task dynamics, resulting in poor reconstruction and prediction accuracy. We introduce Mixture-of-World Models (MoW), a scalable architecture that combines modular variational autoencoders for task-adaptive visual compression, a hybrid Transformer-based dynamics model with task-conditioned experts and a shared backbone, and a gradient-based task clustering strategy for efficient parameter allocation. On the Atari 100k benchmark, a single MoW agent trained once on 26 Atari games achieves a mean human-normalized score of 110.4%, competitive with the score of 114.2% achieved by STORM, an ensemble of 26 task-specific models, while using 50% fewer parameters. On Meta-World, MoW achieves a 74.5% average success rate within 300 thousand environment steps, establishing a new state of the art. These results demonstrate that MoW provides a scalable and parameter-efficient foundation for generalist world models.

Figures (15)

• Figure 1: Overview of the Mixture-of-World models (MoW) architecture. Task-specific observations are encoded through specialized VAEs, with dynamics modeled by a mixture-of-Transformer experts routed via task embeddings. The design enables modular latent dynamics handling while maintaining parameter efficiency. Here, at time $t$ for task $k$, $o_{k}^t$, $r_{k}^t$, $c_{k}^t$ and $a_{k}^t$ denote the high-dimensional observation, reward, termination flag, and action, respectively. The stochastic representation $z_{k}^t$ is sampled from the distribution $\mathcal{Z}_{k}^{t}$, which is encoded from the observation $o_{k}^t$. The hidden state $h_{k}^t$ is learned by the hybrid Transformer architecture, while $e_{k}$ represents the learnable task embedding and $N_e$ is the number of experts.
• Figure 2: The MoW enables precise modeling and reconstruction of task-specific dynamics through task clustering and expert routing. We present the reconstructed images decoded by the final imagination states of the $26$ tasks in Atari $100$k, where the horizon is set to $16$ steps. Compared to the vanilla transformer employed in vanilla STORM framework, the MoW demonstrates improved task discrimination and more accurate imagination in multi-task settings.
• Figure 3: Results of MoW on the Atari $100$k benchmark (left) and the Meta-world benchmark (right). Compared to baseline STORM, MoW results a $50\%$ reduction in model size (middle).
• Figure 4: By increasing the number of experts (blue) and clusters of VAEs (yellow), MoW effectively unlocks the parameter scalability.
• Figure 5: Ablation studies on muti-task STORM.