Mezcla de Expertos se Encuentra con el Aprendizaje por Refuerzo en Contexto

Mixture-of-Experts Meets In-Context Reinforcement Learning

June 5, 2025

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

cs.AI

Resumen

El aprendizaje por refuerzo en contexto (ICRL, por sus siglas en inglés) ha surgido como un paradigma prometedor para adaptar agentes de aprendizaje por refuerzo (RL) a tareas posteriores mediante el condicionamiento de indicaciones. Sin embargo, persisten dos desafíos notables para aprovechar plenamente el aprendizaje en contexto dentro de los dominios de RL: la multimodalidad intrínseca de los datos de estado-acción-recompensa y la naturaleza diversa y heterogénea de las tareas de decisión. Para abordar estos desafíos, proponemos T2MIR (Token- y Task-wise MoE para ICRL), un marco innovador que introduce avances arquitectónicos de la mezcla de expertos (MoE) en modelos de decisión basados en transformadores. T2MIR sustituye la capa de avance con dos capas paralelas: un MoE a nivel de token que captura semánticas distintas de los tokens de entrada a través de múltiples modalidades, y un MoE a nivel de tarea que dirige diversas tareas a expertos especializados para gestionar una amplia distribución de tareas con conflictos de gradiente mitigados. Para mejorar el enrutamiento a nivel de tarea, introducimos un método de aprendizaje contrastante que maximiza la información mutua entre la tarea y su representación de enrutador, permitiendo una captura más precisa de la información relevante para la tarea. Las salidas de los dos componentes MoE se concatenan y se alimentan a la siguiente capa. Experimentos exhaustivos muestran que T2MIR facilita significativamente la capacidad de aprendizaje en contexto y supera a varios tipos de líneas base. Llevamos el potencial y la promesa de MoE a ICRL, ofreciendo una mejora arquitectónica simple y escalable para acercar ICRL un paso más hacia los logros en las comunidades de lenguaje y visión. Nuestro código está disponible en 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.