Research map / Reinforcement learning / Visual and representation RL

Mixture-of-Experts Meets In-Context Reinforcement Learning

NeurIPS 2025 · Reinforcement learning › Visual and representation RL · approach: In-context and adaptive RL · 2 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes T2MIR, which replaces transformer feedforward layers with parallel token-wise and task-wise mixture-of-experts layers (the task router trained with an InfoNCE contrastive loss) to improve in-context RL on AD and DPT backbones.

Problem

In-context RL struggles with the multi-modal state-action-reward token data and with gradient conflicts across a diverse, heterogeneous task distribution within a single transformer policy.

Approach

Substitute the transformer feedforward layer with two parallel MoE layers: a token-wise MoE routing each token by modality/semantics, and a task-wise MoE whose router is trained with a contrastive (InfoNCE, MoCo-style momentum) objective to maximize mutual information between task and router representation.

Key results

• T2MIR-AD reaches 90.9 test return on DarkRoom (Mixed), tied with DICP and above AD (57.5) and DPT (72.5).
• T2MIR-DPT achieves -43.2 on Cheetah-Vel (Mixed) vs. -56.1 for DPT and -119.2 for AD.
• T2MIR-AD scores 823.9 on Reach (Mixed), above the best baseline IDT at 808.5.

Measured effects reported in the paper

Surprising findings

• DICP outperforms T2MIR on DarkRoom in learning speed and matches its asymptotic return, which the authors attribute to model-based planning in small state-action spaces.
• Ablating the token-wise MoE hurts T2MIR-AD more than ablating the task-wise MoE, while the reverse holds for T2MIR-DPT.

Limitations

• Evaluated only on relatively lightweight benchmark datasets rather than large-scale environments.
• Gains over strong baselines are modest in some tasks (e.g., Point-Robot differences within ~0.5 return).
• Hand-tuned balance loss weights and expert configurations required.

Experimental setup

Datasets: DarkRoom Point-Robot Cheetah-Vel Walker-Param Reach Push Mixed Medium-Expert Medium

Benchmarks: DarkRoom Point-Robot Cheetah-Vel Walker-Param Reach Push

Metrics: test return gradient cosine similarity

Baselines: AD DPT IDT DICP UNICORN

Builds on: Algorithm Distillation Decision-Pretrained Transformer Shazeer sparse MoE MoCo InfoNCE

Closest papers by meaning

• Multi-Task Reinforcement Learning with Mixture of Orthogonal Experts (ICLR 2024 · similarity 0.90)
• Mastering Massive Multi-Task Reinforcement Learning via Mixture-of-Expert Decision Transformer (ICML 2025 · similarity 0.90)
• XLand-100B: A Large-Scale Multi-Task Dataset for In-Context Reinforcement Learning (ICLR 2025 · similarity 0.89)
• Towards Provable Emergence of In-Context Reinforcement Learning (NeurIPS 2025 · similarity 0.89)
• In-Context Decision Transformer: Reinforcement Learning via Hierarchical Chain-of-Thought (ICML 2024 · similarity 0.89)
• Vintix: Action Model via In-Context Reinforcement Learning (ICML 2025 · similarity 0.88)