La Miscela di Esperti incontra l'Apprendimento per Rinforzo in Contesto

Mixture-of-Experts Meets In-Context Reinforcement Learning

June 5, 2025

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

cs.AI

Abstract

L'apprendimento per rinforzo in contesto (In-context Reinforcement Learning, ICRL) è emerso come un paradigma promettente per adattare gli agenti di RL a compiti downstream attraverso il condizionamento dei prompt. Tuttavia, due sfide significative rimangono per sfruttare appieno l'apprendimento in contesto nei domini RL: l'intrinseca multimodalità dei dati stato-azione-ricompensa e la natura eterogenea e diversificata dei compiti decisionali. Per affrontare queste sfide, proponiamo T2MIR (Token- and Task-wise MoE for In-context RL), un framework innovativo che introduce avanzamenti architetturali basati su mixture-of-experts (MoE) nei modelli decisionali basati su transformer. T2MIR sostituisce il livello feedforward con due livelli paralleli: un MoE token-wise che cattura le distinte semantiche dei token di input attraverso più modalità, e un MoE task-wise che indirizza compiti diversi a esperti specializzati per gestire un'ampia distribuzione di compiti con conflitti di gradiente attenuati. Per migliorare il routing task-wise, introduciamo un metodo di apprendimento contrastivo che massimizza l'informazione mutua tra il compito e la sua rappresentazione di routing, consentendo una cattura più precisa delle informazioni rilevanti per il compito. Le uscite dei due componenti MoE vengono concatenate e passate al livello successivo. Esperimenti completi dimostrano che T2MIR facilita significativamente la capacità di apprendimento in contesto e supera vari tipi di baseline. Portiamo il potenziale e la promessa di MoE all'ICRL, offrendo un miglioramento architetturale semplice e scalabile per avvicinare l'ICRL ai risultati ottenuti nelle comunità di linguaggio e visione. Il nostro codice è disponibile su 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.