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

arXiv: https://arxiv.org/abs/2602.01270 · HTML: https://mlsanity.com/paper/2602.01270/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Latent dynamics models and model-based RL, one of 11 directions in World models.
- Builds on 36 of its 47 references that are in the map.
- First posted on arXiv 1 Feb 2026.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)
- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 30 of 36)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics](https://mlsanity.com/paper/1705.07115/) (2017)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks](https://mlsanity.com/paper/1711.02257/) (2017)
- [TD-MPC2: Scalable, Robust World Models for Continuous Control](https://mlsanity.com/paper/2310.16828/) (2023)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Transformers are Sample-Efficient World Models](https://mlsanity.com/paper/2209.00588/) (2022)
- [Diffusion for World Modeling: Visual Details Matter in Atari](https://mlsanity.com/paper/2405.12399/) (2024)
- [Gate-Variants of Gated Recurrent Unit (GRU) Neural Networks](https://mlsanity.com/paper/1701.05923/) (2017)
- [Multi-Game Decision Transformers](https://mlsanity.com/paper/2205.15241/) (2022)
- [Multi-Task Reinforcement Learning with Context-based Representations](https://mlsanity.com/paper/2102.06177/) (2021)
- [Masked World Models for Visual Control](https://mlsanity.com/paper/2206.14244/) (2022)
- [Mastering Complex Control in MOBA Games with Deep Reinforcement Learning](https://mlsanity.com/paper/1912.09729/) (2019)
- [Multi-Task Reinforcement Learning with Soft Modularization](https://mlsanity.com/paper/2003.13661/) (2020)
- [TransDreamer: Reinforcement Learning with Transformer World Models](https://mlsanity.com/paper/2202.09481/) (2022)
- [Transformer-based World Models Are Happy With 100k Interactions](https://mlsanity.com/paper/2303.07109/) (2023)
- [STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning](https://mlsanity.com/paper/2310.09615/) (2023)
- [Multi-Task Reinforcement Learning with Mixture of Orthogonal Experts](https://mlsanity.com/paper/2311.11385/) (2023)
- [Efficient World Models with Context-Aware Tokenization](https://mlsanity.com/paper/2406.19320/) (2024)

## Built on by (most cited first, 1 of 1)

- [MA-WAM: Multi-Agent World-Action Model for Test-Time Planning](https://mlsanity.com/paper/2609.31281/) (2026)

## Abstract (arXiv)

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.