Mixture-of-World Models: Scaling Multi-Task Reinforcement Learning with Modular Latent Dynamics
Authors: Boxuan Zhang
Weipu Zhang
Zhaohan Feng
Wei Xiao
Jian Sun
Jie Chen
Gang Wang
Venue: ICLR 2026 Poster (2026)
Keywords: Multi-task reinforcement learning
world model
transformer
mixture-of-world models
Primary area: reinforcement learning
TL;DR: This paper proposes Mixture-of-World models (MoW), a novel and sample-efficient world model architecture for multi-task reinforcement learning.

Abstract: A fundamental challenge in multi-task reinforcement learning (MTRL) is achieving sample efficiency in visual domains where tasks exhibit significant heterogeneity in both observations and dynamics. Model-based RL (MBRL) offers a promising path to sample efficiency through world models, but standard monolithic architectures struggle to capture diverse task dynamics, leading to poor reconstruction and prediction accuracy. We introduce the mixture-of-world models (MoW), a scalable architecture that integrates three key components: i) modular VAEs for task-adaptive visual compression, ii) a hybrid Transformer-based dynamics model combining task-conditioned experts with a shared backbone, and, iii) a gradient-based task clustering strategy for efficient parameter allocation. On the Atari 100k benchmark,  **a single MoW agent** (trained once over Atari $26$ games) achieves a mean human-normalized score of $\mathbf{110.4}$%, competitive with the  $\mathbf{114.2}$% achieved by the recent STORM-an ensemble of $26$ task-specific models-while requiring $\mathbf{50}$% fewer parameters. On Meta-World, MoW attains a $\mathbf{74.5}$% average success rate within $300$K steps, establishing a new state-of-the-art. These results demonstrate that MoW provides a scalable and parameter-efficient foundation for generalist world models.

Published as a conference paper at ICLR 2026
MIXTURE-OF-WORLI MODELS:
SCALING MULTI-TASK REINFORCEMENT LEARNING WITH MODULAR LATENT DYNAMICS
Boxuan Zhang' Weipu Zhang Zhaohan Feng' Wei Xiao Jian Sun Jie Chen" Gang Wang" National Key Lab of Autonomous Intelligent Unmanned Systems, Beijing Institute of Technology "Department of Control Science and Engineering, Harbin Institute of Technology "Jiangxing Intelligence Inc.
(zhangboxuan, zhfeng, xiaowei, sun junjian, chenjie, gangwang) edu. cn we we seipuagang academic.@gmall com com
ABSTRACT
A fundamental challenge in multi-task reinforcement learning (MTRL) is achiev- ing sample efficiency in visual domains where tasks exhibit significant hetero- geneity in both observations and dynamics. Model-based RL (MBRL) offers a promising path to sample efficiency through world models, but standard mono- lithic architectures struggle to capture diverse task dynamics, leading to poor re- construction and prediction accuracy. We introduce the mixture-of world models (MoW), scalable architecture that integrates three key components: i) modu- lar VAEs for task-adaptive visual compression, ii) a hybrid Transformer- based dynamics model combining task-conditioned experts with shared backbone, and, iii) a gradient-
task clustering strategy for efficient parameter alloca- tion. On the Atari 100K benchmark, single MoW agent (trained once over Atari 26 games) achieves mean human-normalized score of 110.4%, competi- tive with the score 114.2% achieved by the recent STORM an ensemble of 26 task-specific models while requiring 50% fewer parameters. On Meta- World, MoW attains a 74.5% average success rate within 300K steps, establishing a new state-of-the-art. These results demonstrate that MoW provides scalable and parfmeter-effficiett foundation for generalist world models.
INTRODUCTION
The pursuit of sample-efficient reinforcement learning (RL) has led to significant progress in com- plex domains, from game-playing (Silver et al., 2016 Hansen et al. 2024) to robotic control (Ye) et al. 2020). However, these advances often confined to single-task settings, where the learned policies lack generalization (Yu et al.] [2020a] Feng et al. 2026). Multi-task RL (MTRL) aims to overcome this by learning unified policy across diverse task suite (Yang et al. 2020). While model-free methods have been extended to MTRL (Xing et al. 2025). they inherit the poor sample efficiency of their single-task counterparts (Mnih et al.55 Schulman al. (017). limiting their applicability in real-world scenarios where data is scarce (Xia et 2018 Huang et al. 2025).
In single-task settings, model-based RL (MBRL) has emerged as a path achieving greater sample efficiency. World models learn compact latent representations of environment dynamics, enabling policy optimization entirely through "latent imagination" (Ha & Schmidhuber 2018 Hafner al.)
2019 2025). These methods typically combine variational autoencoder (VAE) with latent recur- rent state-space model (RSSM) using e.g., a recurrent neural network (RNN) (Hafner et al. 2019)
or Transformer (Vaswani 2017). and have shown remarkable success (Zhang et al. 2025
Hafner
jet al.
2025 Micheli et al. 2022 Zhang et al. 2023).
Scaling world models to multi-task, visual domains presents a fundamental challenge: the inherent visual and dynamic heterogeneity across tasks. single world model must simultaneously capture
Cooresponning author

---

