Mixture-of-Experts trifft auf In-Context Reinforcement Learning

Mixture-of-Experts Meets In-Context Reinforcement Learning

June 5, 2025

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

cs.AI

Zusammenfassung

In-context Reinforcement Learning (ICRL) hat sich als vielversprechendes Paradigma etabliert, um RL-Agenten durch Prompt-Konditionierung an nachgelagerte Aufgaben anzupassen. Dennoch bleiben zwei wesentliche Herausforderungen bestehen, um In-context Learning in RL-Domänen vollständig nutzbar zu machen: die inhärente Multimodalität der Zustand-Aktion-Belohnungs-Daten und die vielfältige, heterogene Natur von Entscheidungsaufgaben. Um diese Herausforderungen zu bewältigen, schlagen wir T2MIR (Token- und Task-wise MoE für In-context RL) vor, ein innovatives Framework, das architektonische Fortschritte im Bereich Mixture-of-Experts (MoE) in transformerbasierte Entscheidungsmodelle einführt. T2MIR ersetzt die Feedforward-Schicht durch zwei parallele Schichten: ein token-weises MoE, das unterschiedliche Semantiken von Eingabe-Tokens über mehrere Modalitäten hinweg erfasst, und ein task-weises MoE, das verschiedene Aufgaben an spezialisierte Experten weiterleitet, um eine breite Aufgabenverteilung mit reduzierten Gradientenkonflikten zu bewältigen. Um das task-weise Routing zu verbessern, führen wir eine kontrastive Lernmethode ein, die die gegenseitige Information zwischen der Aufgabe und ihrer Router-Repräsentation maximiert und so eine präzisere Erfassung aufgabenrelevanter Informationen ermöglicht. Die Ausgaben der beiden MoE-Komponenten werden verkettet und in die nächste Schicht eingespeist. Umfassende Experimente zeigen, dass T2MIR die In-context-Lernfähigkeit signifikant fördert und verschiedene Arten von Baselines übertrifft. Wir bringen das Potenzial und die Versprechen von MoE in die ICRL ein und bieten eine einfache und skalierbare architektonische Verbesserung, um ICRL einen Schritt näher an die Erfolge in den Bereichen Sprache und Vision zu bringen. Unser Code ist verfügbar unter 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.