Mastering Massive Multi-Task Reinforcement Learning via Mixture-of-Expert Decision Transformer

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Kong, Yilun, Ma, Guozheng, Zhao, Qi, Wang, Haoyu, Shen, Li, Wang, Xueqian, Tao, Dacheng
Formato: Preprint
Publicado: 2025
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866908385360740352
author Kong, Yilun
Ma, Guozheng
Zhao, Qi
Wang, Haoyu
Shen, Li
Wang, Xueqian
Tao, Dacheng
author_facet Kong, Yilun
Ma, Guozheng
Zhao, Qi
Wang, Haoyu
Shen, Li
Wang, Xueqian
Tao, Dacheng
contents Despite recent advancements in offline multi-task reinforcement learning (MTRL) have harnessed the powerful capabilities of the Transformer architecture, most approaches focus on a limited number of tasks, with scaling to extremely massive tasks remaining a formidable challenge. In this paper, we first revisit the key impact of task numbers on current MTRL method, and further reveal that naively expanding the parameters proves insufficient to counteract the performance degradation as the number of tasks escalates. Building upon these insights, we propose M3DT, a novel mixture-of-experts (MoE) framework that tackles task scalability by further unlocking the model's parameter scalability. Specifically, we enhance both the architecture and the optimization of the agent, where we strengthen the Decision Transformer (DT) backbone with MoE to reduce task load on parameter subsets, and introduce a three-stage training mechanism to facilitate efficient training with optimal performance. Experimental results show that, by increasing the number of experts, M3DT not only consistently enhances its performance as model expansion on the fixed task numbers, but also exhibits remarkable task scalability, successfully extending to 160 tasks with superior performance.
format Preprint
id arxiv_https___arxiv_org_abs_2505_24378
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Mastering Massive Multi-Task Reinforcement Learning via Mixture-of-Expert Decision Transformer
Kong, Yilun
Ma, Guozheng
Zhao, Qi
Wang, Haoyu
Shen, Li
Wang, Xueqian
Tao, Dacheng
Machine Learning
Artificial Intelligence
Despite recent advancements in offline multi-task reinforcement learning (MTRL) have harnessed the powerful capabilities of the Transformer architecture, most approaches focus on a limited number of tasks, with scaling to extremely massive tasks remaining a formidable challenge. In this paper, we first revisit the key impact of task numbers on current MTRL method, and further reveal that naively expanding the parameters proves insufficient to counteract the performance degradation as the number of tasks escalates. Building upon these insights, we propose M3DT, a novel mixture-of-experts (MoE) framework that tackles task scalability by further unlocking the model's parameter scalability. Specifically, we enhance both the architecture and the optimization of the agent, where we strengthen the Decision Transformer (DT) backbone with MoE to reduce task load on parameter subsets, and introduce a three-stage training mechanism to facilitate efficient training with optimal performance. Experimental results show that, by increasing the number of experts, M3DT not only consistently enhances its performance as model expansion on the fixed task numbers, but also exhibits remarkable task scalability, successfully extending to 160 tasks with superior performance.
title Mastering Massive Multi-Task Reinforcement Learning via Mixture-of-Expert Decision Transformer
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2505.24378