Смесь экспертов встречается с обучением с подкреплением в контексте

Mixture-of-Experts Meets In-Context Reinforcement Learning

June 5, 2025

Авторы: Wenhao Wu, Fuhong Liu, Haoru Li, Zican Hu, Daoyi Dong, Chunlin Chen, Zhi Wang

cs.AI

Аннотация

Обучение с подкреплением в контексте (In-context Reinforcement Learning, ICRL) стало перспективной парадигмой для адаптации агентов RL к последующим задачам через кондиционирование подсказок. Однако две значительные проблемы остаются для полного использования обучения в контексте в областях RL: внутренняя многомодальность данных состояния-действия-вознаграждения и разнородный характер задач принятия решений. Для решения этих проблем мы предлагаем T2MIR (Token- and Task-wise MoE for In-context RL), инновационную архитектуру, которая внедряет усовершенствования смеси экспертов (Mixture-of-Experts, MoE) в трансформерные модели принятия решений. T2MIR заменяет полносвязный слой двумя параллельными слоями: токен-ориентированным MoE, который захватывает различные семантики входных токенов в нескольких модальностях, и задаче-ориентированным MoE, который направляет разнообразные задачи к специализированным экспертам для управления широким распределением задач с уменьшением конфликтов градиентов. Для улучшения задаче-ориентированной маршрутизации мы вводим метод контрастного обучения, который максимизирует взаимную информацию между задачей и её представлением в маршрутизаторе, что позволяет более точно захватывать информацию, релевантную задаче. Выходы двух компонентов MoE объединяются и передаются на следующий слой. Комплексные эксперименты показывают, что T2MIR значительно улучшает способность к обучению в контексте и превосходит различные типы базовых моделей. Мы раскрываем потенциал и перспективы MoE для ICRL, предлагая простое и масштабируемое архитектурное усовершенствование, которое приближает ICRL к достижениям в сообществах обработки языка и компьютерного зрения. Наш код доступен по адресу 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.