Mélange d'Experts Rencontre l'Apprentissage par Renforcement en Contexte

Mixture-of-Experts Meets In-Context Reinforcement Learning

June 5, 2025

Auteurs: Wenhao Wu, Fuhong Liu, Haoru Li, Zican Hu, Daoyi Dong, Chunlin Chen, Zhi Wang

cs.AI

Résumé

L'apprentissage par renforcement en contexte (ICRL) s'est imposé comme un paradigme prometteur pour adapter les agents d'apprentissage par renforcement (RL) à des tâches en aval grâce au conditionnement par prompts. Cependant, deux défis majeurs persistent pour exploiter pleinement l'apprentissage en contexte dans les domaines du RL : la multimodalité intrinsèque des données état-action-récompense et la nature diverse et hétérogène des tâches de décision. Pour relever ces défis, nous proposons T2MIR (Token- and Task-wise MoE for In-context RL), un cadre innovant qui intègre des avancées architecturales de type mixture-of-experts (MoE) dans des modèles de décision basés sur des transformeurs. T2MIR remplace la couche feedforward par deux couches parallèles : un MoE token-wise qui capture les sémantiques distinctes des tokens d'entrée à travers plusieurs modalités, et un MoE task-wise qui achemine les tâches variées vers des experts spécialisés pour gérer une large distribution de tâches tout en atténuant les conflits de gradients. Pour améliorer le routage task-wise, nous introduisons une méthode d'apprentissage contrastif qui maximise l'information mutuelle entre la tâche et sa représentation de routage, permettant une capture plus précise des informations pertinentes à la tâche. Les sorties des deux composants MoE sont concaténées et transmises à la couche suivante. Des expériences approfondies montrent que T2MIR facilite significativement la capacité d'apprentissage en contexte et surpasse divers types de modèles de référence. Nous apportons le potentiel et les promesses du MoE à l'ICRL, offrant une amélioration architecturale simple et évolutive pour rapprocher l'ICRL des réalisations des communautés du langage et de la vision. Notre code est disponible à l'adresse https://github.com/NJU-RL/T2MIR.

English

In-context reinforcement learning (ICRL) has emerged as a promising paradigm for adapting RL agents to downstream tasks through prompt conditioning. However, two notable challenges remain in fully harnessing in-context learning within RL domains: the intrinsic multi-modality of the state-action-reward data and the diverse, heterogeneous nature of decision tasks. To tackle these challenges, we propose T2MIR (Token- and Task-wise MoE for In-context RL), an innovative framework that introduces architectural advances of mixture-of-experts (MoE) into transformer-based decision models. T2MIR substitutes the feedforward layer with two parallel layers: a token-wise MoE that captures distinct semantics of input tokens across multiple modalities, and a task-wise MoE that routes diverse tasks to specialized experts for managing a broad task distribution with alleviated gradient conflicts. To enhance task-wise routing, we introduce a contrastive learning method that maximizes the mutual information between the task and its router representation, enabling more precise capture of task-relevant information. The outputs of two MoE components are concatenated and fed into the next layer. Comprehensive experiments show that T2MIR significantly facilitates in-context learning capacity and outperforms various types of baselines. We bring the potential and promise of MoE to ICRL, offering a simple and scalable architectural enhancement to advance ICRL one step closer toward achievements in language and vision communities. Our code is available at https://github.com/NJU-RL/T2MIR.