Published as a conference paper at ICLR 2026
f
2++1
14.
Shared former
Transformer Blocks
Agent
Agent
Router
Expert
Expert 2
Expert Ne-1
E Export Nr
Enci
Dect
Figure 1: Overview of the cekture-offworlde models (MoW) architecture. Task-specific observations are encoded through specialized VAEs, with dynamics modeled by mixture-off Transformer routed via task embeddings. The design enables modular latent dynamics handling while maintain- ing parameter efficiency. Here, at time for task k, o r o and oj denote the high-dimensional observation, reward, termination flag, and action, respectively. The stochastic representation zf is sampled from the distribution Z which is encoded from the observation of. The hidden state hy learned by the hybrid Transformer architecture, while eg represents the learnable task embedding and No is the number of experts.
-frxturrof-Trrnsffomme experts
disparate visual features while accurately predicting task-specific Dynamics-objectivee that often conflict in practice, straining the capacity of monolithic architectures (Yu et al. (0220b) Simply conditioning shared model on task identity often proves insufficient without abundant expert data (Hansen al., 2024). Crucially, for visual MBRL,, the model must achieve high-fidelity recon- struction to enable effective policy learning in the latent space challenge that remains largely unaddressed MTRL.
In this work, we propose the mixture-of-world models (MoW), novel architecture designed for parameter- and sample-efficient MTRL in visual domains. The MoW explicitly addresses multi- task heterogeneity through modular design. It employs a set of task-specific VAEs for visual encoding and a opatial-temporaa dynamics core composed guided by learned task embeddings. To ensure effective module specialization, we introduce task prediction loss that enhances task discrimination the latent space and an expert balance loss to prevent module collapse. Furthermore, develop a gradient-based task clustering strategy during a warm-up phase to share components across tasks, significantly improving parameter efficiency.
a mixture expert Transformers,
We evaluate MoW and Meta- World (continuous control). Our results demonstrate that a single MoW agent achieves suman-normalizer score performance competitive with an ensemble of 26 task-speeific models on Atari while using half the parameters, and 74.5% success rate on Meta- World within 300K steps, setting new ofate-of-the-art. Ablation studies confirm the contribution of each component.
challenging benchmarks: the Atari 100K benchmark (discrete control)
Our main contributions are summarized below.
The MoW architecture: A mixture-of of world models with modular VAEs and transformer experts that captures diverse task dynamics while maintaining reconstruction fidelity.
Specialized losses: Task prediction and expert balance losses that enable effective task diserimination and balanced utilization across expert components.
Gradient- clustering: A warm-up strategy that optimizes parameter sharing by clus- tering tasks based gradient similarity, enhancing model efficiency.
State-of-the--ar performance: Empirical validation on Atari 100K and WorldWorld demonstrating superior sample efficiency and parameter scaling.
2

---

Published as a conference paper at ICLR 2026
2 RELATED WORK
2.1 MULTI-TASK REINFORCEMENT LEARNING
Multi-task reinforcement learning seeks to develop agents capable of generalizing across diverse set of tasks by leveraging shared experience and representations. A dominant paradigm in online MTRL employs task-conditioned policies. where a unified network or value function is conditioned on task identifiers or embeddings (Xu al., 2020 Sodhani al. al.]|]20000 2021b). This approach has success- fully adapted celebrated model-free algorithms like PPO (Schulman et al. 2017) and SAC (Haarnoja et al. 2018) to multi-task settings, though they inherit the sample inefficieccy inherent to model-free RL methods. Offline MTRL methods, such as those based on decision Transformers (Lee et al.) 2022 Kong et al. 2025), learn from fixed datasets of expert demonstrations. While effective within their data constraints, these approaches cannot improve through online interaction and are typically limited to state-based inputs, where errors task conditioning can critically degrade performance when extended to visual inputs.
The mixture-of-expprts (MoE) architecture, widely successful in scaling large language models to trillions of parameters by virtuc of their modular, sparsely activated design (Fedus al. 2022 Lep- ikhin et al. 2020), has recently been applied to MTRL. In such models, only a small subset of experts invoked for each input token, enabling dramatic increases in model capacity without proportional growth in computational cost. Methods like MOORE (Hendawy et al., 2023) and D2R (He et al. 2024 incorporate MoE into SAC, improving knowledge sharing and specialization. However, these approaches remain evaluated primarily on dow-dimensional state spaces and are hampered by the sample inefficiency of model-free learning. The significantly greater challenge of digeedimesi visual inputs with their complex perception and dynamics-remains largely unaddressed by existing model-free MoE methods.
2.2 WORLD MODELS FOR REINFORCEMENT LEARNING
World models have emerged as a powerful approach for improving sample efficiency in RI. by enabling learning from imagined trajectories. Early work like SimPLe (Kaiser al., 2019) demon- strated the potential of learned models for policy learning. The Dreamer series (Hafner et al. 2019 2025 advanced this paradigm with recurrent state-space models that support long-horizon latent imagination. Recently, Transformers have shown strong performance as world models due to their capacity for long-range dependency modeling. IRIS (Micheli al.] 2022) and A-IRIS (Micheli) et al. 2024) employ VQ-VAEs (Van Den Oord et al., 2017) for discrete visual tokenization, while STORM (Zhang et al., (023) achieves state-of-the-art results on Atari with a customized Trans- former architecture. DIAMOND (Alonso et al. 2024) explores diffusion-based world models oper- ating directly in pixel space, improving visual fidelity.
Critically, these advances all focus single-task RL settings. While TD-MPC2 (Hansen et al. 2024) extends model-based principles to multi-task domains, it operates on low-dimensional state inputs rather than more challenging visual observations and relies on expert data. The challenge of building world models that simultaneously handle diverse visual inputs and dynamics across multi- ple tasks-while maintaining reconstruction fidelity for effective latent imagination- open. Our work directly addresses this gap by introducing a modular architecture specifically designed for multi-task visual world models.
3 METHOD
We consider a multi mulli-task setting with K tasks collectively denoted by T = T2. Tx }- Each task T3 partially observable Markov decision process (POMDP) defined by the tuple MR = =OO::A Au, Pi, RR. 7R), where O₂ is the observation space, As. is the action space, Pa is the transition dynamics, R₂ is the reward function, and 7 E [0, the discount factor. The objective is to learn a single policy that maximizes the expected discounted return across all tasks: ΣK B(oL,RE>~T:P D 2 R (of, a.)]-
From the preceding observations and discussions, our key insight is that constructing a mixture of world models across heterogeneous tasks, along with performing imagination and training in paral-
3

---

Published as a conference paper at ICLR 2026
lel, enables the development of more generalizable multi-task agent while significantly possessing the parameter scalability.
3.1 ARCHITECTURE OVERVIEW
The MoW consists of two primary modules: perceptual module that encodes high-dimensional visual observations into a compact latent space, and a temporal module that learns aask-adaeiive dynamics within this latent space. The overall architecture is depicted in Figure [
3.1.1 PERCEPTUAL MODULE
The perceptual module encodes each image observation of into stochastic representation zi using categorical VAE. Each task Tj is assigned specific encoder-decoder pair (46.4. ), where the index 0 E [1, Nm] is determined by gradient- based clustering procedure detailed in Sec- tion To avoid an excessive number of hyperparameters. N' is not only the number of VAEs, but also the number of distribution, reward, continuation predictors and critic networks mentioned later (all components with the subscript is as the predictors and VAEs are corresponding in task k). A learnable task embedding eı (with ||ex|12 = 1) conditions the encoding and decoding pro- cesses. The stochastic representation = is sampled from the categorical posterior distribution ž (comprises 32 categories, each with 32 classes) to serve observation with the straight-through gradient trick (Bengio et al. 2013):
compact representation of the original
Posterior:
2 I0.in (2k o cx) = 2
(1)
Reconstruction:
o% -pain (ok ex).
3.1.2 TEMPORAL MODULE
The temporal module models the dynamics in the latent space, which takes the sequence of the stochastic representations and actions as input, and predicts the next latent state, reward, task index, and continuation flag. The module is composed of mixture of expert Transformers and shared Transformer, while the variables are also conditioned by the task embedding Ck.
Input Representation. The stochastic representation 2 and action af are first concatenated via a multi-layer perceptron (MIP) mæ.j to a unified token m for each expert j. where = 1,2, Ne index the activated expert Transformers.
1:f j.kk = ma
1: 1:t ,0 ,x)-
(2)
Expert Selection. We employ a task-level router function that takes the task embedding en as the input, and produces set of expert indices Js. with their corresponding weights We via a TopK operation applied to softmax output of MLP.
S Softmax (MLP(e+)).
(3)
WR. Jk TopK (Sx, topk =
m4); =
where Sş denotes the ex-pert affinity scores, TopK selects the ni largest scores, Js = 12. .Jk Th denote the activated expert Transformer indices, and Wx [0, "0....,t.. w -..... wx ,...,,.0.00....... is the normalized weight vector of activated experts, where the position of wv indicates the activation of the corresponding expert. The selection of in is deter- mined during the warmup stage by clustering the gradient vectors and is fixed throughout training. In contrast, the expert indices Jy are dynamically determined by a TopK function conditioned on the current task embedding, and changes over time as the world model is trained in an end-to-end manner but remains fixed in a single training step.
Expert and Shared Transformers. The tokens are processed by the activated expert Transformers fe. for j E Jg. The outputs of the experts are concatenated to form token 11:t which is then passed to the shared Transformer Fg to produce the hidden state ht The expert Transformers specialize in modeling task-specifes dynamics, while the shared Transformer captures common knowledge across the entire domain.
11 concat Joatt "iiie...... S. (m 2
(4)
h 1:t (((t11(.)))
4

---

Published as a conference paper at ICLR 2026
Prediction Heads. The hidden state h is used to predict the next latent state distribution, reward, and continuation flag. Additionally, an auxiliary task predictor head is used to encourage task- discriminative feature extraction:
-1+1 So.se (*k (igse))
1 - So.ie - (hi ek),
(5)
c = go.i ex),
== g (ht).
g
An underlying intuition is that the expert selection should remain relatively stochastic in the early stages of training to facilitate balanced expert utilization. As training progresses and converges, the selection should gradually become more deterministic; otherwise, frequent switching of experts may destabilize the end-to-end optimization. To embody this intuition, we anneal the temperature coefficient in the softmax function progressively toward during training to avoid excessively low temperatures that result in nearly one-hot distributions, which would hinder experts' collaboration.
3.2 DESIGN RATIONALE FOR TASK-LEVEL ROUTING
We adopt task-level routing instead of token-level routing, in other words, we use task embeddings eg rather than mixture tokens m to select experts, based on the following considerations. First, in conventional Transformer architectures, MoE is typically applied to the rather than the self-attention layers. This design implies that the MoE module primarily models sparse transformations between data features, rather than capturing the temporal dependencies that self-attention layers focus on. To address this, we decouple the MoE mechanism from the Trans- former architecture and apply it externally, enabling MoW to learn task-specific dynamics across different tasks through the mixture of transformer backbones.
networks
Second, if expert Transformers are selected based on the mixed token, each token along the training trajectory may activate different experts, even within the same task. Such dynamic routing results in each expert observing only partial sequences, leading to the fragmented learning of task dynamics. As a consequence, the experts may fail to effectively model the full temporal structure of the task. In contrast, routing based task embeddings ensures consistent expert activation within each task, allowing individual experts to capture more complete and coherent task dynamics.
3.3 WORLD MODEL LEARNING
We train the MoW in an end-to-end manner using self-supervised learning paradigm, leveraging data from all tasks jointly. For task Tk, the total loss is given by
(o) Σ (0) Crew (Ø) + 27. (d) + Ctack (Ø) +
B1 B (0) 32
(6)
with hyperparameters B1 0.5 and B3 0.1 fixed, where f denotes the reconstruction loss of the original image, 2 (Ø) the prediction loss of the reward, 2 (Ø) the prediction loss of the continuation flag, and Llack (r) represents the prediction loss of task Te:
Li
1o% 122
Crew (Ø) Symlog rk),
( = ChoaryCrosEEn((,... oe),, Clack (Ø) = CrossEnt(k;, k).
Introducing auxiliary predictors and their associated losses has proven effective for shaping the representation of latent states (Burchi & Timofte (025). Therefore, we incorporate an additional task prediction loss Ct (Ø) to explicitly guide the training of hidden states hf. toward being task- discriminative. This ensures that the hidden states contain sufficient information for task identifica- tion, thereby enhancing the accuracy of the imagination process.
The losses Cdvu (Ø) and CTEP (Ø) train the dynamics predictor network to predict the next stochastic distributions from the hidden states by minimizing the Kullback-Leibler (KL) divergence between and the next encoder representation 2
the predictor output distribution 2 92.. ( ((t)) go lof. ex):
cly (d) max KL
sg 2 >112+(1)))
(8)
CTP (Ø) = max (1,
KL, KL [2 sg( (g(22111 I]),
5

---

Published as a conference paper at ICLR 2026
imagined by multi-task STORM
Imagined by MoW
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 imag- ination states of the tasks Atari 100K, 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 multi-task settings.
Figure 2:
where the maximum operator is used to clip the loss below the value of nat 1.44 bits (Chen et al. 2022) and sg(-) is the stop-gradient operator.
To effectively alleviate the problem of gradient conflicts in multi-task learning and mitigate the scale disparities among loss terms, we employ harmonious loss weights to train the MoW.
3.4
HARMONIOUS AND BALANCED LOSS
The central challenges in MTRL lie in two key aspects: i) The first is the presence of gradient conflicts across tasks, which drive the shared parameters in divergent directions. This issue is further exacerbated in the end-to-end training pipelines involving multiple components, often resulting in degraded representations and unstable training dynamics (Chen et al. 2018). While prior works have introduced techniques such gradient projection and conflict-aware optimization to mitigate this problem, these approaches typically incur additional model complexity. ii) The second challenge concerns the sensitivity of performance to the relative weighting of task-specific losses. Manually tuning these weights is both non-trivial and resource-intensive,, thereby limiting the practicality of MTRL in real-world scenarios (Kendall al. 2018).
To address these challenges, the MoW is trained end-to-end by incorporating the harmonious weights (Ma et al. (2024), which is applied on top of the individual task-specific losses (defined in Equation 19. This loss dynamically adjusts task weights to mitigate potential gradient conflicts during training, defined as follows
K Σ k=1 o
Cx(()=
-ー
Lx(d) + In(1 + tO))
(9)
where Cx(()) is the harmonized loss and the rectified scale is equal to in Appendix A.7 Finally, the balance loss is added to the rectified loss
< 1 (derived
1 1+//44///2.)
((0) (x( + O.1
(10)
k= We KN 1N. where WR is the normalized weight vector and 1x. N. is vector consisting of ones. Since the MoW does not explicitly apply experts' weights but in- stead combine expert features via concatenation, the balance loss primarily serves encourage the activation of all experts, rather than enforcing identical expert selection distributions across tasks.
with (bal
3.5 AGENT LEARNING
The agent learns entirely through the imagination trajectories enabled by the world model, making accurate reconstruction of the raw observations a critical factor for training efficient agents. As shown in Fig. compared to the multi-task vanilla STORM, the MoW produces imagined trajec- tories that reconstruct the raw images more faithfully while avoiding confusion across tasks. We provide detailed reconstruction results of MoW across different tasks in Appendix A.6

---

Published as a conference paper at ICLR 2026
To initiate the imagination process, short context trajectory is randomly sampled from the replay buffer to compute the initial posterior distribution Z During inference, rather than drawing sam- ples from the posterior, the latent representation z is sampled from the prior distribution 2g- To accelerate inference, we incorporate the key-value (KV) caching mechanism within the Transformer architecture (Chen) 2022), which allows for faster autoregressive rollout by reusing previously com- puted attention keys and values. The agent's state is formed by concatenating z hi. and ex.
State:
o = [= 1 exl,
Critic:
Visio (8%) 2 ETTO A Σ Tirk
(11)
T=0
Actor: at (a (x))-
Other core training configurations follow the settings established in the STORM framework, for all tasks, the agent shares the actor network, while the number of critic networks matches the number of VAEs N.- For task Tk. the corresponding critic network index is also determined through clustering the gradient vector during the warm-up stage. The complete loss of the actor-critic algorithm described in Equation 12 and Equation| [4| where ff corresponds to the predicted reward, and o represents the predicted continuation flag. The critic networks learn to approximate the distribution of the X-return estimates R2 using the maximum likelihood loss
(1b)
Σ b,t
Caym (Voisa (st) sg( (R))) + Laym (Veise (s sg(V MURMA iu (sf)))
(12)
where Laym represents the symlog two-hot loss and VORMA is the exponential moving average of the critic network vb for stabilizing training and preventing overfitting. The A-return R2 (Sutton & Hafner et al. (025) is recursively defined as follows
Barto 2018
RA, ri + Ach (1 V) (sttt) + AR
(13)
R = (s
The actor network is trained using the surrogate loss function
R (s,
(())
Σ
sg
In mo(af(s)) 1 - (
(14)
ma
R.t
where (-) denotes the entropy of the policy distribution, and the constant 17 represents the coeffi- cient for entropy loss. The normalization ratio S is defined in Equation 15
computed as the range
between the 95 and 5 percentiles of the A-return R,A,A across batch
S = percentile 95) Rercenttile(RR,e. 5).
(15)
3.6 WARMUPING MoW
World model-based RL methods typically perform a number of interactions with the environment using random actions prior to the world model training stage, populating the replay buffer to enable a stable initialization phase. During the warmup stage MoW, after collecting sufficient amount of random interaction data, we fix the replay buffer and refrain from updating the agent. Instead, we perform number of world model training steps to initialize the dynamics representation. In this process, the expert Transformer modules are trained following the same end-to-end manner as previously described in Equation] TO However, only a single set of VAE and its associated predictors participate in the warmup training stage. Upon completion of the warm-up stage, the parameters of these components are copied to the remaining modules with identical architectures, after which the standard iterative updates of the world model and agent, as classic world model-based RI. After the warmup stage, we extract the gradient vector associated with each task. These task-specific gradient vectors are then clustered to determine which tasks should share the i-th set of VAE, predictor and critic network. This task grouping strategy based on the gradient correspondence is similar to Ha (Hu et al. 2024). In practice, due to the large number of trainable parameters in our mixture world model. reduce the computational cost by averaging the gradient vectors from each neural network layer. For a more comprehensive understanding of training the mixture world model, provide the pseudo-code implementation in Appendix A.12
7

---

Published as a conference paper at ICLR 2026
Atari 100k
Model Size
Meta- -ert--Wrrid tasks
tasks Atari 100k)
Soors
Model 4020 (1465
2032.5
$28 2
332
143.5
Ensfronment
Agent
World model
Environment
Mucture-oFWorrd Models
Singeetak STORM
Figure 3: Results of MoW on the Atari 100K benchmark (left) and the Meta-world benchmark (right). Compared to baseline STORM, MoW results a 50% reduction in model size (middle).
Table 1: Results on MetaWorld MT50
Algorithms
Input
Success Rate (Episode Length)
Total Environment Steps
MTSAC
49.3 + 1.5 (150)
CARE
State
50.8 + 1.0 (150)
20M
PaCo
57.3 + 1.3 (150)
MOORE
State
72.9 + 3.3 (150)
100M
MoW (ours)
Image
74.5 + 1.1 (500)
15M
4
EXPERIMENTS
We evaluate the effectiveness of MoW in improving the sample efficiency of baseline MBRL meth- ods across two representative benchmarks: Meta-world with robotic manipulation tasks (continuous action spaces) and Atari 100K with video games with discrete action spaces. Experimental results are shown in Figure 3 and demonstrate that our architecture not only substantially enhances sample efficiency but also achieves competitive performance with significantly fewer parameters compared to baselines. In our experiments, we use a machine with NVIDIA 4090 graphics cards with CPU cores and 24 GB RAM. We use distributed data parallel (DDP) in Pytorch to enable multi-GPU parallel multi-task training, and training MoW on each Atari game for 100K steps took roughly 3.5 hours. Hyperparameters can be found in Appendix A.9
4.1 ATARI 100K EXPERIMENTS
The Atari 100K benchmark is widely adopted testbed for RL agents. To determine the performance of human player H, the player is allowed to become familiar with the game under the same sample constraint. The MoW achieves human normalized score of 110.4% using a single unified model. To ensure the faimess and quality of our results, we also reproduced STORM results using the official code and configurations, which is a human normalized score of 114.2%. Furthermore, compared to STORM, our primary baseline, which requires 1,977.5 MB to support 26 tasks, the MoW achieves comparable performance with a significantly more compact model of only 972.5 MB. Here, MB refers to the storage size of the actual model checkpoint (.pth file), corresponding to 50% reduction in model size. Details are provided in Appendix A.10.1
the Atari 100K benchmark
4.2 META- WORLD EXPERIMENTS
Pora-Worldii is a benchmark of 50 robotic manipulation tasks with fine-grained observation details, such small target objects. In all tasks, the episode length is 500 environment steps with no action repeat. this benchmark, performance is primarily measured by the average success rate across all tasks, offering holistic evaluation of the robot system's 's adaptability and competence in varied task settings. We compare MoW with the official results of state-of of-the-art MTRL methods based on state inputs. As stated in Table I although MoW makes decisions solely from the visual observations-a considerably more challenging setting than state-based decision maaisiinnmmkiig... it still achieves a success rate of 74.5% within only 300K steps per task (15M steps in total). Notably, this surpasses the
8

---

Published as a conference paper at ICLR 2026
performance of classic model-free baselines trained for 20M steps, highlighting MoW's remarkable sample efficiency.
4.3 PARAMETER SCALABILITY
Similar to the MoE architecture, MoW exhibits strong scalability a large number of tasks, owing to its parameter-effiient design. Specif- ically, MoW enhances overall performance by increasing the number of expert modules, ef- fectively reducing the task load per expert. As illustrated in Figure 4 MoW achieves compet-
Parameter scalability on Atari 100k
(12)
Human normalized mean
Human Score
(9)
0.0
(6)
itive results across 26 tasks in the Atari 100K benchmark with only 12 experts. Within sonable range, scaling up the model size leads substantial performance gains. Notably, increasing the number of expert Transform- ers yields more pronounced improvements, as these modules directly facilitate task-specific dynamics modeling. While adding additional VAEs also influences performance by improv- ing reconstruction precision, their impact is comparatively modest when the number of ex- pert Transformers is held constant.
0.2
Scalability
Scalability of experts
398
541
685 Model size (MB)
770
829
Figure 4: By increasing the number experts (blue) and clusters of VAEs (yellow), MoW ef- fectively unlocks the parameter scalability.
5 ABLATION STUDIES
To further investigate the advantages of MoW, we first extended the vanilla STORM frame- work a multi-task setting by incorporating task embeddings. The resulting model failed train effectively, indicating that the baseline is not readily adaptable to multi-task scenar- ios. To ensure this limitation was not due to the VAE, we introduced multiple VAES into the vanilla STORM baseline. The results show that the vanilla Transformer configuration did not achieve performance comparable to our mix-
Gamer Mean
Mixture World Models
Human normalized mean
1.0
Multi task STORM
Human Score
Multi-tsk STORM (12 VAES)
0.6
ture architecture. As result, the agent unable to learn effectively within the latent space. In particular, suffered from signifi- cantly higher dynamics modeling loss, suggest- ing that the imagination process was severely impaired (as shown in Figure
50k Environment
100k
Figure 5: Ablation studies on muti-task STORM.
CONCLUSION
In this work, we introduced a MTRL approach based on the modular mixture-of-world models (MoW) architecture. By combining multiple categorical VAEs, parallel expert Transformers, and shared Transformer, the MoW effectively captures diverse task dynamics while maintaining high sample efficiency. The use of task prediction and expert balance losses further enhances task dis- crimination and expert utilization, respectively, Additionally, the proposed warmup stage promotes efficient module sharing through gradient-based task clustering, substantially reducing model com- plexity. Substantial experiments on Atari 100K and Meta- World benchmarks demonstrate that the MoW not only surpasses prior world model baselines, but also achieves superior performance with significantly fewer parameters. These results underscore the potential of MoW as scalable and efficient framework for general MTRI..

---

Published as a conference paper at ICLR 2026
ACKNOWLEDDMEENTS
This work was supported in part by the National Natural Science Foundation of China under Grants 62495095, 62495090, and U23B2059.
REFERENCES
Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos Storkey, Tim Pearce, and François Fleuret. Diffusion for world modeling: Visual details matter atari. arXiv:2405. 12399, 2024.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating propagating gradients through stochastic neurons for conditional computation. orXiv:1308.3433 2013.
Maxime Burchi and Radu Timofte. Learning transformer-based world models with contrastive pre- dictive coding. arXhv:2503.04116 2025.
Carol
Chen.
Transformer inference arithmetic.
https: 1y/log// tewartihhiii/////mmmmmmmmm 2022. Accessed: 2025-07-17.
Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. Transdreamer: Reinforcement learning with transformer world models. arXiv:2202.094881 2022.
Zhao Chen, Vijay Badrinarayanan, Chen- Yu Lee, and Andrew Rabinovich. Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks. In Proceedings of the 35th International Conference Machine Learning, pp. 799-803, Stockholm, Sweden, 2018. PMLR.
Rahul Dey and Fathi M Salem. Gate-variants of gated recurrent unit (GRU) neural networks. In International Midwest Symposium on Circuits and Systems, pp. 1597-1600. IEEE, 2017.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-39 2022.
Zhaohan Feng. Ruiqi Xue, Lei Yuan, Yang Yu, Ning Ding, Meiqin Liu, Bingzhao Gao, Jian Sun, Xinhu Zheng, and Gang Wang. Multi-agent embodied AI: Advances and future directions. Sci- ence China Information Sciences, 69(5):151202, 2026.
David Ha and Jürgen Schmidhuber. World models. arXi:: 2018.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International Confer- ence on Machine Learning, pp. 1861-1870. PMLR, 2018.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. 1912. 0/00003, 2019.
Danijar Hafner, Jurgis Pasukonis. Jimmy Ba, and Timothy Lillicrap. Mastering diverse control tasks through world models. Nature, 140:647-653, 2025.
Nicklas Hansen, Hao Su, and Xiaolong Wang. Td-mpc2: Scalable, robust world models for contin- uous control. In International Conference on Learning Representations, 2024.
Jinmin He, Kai Li, Yifan Zang, Haobo Fu, Qiang Fu, Junliang Xing, and Jian Cheng. Not all tasks are equally difficult: Multi-task deep reinforcement learning with dynamic depth routing. In Proceedings of the AAAI Conference Artificial Intelligence, volume 38, pp. 12376-12384, 2024.
Ahmed Hendawy, Jan Peters, and Carlo D' D'Eramo. Multi-task reinforcement learning with mixture of orthogonal experts. arXiv preprint 2rXiv:23/1.11333 2023.
Shengchao Hu, Ziqing Fan, Li Shen, Ya Zhang, Yanfeng Wang, and Dacheng Tao. Harmodt: Harmony multi- task decision transformer for offline reinforcement learning. arXiv:2405.1800 2024.
10

---

Published as a conference paper at ICLR 2026
Jincai Huang, Yongjun Xu, Qi Wang. Qi Cheems Wang, Xingxing Liang, Fei Wang, Zhao Zhang, Wei Wei, Boxuan Zhang, Libo Huang, et al. Foundation models and intelligent decision decision-making: Progress, challenges, and perspectives. The Innovation, 2025.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pp. 448-456. pmlr, 2015.
Lukasz Kaiser, Mohammad Babacizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model-based reinforcement learning for atari. arXiv: 1903.0044, 2019.
Alex Kendall, Yarin Gal, and Roberto Cipolla. Multi-task learning using uncertainty to weigh losses for scene geometry and semantics. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 7482-7491, 2018.
Yilun Kong, Guozheng Ma, Qi Zhao, Haoyu Wang, Li Shen, Xueqian Wang, and Dacheng Tao. Mastering massive multi-task reinforcement learning via mixture-of-cxpert decision transformer. arXiv:2505.24378, 2025.
Yann LeCun, Berhhard Boser, John S Denker, Donnie Henderson, Richard Howard, Wayne Hub- bard, and Lawrence D Jackel. Backpropagation applied to handwritten zip code recognition. Neural computation, 1(4):541-551, 1989.
Kuang-Huei Lee, Ofir Nachum, Mengjiao Sherry Yang, Lisa Lee, Daniel Freeman, Sergio Guadar- rama, Ian Fischer, Winnie Xu, Eric Jang, Henryk Michalewski, et al. Multi-game decision trans- formers. In Advances in neural information processing systems, volume 35, pp. 27921-27936, 2022.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu. Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv:2006. 16668, 2020.
Haoyu Ma, Jialong Wu, Ningya Feng, Chenjun Xiao, Dong Li, Jianye Hao, Jianmin Wang, and Mingsheng Long. Harmonydream: Task harmonization inside world models. In International Conference on Machine Learning, pp. 33983-34007. PMLR, 2024.
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world mod- els. arXiv:2209.0055 2022.
Vincent Micheli, Eloi Alonso, and François Fleuret. Efficient world models with context-aware tokenization. arXiv:2406. /9320, 2024.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature, 51 (177554):559-5333 2015.
Jan Robine, Marc Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-based world mod- els are happy with 100k interactions. arXiv:2303.07100 2023.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov, Proximal policy optimization algorithms.
2017.
Younggyo Seo, Danijar Hafner, Hao Liu, Fangchen Liu, Stephen James, Kimin Lee, and Pieter Abbeel. Masked world models for visual control. In Conference on Robot Learning, pp. 1332- 1344. PMLR, 2023.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panncershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. Nature, 529(7587):484
2016. Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-tast reinforcement learning with context- based representations. In International conference on machine learning, pp. 9767-9779. PMLR, 2021a.
11

---

Published as a conference paper at ICLR 2026
Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-task reinforcement learning with context- based representations. In International Conference on Machine Learning, pp. 9767-9779. PMLR, 2021b.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Aaron Van Den Oord, Oriol Vinyals, al. Neural discrete representation learning. Advances in Neural Information Processing Systems, volume 30, pp. 6309 6318. Curran Associates, Inc., 2017.
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, pp. 1998-6008 Curran Associates, Inc., 2017.
Fei Xia, Amir R Zamir, Zhiyang He, Alexander Sax, Jitendra Malik, and Silvio Savarese. Gibson env: Real-world perception for embodied agents. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 9068-9079, 2018.
Jiaxu Xing, Ismail Geles, Yunlong Song, Elie Aljalbout, and Davide Scaramuzza. Multi-task re- inforcement learning for quadrotors. IEEE Robolics and Automation Letters, 1033:2112-2119, 2025.
Zhiyuan Xu, Kun Wu, Zhengping Che, Jian Tang, and Jieping Ye. Knowledge transfer in multi- task deep reinforcement learning for continuous control. In Advances in Neural Information Processing Systems, volume 33, pp. 15146-15155, 2020.
Ruihan Yang, Huazhe Xu, Yi Wu, and Xiaolong Wang. Multi-task reinforcement learning with soft modularization. In Advances in Neural Information Processing Systems, volume 33, pp. 4767- 4777, 2020.
Deheng Ye, Zhao Liu, Mingfei Sun, Bei Shi, Peilin Zhao, Hao Wu, Hongsheng Yu, Shaojie Yang, Xipeng Wu, Qingwei Guo, et al. Mastering complex control in moba games with deep reinforce- ment learning. In Proceedings of the AAAI conference artificial intelligence, volume 34, pp- 2020.
Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. In Advances in Neural Information Processing Systems, volume 33, pp. 15534-1555 15547. Curran Associates, Inc., 2020a.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference robot learning, pp. 1094-1100, PMLR, 2020b.
Matthew D Zeiler, Dilip Krishnan, Graham W Taylor, and Rob Fergus. Deconvolutional networks. In 2010 IEEE Computer Society Conference on computer vision and pattern recognition, pp- 25355 IEEE, 2010.
Boxuan Zhang, Runqing Wang. Wei Xiao, Weipu Zhang, Jian Sun, Gao Huang, Jie Chen, and Gang Wang. Dymodreamer: World modeling with dynamic modulation. In Conference on Neural Information Processing Systems, 2025.
Weipu Zhang, Gang Wang, Jian Sun, Yetian Yuan, and Gao Huang. Storm: Efficient stochastic transformer based world models for reinforcement learning. In Advances in Neural Information Processing Systems, volume 37. Curran Associates, Inc., 2023.
12

---

Published as a conference paper at ICLR 2026
A APPENDIX
A.1 FURTHER ANALYSIS
We highlight the distinctions between MoW and recent approaches in the world model and MTRL
as follows:
Table 2: Comparison between MoW and recent approaches.
Attributes
CARE
MOORE
DreamerV3
STORM
MoW (ours)
MBRL/MFRL
Model free
Model free
Model based
Model based
Model based
Setting
Multi-task
Multi-task
Single-task
Single-task
Multi-task
World Model
GRU
Transformer
Hybrid Transformer
Input
State
State MTSAC
Image DreamerV3
Image
Image
Agent training
MTSAC
As DreamerV3
Multi-task STORM
CARE (Sodhani et al., 202 and MOORE (Hendawy al., (023) are model-free multi- task RL methods that suffer from low sample efficiency, whereas DreamerV3 (Hafner et al. 2025 and STORM (Zhang et al. 2023) are workd-model-basee single-task reinforeement learning approaches that exhibit high sample efficiency but lack multi-task implementa- tions. In contrast, MoW a world model- based MTRL method that combines high sample efficiency with strong multi-task performance.
In contrast DreamerV3 (Hafner et al., 2025), which employs GRU (Dey & Salem 2017). and STORM (Zhang et al., 2023). which adopts a vanilla Transformer as the se- quence model, MoW introduces a novel hybrid Transformer architecture that better cap- tures diverse task dynamics while preserving reconstruction fidelity.
CARE (Sodhani et al. 2021 ab and MOORE (Hendawy al. 2023) are state-based methods that build on the classical MTSAC algorithm for decision-making, with their primary focus on introducing MoE architectures for policy networks and task-specific training strategies. DreamerV3 V3 (Hafner et al. 2025) and STORM (Zhang et al., (223), in contrast, adopt tor-critic methods to train the agent on imagined trajectories, with their main contributions centered encoding and predicting more accurate image observations by the hidden rep- resentations. MoW extends the architecture of STORM by leveraging gradient-based clus- tering allocate distinct critic networks different tasks, while still employing similar actor-critic training paradigm for the agent.
A.2 THE USE OF LARGE LANGUAGE MODELS (LLMs)
LLMs are employed to assist in detecting and rectifying potential grammatical errors the manuscript.
13

---

Published as a conference paper at ICLR 2026
A.3 GRADIENT-BASED CLUSTERING
The warmup stage in MoW is a short offline world-model self-supervised training phase on fixed replay buffer (as we shown in Algorithm [ Before clustering, we first collect sufficient amount of random interaction data for all tasks and then train a single VAE-predictor stack on this fixed dataset while keeping the agent frozen (Sec. 3.6). This stage corresponds to self-supervised learning of reconstruction, reward, continuation, dynamic and task prediction losses, rather than unstable online policy optimization. we shown in the Figure 6 these losses converge within a few thousand optimization steps, indicating that the world model has already fitted the warmup dataset well. In our experiments, varying the warmup length within reasonable range (e.g., 5000 steps) has only minor impact on final Atari and Meta-Word performance, suggesting that MoW is not overly sensitive to the exact duration of warmup, as long as the losses have essentially stabilized.
during Warm-up up Stage
2/09
1:00
2:08
180
180
X08
Staps
Figure 6: The gradient norm curves in the warmup stage of MoW.
14

---

Published as a conference paper at ICLR 2026
A.4 CHALLENGES IN DIGH-DIMENSIONAL OBSERVATIONS
Reinforcement learning in high-dimensional observation spaces, such as those encountered in vi- sual tasks, presents several challenges compared to state-based RL, where observations typi- cally simpler and lower-dimensional. These challenges arise due to the complexities of learning and processing nigh-dimensional data, such as images or videos, and the need for efficient gener- alization across multiple tasks. In state-based RL, the observation space is often relatively small (e.g., dimension vector representing the agent's position and velocity (Sodhani et al. (021b), ) and traditional MTRL algorithms can work effectively by directly modeling these state represen- tations. In contrast, high-dimennional observations, such as images, have large input spaces and complex structures that must be compressed into useful representations. In addition, state-based RL typically utilizes a single, simpler model, as state representations are usually more structured and wwwrrdiimesiionl.. However, high-ding -igh-dimensional visual inputs require specialized visual modules, and the online optimization of these modules during training imposes a significant burden on the convergence of reinforcement learning.
To further illustrate the challenges of high-dimensional visual MTRL compared to state- state-based MTRL, we use official results from the classic multi-task method TDMPC2 (Hansen et al.) 2024). and the results is shown in Figure 7 In the 12 tasks of Deepmind Control, the performance of TDMPC2 with visual pixels inputs (the orange lines) is worse and converges to the state- based inputs (the blue lines). It fails to converge after 2M steps of single-task interac- tion. To ensure fairness, the only difference between these two models is the addition of a shallow convolutional encoder in the visual input version, with all other parameters being identical.
slowly compared
Engorgeeeery
Dnvvemeee strp
Expreent
Eecterent
Figure 7: Results of pixels-input and states-input TDMPC2 on DMC suite.
15

---

Published as a conference paper at ICLR 2026
A.5 ADDITIONAL ABLATION STUDIES
Due to the page limitations in the main text, we provide additional ablation on all of our design choices for MoW in this section, including the clusters of VAE-predictors, number of experts, cas- caded expert-shared transformers, task predictor head, balanced loss, and the gradient-based task clustering. The results are shown in Figure 8 andd 9
Clusters of VA TAE-predictors. We ablate the number of clusters of VAE-predictors Nw from 12 to 3 and observe a clear performance degradation as the clusters count decreases. With fewer clusters, each VAE-predictor must capture the visual characteristics of a larger set of tasks, which slows down convergence and leads lower overall performance.
Number of Expert Transformers. We ablate the number of expert Transformers Ne from 12 to 3 (when N₂ = 1, the model degenerates into a single shared Transformer) and observe clear per- formance degradation as the experts count decreases. With fewer experts, it becomes more difficult to capture the diverse dynamics across tasks, which reduces the quality of imagined trajectories generated by the world model and ultimately makes policy learning more challenging.
Expert and Shared Transformers. We ablate the shared Transformer and observe a modest perfor- mance drop, tasks in Atari 100K share relatively little inter-task structure. In contrast, removing the expert Transformer leads to a much larger degradation, since the expert Transformer is respon- sible for capturing sasc-spec dynamics and generating accurate multi-task imagined trajectories.
Task Prediction. We ablate the task predictor by removing both the prediction head and its asso- ciated loss. The task predictor plays important role in maintaining task-disrimminative structure in the learned latent space, which benefits both world-model learning and downstream RL perfor- mance. Consistent with this, removing the task predictor leads noticeable decline in overall performance.
Harmonious Loss. We ablate the harmonious loss by removing the harmonious weight in Equation 9 Without this operation, the world-model training faces the gradient conflicting problem, and the end-to-end loss is dominated by tasks with larger loss magnitudes, leading to unstable training and reduce the RL performance. This confirms that harmonizing the multi-task losses is important for stable and effective multi-task visual world-model learning.
Balance Loss. To evaluate whether the expert- balenne term is necessary to prevent overuse of a few experts, we remove the balance loss and observe that the training becomes sensitive to initialization. The balance loss is critical for ensuring all experts are utilized and prevents collapse into degenerate expert usage.
Gradient-based Task Clustering. We use randomly partition tasks into the same number of clus- ters (Nm 12) to ablate the gradient-basee warmup clustering. The results show that the gradient- based clustering is necessary for allocating shared parameters meaningfully and avoiding interfer- ence across unrelated tasks.
Clusters VAE-predictars
Number
Expert Shared Trarsfommess
Prediction
mwan moan
sas cestiiiss VAG-ore cara 3 (AGoodd carr
Popert Dpetta
Nocture Tunulamar
330.0%
Human normaliien
3 Dipett
58.5%
13.55%
75.3%
52 52.3%
Harmonicus Locs
Balance Less
Gradient based Task Clustering
3.5
16
Hurman surmaltredd
1.
110.4%
1.2
110
10..0%
32
110.4%
10. 9%
Figure 8: Results of additional ablation studies
16

---

Published as a conference paper at ICLR 2026
Clusters VAE-predictors
Number Experts
Expert and Shared Transforers
Task Prediction
Leant::
preeters
Exsorb
Exaort Tenibreer
Human normaliio
Harmenibus
Balance
Gradient based Clustering
cormalzeed mose
Maa7
NUM
0.2
62-
0.4
6.-
100k
30
Figure 9: Results of additional ablation studies
17

---

Published as a conference paper at ICLR 2026
A.6 RECONSTRUCTION BY IMAGINATION
We provide Multi-step predictions on several environments in Atari games. MoW utilizes 8 observa- tions and actions as contextual input, enabling the imagination of future events in auto-regressive manner.
Imagination
Environment: Alien
Imagination
Environment: Amidar
Imagination
Environment: Assault
Imagination
Environment: Asterix
Imagination
Environment: BankHelst
Imagination
Environment:
BattleZone
Imagination
Environment: Boxing
Imagination
Environment:
Breakout
Imagination
Environment: ChopperCommand
Imagination
Environment: CrazyClimber
Figure 10: Multi-step predictions on several environments in Atari games.
18

---

Published as a conference paper at ICLR 2026
Imagination
Environment: DemonAttack
Imagination
Environment:
Freeway
Imagination
Environment:
Frostbite
Imagination
Environment: Gopher
Imagination
Environment: Hero
Imagination
Environment: Jamesbond
Imagination
Environment: Kangaroo
Imagination
Krull
Imagination
Environment:
FuMastrr
Imagination
Environment: MsPacman
Figure 11: Multi-step predictions on several environments in Atari games.
19

---

Published as a conference paper at ICLR 2026
Imagination
Environment: Pang
Imagination
Environment: PrivateEye
Imagination
Environment: Obert
Imagination
Environment:
RoadRunner
Imagination
Environment: Seaquest
Imagination
Environment: UpNDown
Figure 12: Multi-step predictions several environments in Atari games.
20

---

Published as a conference paper at ICLR 2026
A.7 HARMONIOUS Loss
Adjusting the coefficients of different loss terms has the potential to significantly enhance the per- formance of world model-based reinforcement learning algorithms. The harmonious loss strategy dynamically adjusts the dominance between the observation and the reward modeling during world model training process. The harmonious loss is constructed as:
Cx(0) Σ cx(0) +n(1 +kk) k=1
(16)
The reason why we use In(1 + σk) as regularization terms that о is parameterized as dk = exp(ax) 0 to optimize parameters ds free of sign constraint. Since loss item with small val- ues, such as the reward loss, can lead to extremely large coefficient 1/ Rs L > 1, which potentially hurt training stability. The derivation of the harmonious loss scale follows.
To minimize [( on)]. we force the the partial derivative w.r.t. to O:
V=E(CH(0,0)]
V. B[Cx] + In(1 + ov))
20-222/0 B[Cx]] 1 ok
(17)
E[Ca] + V E + 4E(CA]
o
Therefore the learnable loss weight, the rectified harmonious loss, approximates the analytic loss weight:
(18)
E[Ca] + [ CA
and equivalently, the harmonized loss scale is:
1,
(19)
1
+ ECA
and add the regularization term In(1 + σR)) results in the 4/ 4/E[Ca] in the /I + 4/E[Cx] term, which prevents the loss weight from getting extremely large when faced with small [[[x].
21

---

Published as a conference paper at ICLR 2026
A.8 DETAILS OF MODEL STRUCTURE
Table and Table 4 show the structure of the image encoder and decoder. The sizes of the submod- ules omitted they can be directly inferred from the tensor shapes. ReLU denotes the rectified linear unit activation, and Linear corresponds to a fully connected layer. Flatten and Reshape oper- ations are applied solely to modify the indexing of tensors while preserving both the data and their original ordering. Conv refers to CNN layer(LeCun et al. 1989). characterized by kernel = 4, stride 2, and padding
and
DeConv denotes a transpose CNN layer (Zeiler af., 2010). also characterized by kernel = 4, stride = 2, and padding = 1. BN denotes the batch norma ization layer (loffe & Szegedy 2015). The variable E represents the dimension of task embedding Ek.
Table 3: Structure of the image encoder.
Submodule
Output tensor shape
Input image (o
3 64 64
Convl + BN1 ReLU
32 x 32 32
Conv2 + BN2 ReLU
64 16 16
Conv3 + BN3 ReLU
128 x 8 x 8
Conv4 + BN4 ReLU
256 x 4
Flatten
4096
Linear
1024
Table 4: Structure of the image decoder.
Submodule
Output tensor shape
Random sample (2()
32 x 32
Flatten and task embedded
1024 E 4096
BNO+eL+
Reshape DeConyl + BN1 ReLU
256 4 x 4
128 x 8 x 8
DeConv2 + BN2 ReLU
64 x 16
16
DeConv3 + BN3 ReLU
32 x 32 x 32
DeConv4 (produce of))
3 64 64
1st ex). Concatenate denotes combining the last di- Table 5: Action mixer m Ma a
mension of tensors and merging them into one new tensor. The variable A represents the action dimension, which is 18 for Atari and for Meta-worid. D denotes the feature dimension of the Transformer. LN is an abbreviation for layer normalization
Submodule
Output tensor shape
Random sample (zf), action (af), task embedding (ex)
32 x 32, A, E 1024 + E + A
Reshape and concatenate
Linear I +LNN+ Linear2+LN2+ReLU
D
D
22

---

Published as a conference paper at ICLR 2026
Table 6: Transformer block. Dropout (Vaswani 2017) is used in each Transformer submodule to reduce overfitting. We also apply Dropout to attention weights in the MHSA module.
Submodule
Module alias
Output tensor shape
Input features (label as x1)
32 x 32
Multi-head self attention
Linearl + Dropout
MHSA
TxD D
Residual (add æ1)
LNael as x22)
Linear2 ReLU
T 2D
Linear3 + Dropout
T D T D
FFN
Residual (add x2)
LN2
TX D
Table 7: Expert Transformer 1/12 = Sou (m where j is the index the expert.
Submodule
Output tensor shape
Input (m²);)
T x D T x D Tx
Transformer blocks xM
Output (g)))
Table 8: Shared Transformer h!t = Foll ex).
Submodule
Output tensor shape
Input (1 ec) Reshape Transformer blocks xM
T' x (n x D + E)
T
T x D
Output (h):))
T
Table 9: Pure MLP structures. A layer MLP corresponds to fully-connected layer, 255 is the size of the bucket of symlog two-hot loss, and K is the number of tasks.
Submodule
MLP layers
Input/ MLP hidden/ Output dimension
Task embedding (ek)
1/ /96
Dynamic predictor (g
1
D + E 11024
Reward predictor (go.i))
3
D D/D/255
Continuation predictor (90.5
3
D/D/1
Task predictor (g
1
D/-/K
Policy network (Vy.i.)
3
D/D/A D/D/255
Actor network (Ty (af s(.))
3
23

---

Published as a conference paper at ICLR 2026
A.9 HYPERPARAMETERS
Table 10 detail hyperparameters of the optimization and environment, as well as hyperparameters shared by multiple components. Hyerparameters.
Table 10: Hyperparameters
Name
Value
World model training batch size
16
World model training batch length
64
Imagination batch size
1024
Imagination context length
8
Imagination horizon
16
Warmup steps
5000
Environment context length
16(Atari100K) (3taaata-or)ld)
Update world model every step Update agent model every env step
1
Number experts
12(Atari100K) /20(Meta-world)
Number of activated experts
3(A (Atari100K) /4(Meta-worll)
Number of VAE clusters
Learning rate
10-1
World model gradient clipping
1000
Adam epsilon
10-8
Optimizer
Adam
Expert Transformer layers
2
Shared Transformer layers
2
Expert Transformer heads
8
Shared Transformer heads
8
Dropout
0.1
Discount horizon
333
Return lambda
0.95
Critic EMA decay
0.98
Critic EMA regularizer
1
Actor entropy scale
3 x 10-4
Learning rate
3 10-5
Actor-critic gradient clipping
100
24

---

Published as a conference paper at ICLR 2026
A.10 EXPERIMENTAL DETAILS
A.10.1
ATARI 100K CURVES AND AGGREGATE SCORES
We evaluate agent performance 100 episodes using the final model checkpoint, saved every 2500 steps.
1200
400
65-
36x boorg
battl 2oka
Chapper crrsard
urayy cinber
cald
0000
506
3/4
10x
50
devan atace
1eewwy
Fipter
-441
2002
300-
1940
1880
W
1380
Vieum Word Redet
2rgit 5T55
Figure 13: Evaluation curves of MoW and single-task STORM the Atari100K benchmark for individual games (400K environment steps). The solid lines represent the average scores over seeds, and the filled areas indicate the standard deviation across these seeds.
Table I1 summarizes the aggregate scores of MoW and some Transformer-based world model base- lines (Robine al. 2023 Micheli al. 2022
A.10.2 META-WORLD
The Meta-Wil World benchmark encompasses diverse array of 50 distinct manipulation tasks, unified by shared dynamics. These tasks involve a Sawyer robot engaging with variety of objects, each distinguished by unique shapes, joints, and connective properties. The complexity this bench- mark lies in the heterogeneity of the observation spaces and reward functions across tasks, as the robot required to manipulate different objects towards varying objectives. The robot operates with a 4-dimensional fine-grained action input at each timestep, which controls the 3D positional movements its end effector and modulates the gripper's 's openness. In order to use a single cam- the modified corner2 camera viewpoint for all
viewpoint consistently over all 50 tasks, tasks. Specifically, we adjusted the camera position with env env.mode www.model.cam pos[2][]]=00.75 0.075, 0.7], which enables solve non-zero success rate all tasks. Maximum episode length for
25

---

Published as a conference paper at ICLR 2026
Table 11: Game scores and overall human-normalized performance the 26 games in the Atari 100K benchmark.
Game
Random
Tuman
STORM (origina))
STORM treprodaced)
MoN
TWM
IRIS
Elticien/Zem
BBr
EfficienZeeee22
(ours)
Setting
Suagle task
Single task
Multi task
Single task
Single task
Single task
Single task
Single task
Alien
238
7128
964
748
511
675
4.00
809
1175
1558
Amidar
6
1720
205
122
142
245
222
1534
7one
24460
Boxing
Breakoot
30
20
371
Chopper Command
811
7388
1988
1644
1504
1697
1365
1117
7589
1197
Ciury Climber
10780
35820
667766
79196
64554
71870
59234
83940
58452
112363
Demon Attack
152
1971
165
324
173
350
2034
13004
13341
22773
Frezway
0
30
24
31
22
26
Frosthit
65
4335
1316
365
839
1476
250
296
2365
1136
Gopber
258
2413
5307
1673
2236
Hoin
7254
Kangaroo
5035
Kung
Ms. Pacman
307
6952
2673
1890
2492
1588
909
[28]
2008
Pong
-21
13
18
13
19
15
17
Privale Fye
69571
7781
100
92.
K7
100
100
Qbert
164
13455
4522
2910
3641
3331
746
13762
4487
1605S
Road Ranner
12
7545
17564
14841
17169
9109
9615
17751
35427
27516
Seocuest
42055
525
557
431
774
661
1100
1235
1974
Up Down
533
11893
7985
6128
6943
19983
3546
17264
12102
15224
Human Mean
Median
Meta-worte world tasks is 500. We follow the environment setup of Meta-World as specified in Masked World Model (MWM) (Seo et al. 2023). with the exception of using different action repeat setting (2 in MWM but in MoW). All curves shown in Figure 14
26

---

Published as a conference paper at ICLR 2026
Bapthhgg
434
63
04
110
20th
200
30
68
180
-206
2006
colfe-attem
cof
0.4
4.4
4.4
Sxonect
WAWA
0.4
2
4.-
2
4.
1106
200
pepAistopads
180287
atcoout
0.3
0.-
0.2
130
120
270
180
-50
1066
miceean
2P483
Insrppes
sidsecies
Picture
Sirghetask-5000
Figure 14: Evaluation curves of MoW on the Meta- World benchmark.
27

---

Published as a conference paper at ICLR 2026
A.11 BENCHMARKING TD-MPC2
We compare the performance of MoW with TD-MPC2 on the Meta- World benchmark. In the official implementation of TD-MPC2, the multi-task learning setup based on large offline datasets (545M transitions), where the model is trained on the fixed state-based data rather than learning directly from online interactions. Moreover, TD-MPC2 is designed for state-based inputs, unlike MoW, which is tailored to handle high-dimensional visual inputs. To make a fair comparison in the context of visual reinforcement learning (RL), we adapted the official state-based implementation of TD- MPC2 by incorporating a shallow convolutional network to handle visual inputs (this operation is same to the visual RL experiments in TD-MPC2). We then tested this modified version of TD-MC22 against MoW in the same World tasks. This modification allowed us to assess the performance of TD-MPC2 in a visual setting, providing insights into its capability to scale to visual RL tasks compared to MoW's more native approach handling visual inputs and multi-task learning.
TDMPC-2 achieves success rate in the same environment as MoW. The results shown in Figure [5 demonstrate that while TD-MPC2 excels in continuous control tasks, MoW offers more scalable and parameter-efficient approach for multi-task reinforcement learning that can effectively handle both continuous and discrete action spaces, especially when high-dimensional visual inputs are involved.
TD-MPC2 and its Challenge with Discrete Actions: The core of TD-MPC2's methodology lies on Model Predictive Path Integral (MPPI) control, which involves optimizing over continuous action sequences maximize expected returns. The key advantage of this approach is its ability to perform local trajectory optimization in a learned world model, where actions are treated as contin- uous values. The model predicts continuous action trajectories over short planning horizon and uses these predictions to guide policy execution. The trajectory optimization process inherently as- sumes continuous actions, as involves sampling action sequences from multivariate Gaussian distribution and updating their mean and variance iteratively. The Gaussian distribution allows the model to effectively explore action space by adjusting the mean and variance of actions, and this is inherently designed for continuous action dimensions. In discrete action spaces, however, unlike the continuous action spaces where actions are represented as continuous variables that can be ad- justed smoothly, discrete actions limited finite set of predefined choices. This discrete nature doesn't align well with the sampling mechanism in MPC, which relies on continuous action adjust- ments. Efficient Monte Carlo Tree Search (MCTS) or other searching methods used in the discrete action settings would require specific adjustments for handling discrete action sets, and standard MPC planning would not scale properly without a discrete action search space method.
6.-
Average Success Rate
6.6
1.5
611
50.
Environment steps
Mixture-of-ooll Models
TD-MPC2
Figure 15: Comparison of MoW and TD-MPC2 on the Meta- World benchmark.
28

---

Published as a conference paper at ICLR 2026
A.12 PSEUDOCODE
For a more comprehensive understanding of training the mixture world model, we provide the pseu-
docode implementation in Algorithm 1
Algorithm Mixture World Model-based Reinforcement Learning
Input:
Task set T (71, T.. ]
replay buffer B = [1,, B2, Bg
warm-up steps Tw
total number of environment steps Trotal.
Initialize: Trained mixture world model parameters 6,
actor-critic parameters and v
actor-critic policy TT.
replay buffer B = 0,
number of environment steps
number of warmup steps tu = 0
// Warmup stage
Repeat:
// J. Collect data from environment
B = for Tyy a }
t+ 1
Until Each Bę has sufficient data for warmup training.
Repeat:
2. Offline training of the mixture world model
D/ sample-trajectory! BB))
Φ apdate- (
tx
Until L. Tw
Training stage
Repeat:
/1 3. Collect data from environment
B too,r rf, a )
Update mixture world model
D₁ sample-trajectory (Bx)
0 = world-molool
5. Sample observation for imagination
0 sample-obs B&)
wordd-modell-maaiine(( O TT))
// 6. Learning in the imagination hypdate-policy11m 0, v))
1+ =
Until t Teoaal-
29


Reviews:
[{"review_id": "zYAavMoWid", "reviewer": "Reviewer_XJhM", "summary": "This paper addresses the challenge of sample efficiency in Multi-Task Reinforcement Learning (MTRL). The authors propose Mixture-of-World (MoW) Models, a modular architecture integrating:\n- Modular VAEs specifically clustered for task-adaptive visual compression.\n- A hybrid dynamics model combining task-conditioned MoE (Mixture of Experts) Transformers with a shared Transformer backbone.\n- Gradient-based task clustering to efficiently allocate parameters during a warmup phase.", "questions": "1. Given the high variance of gradients early in standard RL training (especially from pixels), how stable are the resulting task clusters across different random seeds? Did you observe cases where a poor warmup led to irrecoverable sub-optimal clusterings?\n2. How sensitive is MoW to the fixed hyperparameters? Were these tuned specifically for Atari and Meta-World, or do you expect them to generalize to new domains?\n3. Does task-level routing limit zero-shot generalization to new tasks? If a new task is a composition of two existing tasks, your current architecture forces it to choose just one set of experts, whereas token-level might allow it to interpolate.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Tailored MoE specifically for World Model Dynamics. Standard MoE in Large Language Models often uses token-level routing. The authors rightly identify that for dynamics modeling, token-level routing can lead to \"fragmented learning,\" where sequential temporal dependencies are broken if consecutive tokens route to different experts.\n2. End-to-end MTRL often suffers from gradient conflicts and tasks dominating the loss landscape. MoW integrates distinct mechanisms to stabilize this.\n3. Parameter Efficiency via Gradient Clustering. Instead of naively assigning a VAE per task (expensive) or one VAE for all (insufficient for heterogeneous visuals), they use a warmup phase to cluster tasks based on gradient similarity. This allows intelligent parameter sharing for the perceptual modules, balancing capacity with efficiency", "weaknesses": "1. Dependency on Warmup Quality. The gradient-based clustering relies heavily on a \"warmup stage\" where a single VAE/predictor set is trained. If the initial warmup yields noisy gradients (common in early RL from pixels), the resulting clusters might be suboptimal and fixed for the rest of training. The paper does not deeply analyze the sensitivity of final performance to the duration or stability of this warmup phase.\n2. Architectural Complexity and Tuning. The system is highly complex, involving multiple specialized losses (reconstruction, reward, continuation, task prediction, dynamics KL, representation KL, harmonious weighting, expert balance). While effective, this increases the hyperparameter surface significantly. Reproducibility outside specifically tuned benchmarks might be challenging without robust auto-tuning mechanisms for these loss components.\n3. Routing Rigidity. While task-level routing prevents temporal fragmentation, it might be too rigid for tasks that share partial sub-dynamics. Task-level routing forces it to choose one set of experts for the entire episode.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Psr0T0px0w", "reviewer": "Reviewer_hP6A", "summary": "This paper proposes Mixture-of-World Models (MoW) for multi-task reinforcement learning. The architecture employs task-cluster-specific VAEs combined with a hybrid transformer containing both shared and expert modules. Task grouping is performed through gradient-based clustering, and the training incorporates an expert balance loss and adaptive loss weighting to stabilize learning. Experiments on Atari100k and Meta-World benchmarks show that MoW achieves comparable or improved performance with a more compact model design.", "questions": "1. While there are few studies on online multi-task reinforcement learning on Atari, there exist some offline multi-task works [1, 2]. Could MoW be extended to offline settings, and would its modular design still provide efficiency or performance advantages there? (I understand this would be a substantial effort and not necessarily feasible during the review period.)\n2. In Line 372, how exactly is the gradient vector per task computed? Is it obtained by averaging gradients over all mini-batches associated with each task?\n\n[1] Multi-Game Decision Transformers\n\n[2] Scaling Offline Model-Based RL via Jointly-Optimized World-Action Model Pretraining", "limitations": "", "rating": 2, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. To the best of my knowledge, modular architectures for world models are still underexplored, making this paper’s direction novel and valuable.\n2. The design choices—such as task-level routing, the combination of fixed task clusters for VAEs with dynamic expert routing for transformers—are conceptually interesting and insightful.\n3. The qualitative comparison in Figure 2 shows that MoW produces notably higher-quality imagined rollouts than the multi-task STORM baseline.", "weaknesses": "1. **The ablation study is severely incomplete.** The method introduces many components—such as cluster-specific VAEs, cascaded expert-shared transformers, an auxiliary task predictor head, balanced loss, and gradient-based task clustering—but none of these components are ablated to clarify their individual contributions to performance. Without such analysis, it is hard to assess which design elements are truly essential.\n2. The overall architecture is highly complex, introducing a large number of new hyperparameters, which may make reproducibility and tuning difficult.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iClae7upnJ", "reviewer": "Reviewer_ujJn", "summary": "This paper proposes a well-motivated architecture—Mixture-of-World Models (MoW)—for multi-task visual RL that marries task-specific VAEs with a mixture-of-expert Transformers and a shared backbone, plus task-prediction and expert-balance losses and a gradient-based task clustering warm-up. The approach targets the core pain point of heterogeneous observations/dynamics and reports strong results: a single model reaches 110.4% human-normalized score on Atari-100k, competitive with STORM’s 114.2% while using ~50% fewer parameters, and achieves 74.5% average success on Meta-World MT50 within 300k steps per task. Overall, this is a solid contribution to scalable, parameter-efficient world models for MTRL with promising empirical evidence.", "questions": "1. What does $m_{\\phi,j}$ represent in Equation 2?\n2. Is $W_k$ sorted? To my perspective, it is not, so you should verify the definition of $W_k$.\n3. What does $N_m$ represent?\n4. What does $l_k^{1:t}$ represent?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "**Clear modular design**. MoW uses reasonable modular architecture designs.\n\n**Promising parameter scalability**. MoW demonstrates promising parameter scaling capability. However, it would be interesting if the author can further scale the parameters to investigate if the performance could be further exponentially improved.", "weaknesses": "**Natation clarity**. It would be beneficial to clarify the notations in the figures, for example, Fig. 1, which could significantly enhance the readability of this paper. It would also be beneficial if you could explain the notation a little bit when it first appears. Details see questions.\n\n**Lack of experimental evidence**. The author didn't demonstrate the common Atari-100k performance table, including mainstream baselines.\n\n**Insufficient baselines.** STORM is actually not a multi-task RL baseline. I personally cannot understand why the authors include more single-task baselines as comparisons.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kgSPKO093l", "reviewer": "Reviewer_uU94", "summary": "The Mixture-of-World Models (MoW) is a new, scalable, and parameter-efficient architecture designed for multi-task reinforcement learning (MTRL), specifically targeting visual environments where tasks have highly different visual observations and dynamics. The core challenge is that traditional model-based RL (MBRL) world models struggle to generalize across diverse tasks, leading to poor performance. MoW overcomes this by integrating three key modular components: a) Modular Variational Autoencoders (VAEs) for task-adaptive visual data compression, b) a hybrid Transformer-based dynamics model that shares a core backbone while utilizing task-conditioned specialized experts, and c) a gradient-based task clustering mechanism to efficiently allocate parameters. A single MoW agent demonstrated strong results, achieving a competitive human-normalized score of 110.4% on the Atari 100k benchmark while requiring 50% fewer parameters than the previous state-of-the-art ensemble method (STORM). It also established a new state-of-the-art success rate of 74.5% on the Meta-World benchmark, proving MoW's effectiveness as a foundation for building scalable, generalist world models.", "questions": "- Please benchmark TD-MPC2 on both the Atari 100k and Meta-World benchmarks and compare its performance with MoW. I believe this will highlight the potential of the proposed method.\n- The inclusion of the single-task STORM training curves is highly recommended. Could you please provide these for Figure 3?\n- I believe there should be consistency in the learning settings across MOORE, the baselines, and MoW. Please consider updating the empirical results on Meta-World to account for these aspects. \n- To strengthen the paper, it would be beneficial to demonstrate how MOORE and the other baselines fail to effectively learn from raw visual inputs, which would highlight the advantage of the proposed MoW architecture.\n- What are the model sizes (e.g., number of parameters) for MoW and the baselines when benchmarking on Meta-World?\n- An ablation study is requested to demonstrate the effect of changing the different algorithmic components on overall performance.\n- Could you please provide the individual task learning curves in the appendix for the Meta-World tasks?\n\nI acknowledge that these requests constitute a substantial body of work for the rebuttal period. However, addressing these crucial empirical and validation points is necessary, in my opinion, for the paper to be ready for publication and subsequent acceptance.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 1, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written and structured.\n- The work is well-motivated: The authors clearly state the general problem of multi-task RL (MTRL) with visual inputs. They effectively motivate their proposed solution of using model-based RL (MBRL) via a world model, specifically addressing the technical obstacle of learning diverse task dynamics with a single model through the use of a Mixture of Experts (MoE).\n- The learning algorithm integrates several interesting components, such as the annealing of the temperature coefficient in the softmax function, which is used to prevent distribution collapse and limit parameter sharing.\n- The results on Meta-World are strong, with MoW's performance using visual inputs clearly outperforming its results using state vectors. However, I would have expected the baselines to be benchmarked with visual inputs as well, in order to demonstrate their failure to handle high-dimensional visual inputs.", "weaknesses": "- My main criticism is the poor empirical evaluation, which makes the submission incomplete.\n- Although the two benchmarks, Atari 100k and Meta-World, are diverse, the baselines provided are very sparse.\n- For instance, the learning curves for single-task STORM [1] are missing in Figure 3, where only the MoW results are shown.\n- Since TD-MPC2 [2] was highlighted in the related work section, I would have expected to see it benchmarked against the proposed algorithm on both benchmarks.\n- Since the baseline results on Meta-World were taken from MOORE [3], there are two issues to address.\n- One issue is the total number of environment steps for the baselines. The appendix of MOORE stated 100M (2M per task), which contradicts the 20M cited here. I believe the 20M is a typo in the MOORE paper, as recently highlighted by other work [4]. Although the author is not to blame for this, the correct value should be stated.\n- A more crucial issue is the mismatch between the proposed approach's hyperparameters and the MOORE baseline results (e.g., MOORE used a planning horizon of 150 instead of 500). I advise either rerunning the baselines (which I recommend) or rerunning MoW with the exact hyperparameters stated in the MOORE Appendix.\n- While MoW is clearly more parameter-efficient than single-task STORM, the model size (number of parameters) for MoW and the baselines on Meta-World was not stated. I suspect that MoW uses a very large model relative to MOORE and the other baselines.\n- Since the proposed algorithm contains many algorithmic components, I would have expected a full ablation study detailing their individual effect on performance.\n\n[1] Zhang, Weipu, et al. \"Storm: Efficient stochastic transformer based world models for reinforcement learning.\" Advances in Neural Information Processing Systems 36 (2023): 27147-27166.\n\n[2] Hansen, Nicklas, Hao Su, and Xiaolong Wang. \"Td-mpc2: Scalable, robust world models for continuous control.\" arXiv preprint arXiv:2310.16828 (2023).\n\n[3] Hendawy, Ahmed, Jan Peters, and Carlo D'Eramo. \"Multi-Task Reinforcement Learning with Mixture of Orthogonal Experts.\" The Twelfth International Conference on Learning Representations.\n\n[4] McLean, Reginald, et al. \"Multi-Task Reinforcement Learning Enables Parameter Scaling.\" arXiv preprint arXiv:2503.05126 (2025).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)