# Mixture-of-Experts Meets In-Context Reinforcement Learning arXiv: https://arxiv.org/abs/2506.05426 · HTML: https://mlsanity.com/paper/2506.05426/ ## Key facts - Cited by 8 papers in the map (data through 14 Oct 2026). - Placed in MoE in reinforcement learning and robotics, one of 12 directions in Mixture-of-experts. - Builds on 63 of its 77 references that are in the map. - First posted on arXiv 5 Jun 2025. ## Placed in - Mixture-of-experts / [MoE in reinforcement learning and robotics](https://mlsanity.com/moe/moe-in-reinforcement-learning-and-robotics/) - Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/) ## Builds on (most cited first, 30 of 63) - [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020) - [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017) - [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023) - [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018) - [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019) - [DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning](https://mlsanity.com/paper/2501.12948/) (2025) - [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017) - [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018) - [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022) - [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017) - [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016) - [Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context](https://mlsanity.com/paper/2403.05530/) (2024) - [OpenVLA: An Open-Source Vision-Language-Action Model](https://mlsanity.com/paper/2406.09246/) (2024) - [π0.5: a Vision-Language-Action Model with Open-World Generalization](https://mlsanity.com/paper/2504.16054/) (2025) - [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021) - [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019) - [Training Diffusion Models with Reinforcement Learning](https://mlsanity.com/paper/2305.13301/) (2023) - [DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models](https://mlsanity.com/paper/2401.06066/) (2024) - [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020) - [Is Conditional Generative Modeling all you need for Decision-Making?](https://mlsanity.com/paper/2211.15657/) (2022) - [RL2: Fast Reinforcement Learning via Slow Reinforcement Learning](https://mlsanity.com/paper/1611.02779/) (2016) - [The Entropy Mechanism of Reinforcement Learning for Reasoning Language Models](https://mlsanity.com/paper/2505.22617/) (2025) - [Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables](https://mlsanity.com/paper/1903.08254/) (2019) - [Learning to Reason under Off-Policy Guidance](https://mlsanity.com/paper/2504.14945/) (2025) - [From Sparse to Soft Mixtures of Experts](https://mlsanity.com/paper/2308.00951/) (2023) - [In-context Reinforcement Learning with Algorithm Distillation](https://mlsanity.com/paper/2210.14215/) (2022) - [A Survey of Zero-shot Generalisation in Deep Reinforcement Learning](https://mlsanity.com/paper/2111.09794/) (2021) - [Prompting Decision Transformer for Few-Shot Policy Generalization](https://mlsanity.com/paper/2206.13499/) (2022) - [RAPHAEL: Text-to-Image Generation via Large Mixture of Diffusion Paths](https://mlsanity.com/paper/2305.18295/) (2023) - [Supervised Pretraining Can Learn In-Context Reinforcement Learning](https://mlsanity.com/paper/2306.14892/) (2023) ## Built on by (most cited first, 8 of 8) - [Diversity-Incentivized Exploration for Versatile Reasoning](https://mlsanity.com/paper/2509.26209/) (2025) - [DyGRO-VLA: Cross-Task Scaling of Vision-Language-Action Models via Dynamic Grouped Residual Optimization](https://mlsanity.com/paper/2605.17486/) (2026) - [VLGOR: Visual-Language Knowledge Guided Offline Reinforcement Learning for Generalizable Agents](https://mlsanity.com/paper/2603.22892/) (2026) - [In-Context Reinforcement Learning under Non-Stationarity: A Survey](https://mlsanity.com/paper/2607.11906/) (2026) - [Robust Exploration in Directed Controller Synthesis via Reinforcement Learning with Soft Mixture-of-Experts](https://mlsanity.com/paper/2602.19244/) (2026) - [WiSDoM: Wireless Sparse Decision Transformer with Mixture-of-Experts for Multi-Task Mobile Network Optimization](https://mlsanity.com/paper/2609.00284/) (2026) - [MA-WAM: Multi-Agent World-Action Model for Test-Time Planning](https://mlsanity.com/paper/2609.31281/) (2026) - [Scalable In-Context Reinforcement Learning with Recurrent Algorithm Distillation](https://mlsanity.com/paper/2609.35333/) (2026) ## Abstract (arXiv) 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.