mlsanityData to 14 Oct 2026

One Model for All Tasks: Leveraging Efficient World Models in Multi-Task Planning

arXiv 2509.07945 ↗First posted 9 Sep 2025cs.LGcited by 3

On the map: World models → Latent predictive world models

also in 1 more
Abstract (arXiv)

In heterogeneous multi-task decision-making, tasks not only exhibit diverse observation and action spaces but also vary substantially in their underlying complexities. While conventional multi-task world models like UniZero excel in single-task settings, we find that when handling a broad and diverse suite of tasks, gradient conflicts and the loss of model plasticity often constrain their sample efficiency. In this work, we address these challenges from two complementary perspectives: the single learning iteration and the overall learning process. First, to mitigate the gradient conflicts, we systematically investigate key architectural designs for extending UniZero. Our investigation identifies a Mixture-of-Experts (MoE) architecture as the most effective approach. We demonstrate, both theoretically and empirically, that this architecture alleviates gradient conflicts by routing task-specific representations to specialized sub-networks. This finding leads to our proposed model, textitScaleZero. Second, to dynamically allocate model capacity throughout the learning process, we introduce an online Dynamic Parameter Scaling (DPS) strategy. This strategy progressively integrates LoRA adapters in response to task-specific progress, enabling adaptive knowledge retention and parameter expansion. Evaluations on a diverse set of standard benchmarks (Atari, DMC, Jericho) demonstrate that ScaleZero, utilizing solely online reinforcement learning with one model, performs on par with specialized single-task agents. With the DPS strategy, it remains competitive while using just 71.5% of the environment interactions. These findings underscore the potential of ScaleZero for effective multi-task planning. Our code is available at textcolormagentahttps://github.com/opendilab/LightZero.

Key facts. Cited by 3 papers in the map (data through 14 Oct 2026). Placed in Latent predictive world models, one of 11 directions in World models. Builds on 54 of its 66 references that are in the map. First posted on arXiv 9 Sep 2025.

Builds on

54 of 66 references.

Closest ancestors · 2

Same direction, newest first

  1. 2024UniZero: Generalized and Efficient Planning with Scalable Latent World Models15 Jun
  2. 2022Temporal Difference Learning for Model Predictive Control9 Mar

Other references · 52

  1. 2017Attention Is All You Needcited by 60,361
  2. 2015Deep Residual Learning for Image Recognitioncited by 52,096
  3. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 30,186
  4. 2019Momentum Contrast for Unsupervised Visual Representation Learningcited by 7,513
  5. 2016Layer Normalizationcited by 5,669
  6. 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,509
  7. 2023QLoRA: Efficient Finetuning of Quantized LLMscited by 3,418
  8. 2013Playing Atari with Deep Reinforcement Learningcited by 3,274
  9. 2019Mastering Atari, Go, Chess and Shogi by Planning with a Learned Modelcited by 1,353
  10. 2020Gradient Surgery for Multi-Task Learningcited by 1,256
  11. 2018IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architecturescited by 883
  12. 2024DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Modelscited by 867
