Mixture-of-Experts Meets In-Context Reinforcement Learning
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.
68527e8ce329b3b5e93f2997
Wenhao Wu
claimed_verified
2025-06-18T10:58:48.394Z
6756970ce110734a48701a08
https://cdn-avatars.huggingface.co/v1/production/uploads/6756970ce110734a48701a08/yn24gxS2LhZHRE3NYRkjr.jpeg
Wenhao Wu
user
Wenhao0
68527e8ce329b3b5e93f2998
Fuhong Liu
68527e8ce329b3b5e93f2999
Haoru Li
68527e8ce329b3b5e93f299a
Zican Hu
claimed_verified
2025-06-20T12:19:52.817Z
662de927b10f82aa4f46557f
/avatars/19878db42ecf8e35456378571eae4643.svg
Zican Hu
user
huzican
68527e8ce329b3b5e93f299b
Daoyi Dong
68527e8ce329b3b5e93f299c
Chunlin Chen
68527e8ce329b3b5e93f299d
Zhi Wang
Mixture-of-Experts Meets In-Context Reinforcement Learning
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.
T2MIR, a framework using token-wise and task-wise MoE in transformer-based decision models, enhances in-context reinforcement learning by addressing multi-modality and task diversity.
Mixture-of-experts (MoE) transformer-based decision models token-wise MoE task-wise MoE contrastive learning mutual information