一个模型搞定多种任务,通过专家路由和动态参数扩展提升效率。
One Model for All Tasks: Leveraging Efficient World Models in Multi-Task Planning
- 用专家混合架构缓解多任务梯度冲突,提升模型适应性。
- 在Atari、DMC等基准上表现媲美专用模型,交互次数减少28.5%。
- 适合需要高效多任务决策的强化学习场景,如机器人控制。
在异构多任务决策中,任务不仅观察与动作空间各异,复杂度也差异显著。传统多任务世界模型如UniZero虽在单任务中表现优异,但在处理广泛多样任务时,常因梯度冲突和模型可塑性下降导致样本效率受限。本文从单次学习迭代和整体学习过程两方面应对挑战:首先,系统研究UniZero的结构设计,发现专家混合(MoE)架构最有效,能通过将任务特定表征路由至专用子网络,理论与实证证明其缓解梯度冲突;由此提出新模型ScaleZero。其次,引入在线动态参数缩放(DPS)策略,根据任务进展逐步集成LoRA适配器,实现自适应知识保留与参数扩展。在Atari、DMC、Jericho等标准基准上的评估表明,ScaleZero仅用在线强化学习与单一模型,性能媲美专用单任务代理;结合DPS策略后,仍保持竞争力,且仅需71.5%的环境交互次数。结果验证了ScaleZero在多任务规划中的潜力。代码已开源:https://github.com/opendilab/LightZero。
原文摘要 · Abstract (English)
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, \textit{ScaleZero}. 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 \textcolor{magenta}{https://github.com/opendilab/LightZero}.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。