38 more
  1. 2017GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networkscited by 834
  2. 2024DoRA: Weight-Decomposed Low-Rank Adaptationcited by 758
  3. 2017Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithmcited by 734
  4. 2023TD-MPC2: Scalable, Robust World Models for Continuous Controlcited by 512
  5. 2019Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variablescited by 389
  6. 2024LoRA+: Efficient Low Rank Adaptation of Large Modelscited by 377
  7. 2020dm_control: Software and Tasks for Continuous Controlcited by 375
  8. 2022Transformers are Sample-Efficient World Modelscited by 328
  9. 2021Mastering Atari Games with Limited Datacited by 237
  10. 2022Multi-Game Decision Transformerscited by 214
  11. 2021Multi-Task Reinforcement Learning with Context-based Representationscited by 174
  12. 2023The Dormant Neuron Phenomenon in Deep Reinforcement Learningcited by 174
  13. 2019Interactive Fiction Games: A Colossal Adventurecited by 168
  14. 2020Multi-Task Reinforcement Learning with Soft Modularizationcited by 143
  15. 2023Transformer-based World Models Are Happy With 100k Interactionscited by 139
  16. 2021Online and Offline Reinforcement Learning by Planning with a Learned Modelcited by 90
  17. 2023LoRAMoE: Alleviate World Knowledge Forgetting in Large Language Models via MoE-Style Plugincited by 82
  18. 2024Mixtures of Experts Unlock Parameter Scaling for Deep RLcited by 78
  19. 2021Learning and Planning in Complex Action Spacescited by 77
  20. 2022Offline Q-Learning on Diverse Multi-Task Data Both Scales And Generalizescited by 73
  21. 2024Smooth Tchebycheff Scalarization for Multi-Objective Optimizationcited by 54
  22. 2022Towards Understanding Mixture of Experts in Deep Learningcited by 52
  23. 2023Recon: Reducing Conflicting Gradients from the Root for Multi-Task Learningcited by 51
  24. 2023Multi-Task Reinforcement Learning with Mixture of Orthogonal Expertscited by 49
  25. 2022PaCo: Parameter-Compositional Multi-Task Reinforcement Learningcited by 46
  26. 2024Sharing Knowledge in Multi-Task Deep Reinforcement Learningcited by 44
  27. 2024MTLoRA: A Low-Rank Adaptation Approach for Efficient Multi-Task Learningcited by 44
  28. 2023MultiLoRA: Democratizing LoRA for Better Multi-Task Learningcited by 32
  29. 2023Learning to Modulate pre-trained Models in RLcited by 28
  30. 2024Unleashing the Power of Multi-Task Learning: A Comprehensive Survey Spanning Traditional, Deep, and Pretrained Foundation Model Erascited by 26
  31. 2023HarmonyDream: Task Harmonization Inside World Modelscited by 22
  32. 2024Jack of All Trades, Master of Some, a Multi-Purpose Transformer Agentcited by 21
  33. 2025LoRI: Reducing Cross-Task Interference in Multi-Task Low-Rank Adaptationcited by 19
  34. 2024MTL-LoRA: Low-Rank Adaptation for Multi-Task Learningcited by 17
  35. 2023Contrastive Modules with Temporal Attention for Multi-Task Reinforcement Learningcited by 10
  36. 2023Not All Tasks Are Equally Difficult: Multi-Task Deep Reinforcement Learning with Dynamic Depth Routingcited by 8
  37. 2023Efficient Multi-Task and Transfer Reinforcement Learning with Parameter-Compositional Frameworkcited by 7
  38. 2023Language Model-In-The-Loop: Data Optimal Approach to Learn-To-Recommend Actions in Text Gamescited by 3

50 most cited of 52 All 52 (JSON)

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2026: 3.

3 papers.

  1. 2026PriorZero: Bridging Language Priors and World Models for Decision Makingcited by 1
  2. 2026Eradicating Negative Transfer in Multi-Physics Foundation Models via Sparse Mixture-of-Experts Routingcited by 0
  3. 2026WMLLM: Self-Evolving Optimization Agents via Predict-Then-Act World Modelingcited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2025Mastering Massive Multi-Task Reinforcement Learning via Mixture-of-Expert Decision Transformer
  2. 2026Mixture-of-World Models: Scaling Multi-Task Reinforcement Learning with Modular Latent Dynamics
  3. 2024Merging Multi-Task Models via Weight-Ensembling Mixture of Experts
  4. 2026Routing-Aligned Fine-Tuning for Multilingual Downstream Tasks in Mixture-of-Experts Models
  5. 2025Mixture-of-Experts Meets In-Context Reinforcement Learning
  6. 2026UniW(A)M: A Unified Mixture-of-Experts Framework for World and World-Action Models
  7. 2024AT-MoE: Adaptive Task-planning Mixture of Experts via LoRA Approach
  8. 2024Efficient and Effective Weight-Ensembling Mixture of Experts for Multi-Task Model Merging