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

TL;DR

The paper introduces Mixture-of-World Models, a scalable, modular architecture for multi-task reinforcement learning that improves sample efficiency and achieves state-of-the-art results on Atari and Meta-World benchmarks.

Contribution

It presents a novel modular architecture combining variational autoencoders, Transformer-based dynamics models, and task clustering for efficient multi-task RL.

Findings

Achieves 110.4% human-normalized score on Atari 100k with fewer parameters.

Sets a new state-of-the-art success rate of 74.5% on Meta-World.

Uses 50% fewer parameters than ensemble models while maintaining performance.

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.