Title: Mixture-of-Experts Meets In-Context Reinforcement Learning
Venue: NeurIPS 2025
Accepted: yes
Keywords: In-context reinforcement learning

mixture-of-experts

offline reinforcement learning

contrastive learning

transformers
Abstract: 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** (**T**oken- and **T**ask-wise **M**oE for **I**n-context **R**L), 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](https://github.com/NJU-RL/T2MIR).


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper tackles the challenge of in-context reinforcement learning (ICRL). To cope with (i) multi-modality of state-action-reward tokens and (ii) the wide variety of decision tasks, the authors replace each transformer feed-forward block with two Mixture-of-Experts (MoE) layers: one that routes individual tokens by modality, and another that routes whole trajectories by task. Evaluated on four benchmark suites—DarkRoom, Point-Robot, MuJoCo, and Meta-World—the proposed architecture surpasses ICRL and offline meta-RL baselines.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n1. The method outperforms competitive ICRL and offline meta-RL baselines across DarkRoom, Point-Robot, multiple MuJoCo locomotion tasks, and Meta-World Reach & Push.\n2. The paper carefully derives the dual MoE routing scheme and its contrastive objective, making the architectural choices easy to follow.\n3. The ablation study and t-SNE visualizations shed light on how token- and task-level experts contribute to the gains.\n\n**Weaknesses**\n1. The paper asserts “the learning efficiency can be impeded by intrinsic gradient conflicts in challenging scenarios with significant task variation,” yet provides no empirical evidence of such conflicts in the pre-training phase of ICRL. Since the online adaptation phase is gradient-free, the relevance of this issue remains unclear.\n2. While the authors motivate token-wise MoE by the heterogeneity of state, action and reward tokens, they do not quantify how serious this mismatch is for existing ICRL backbones.\n3. Experiments focus on relatively simple tasks (e.g., Meta-World Reach/Push). It is uncertain whether the approach scales to harder suites such as the full ML1 or ML10.\n\nQUESTIONS:\n1. Do you have quantitative evidence of gradient conflict to justify the task-wise MoE?\n2. Can you isolate the modality issue so as to substantiate that semantic discrepancy is a major challenge of ICRL?\n3. Have you evaluated T2MIR on more challenging benchmarks?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-23"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a framework named T2MIR which embeds token-wise MoE and task-wise MoE into transformer layers to handle multimodal ICRL data and diverse tasks. The authors choose AD and DPT as backbone models of the framework and evaluated them on five multi-task benchmarks. Across mixed datasets of the five benchmarks, both T2MIR-AD and T2MIR-DPT outperform strong baselines in learning speed and final return.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n - The authors propose a novel architecture to integrate MoE into ICRL problems to tackle multi-modality transition and task heterogeneity.\n - The experiments on five diverse benchmarks and varying data qualities are comprehensive, showing robust performance gains of T2MIR.\n - The ablations and analysis are inisghtful. Ablation studies isolate the contributions of two MoEs; the t-SNE andn gradient-cosine plots illustrate expert specialization and conflict mitigation.\n\nWeakness:\n - The DarkRoom and Point-Robot benchmarks seem to be weak in modern RL research. I suggest the authors consider other Mujoco benchmarks like AntDir and Humanoid, or text based benchmarks like webshop [1].\n - In this paper, the task inference is based on InfoNCE which takes hidden states and past trajectories to infer a task embedding. UNICORN broke a transition tuple (s, a, r, s') into two components: (s, a) that is more related to behavior policy, and (r, s') that is more related to the task [2]. I suggest the authors analyze how each component contributes to the task representation in their work, e.g., given two transitions from two tasks that have the same (s,a), but different (r, s'), how their task representations differ. I believe this analysis is important to show the robustness of the work in OOD scenario.\n\n[1] Yao, Shunyu, et al. \"Webshop: Towards scalable real-world web interaction with grounded language agents.\" Advances in Neural Information Processing Systems 35 (2022): 20744-20757.\n[2] Li, Lanqing, et al. \"Towards an information theoretic framework of context-based offline meta-reinforcement learning.\" arXiv preprint arXiv:2402.02429 (2024).\n\nQUESTIONS:\n1. Clarification of “multi-modality” among states, actions, and rewards - In many environments, these signals are deeply intertwined. For instance, in Cheetah-Vel the state comprises each segment’s angles, velocities, and angular velocities; the action is the torques applied to those segments; and the reward is a deterministic function of the resulting state. Could the authors clarify in what sense such tightly coupled signals are treated as separate “modalities” in their token-wise MoE routing?\n2. Evaluation under IID tasks but OOD offline data - Have the authors tested T2MIR when tasks are sampled IID from P(M) but the offline dataset is out-of-distribution—e.g., collected by a random policy or by a policy trained on a different task? Evaluating under this dataset shift would more rigorously demonstrate T2MIR’s robustness.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work addresses the challenge of designing a scalable in-context RL framework capable of handling multimodality and diverse task distributions within a single transformer model. To this end, it introduces a Mixture of Experts (MoE) architecture that augments the transformer’s ability to manage the complexity arising from multimodality induced by states, actions, and rewards, as well as the generalization difficulties posed by task diversity. Specifically, the proposed framework replaces the standard feedforward layers in causal transformer blocks with two specialized modules: a token-wise MoE and a task-wise MoE. Experimental results across a range of tasks demonstrate the effectiveness of these modifications, and ablation studies further validate the design choices.\n\nSTRENGTHS AND WEAKNESSES:\nThe idea of using the MoE to tackle the multimodality and generalization challenges sounds promising, and as the paper has shown, the empirical evidence suggests that it has the potential to improve the performance of in-context RL. The authors have conducted a good set of experiments addressing various aspects of the work, which I appreciate. The paper is fairly well-written, though improvements can be made to improve the comprehension.  \n\nOn a constructive note, considering the central role of task distributions in this work, the definition of task variation remains somewhat ambiguous. Appendix D.1 implies that these variations may be parameterized, in which case it would be natural to formalize the problem within the framework of a Contextual Markov Decision Process (CMDP). However, if the task variations are not consistently parameterizable, a different formalization may be more suitable. As currently presented, the paper lacks a clear formalism and specification of how tasks differ, making it challenging to interpret the relationships between tasks or understand how their difficulty evolves, specifically in interpreting the results.\n\nSecond, upon examining Figure 1 and the ablation in Section 5.4 on token-wise and task-wise MoE visualizations, I began to question whether learning the gating mechanism in MoEs is actually necessary. For instance, in the left figure of Figure 1—if I'm interpreting it correctly—it seems that three experts specializing in state, action, and reward tokens may suffice to get reasonably competing performance. This raises the possibility of manually defining a fixed gating scheme based on token type. Similarly, in the right panel of Figure 2, a simple heuristic gating based on evenly splitting the velocity range could be applied (this is also why I think the formalism on how task variations should be defined, as it tells you if one could do this or not). I'm curious how such heuristic approaches would perform compared to learned gating. Such an approach would not require separate loss terms like balancing and contrastive loss. It’s possible that I may be overlooking something, as the authors have likely considered this aspect carefully, so I would appreciate the authors' thoughts on the matter.\n\nThe experiments conducted are nice and cover the major important axes of interest. However, it would have been nice to have ablations to show the impact of balancing loss and contrastive loss. The authors provide the intuition as to why they are needed, and I agree with them, but it would be nice to see the impact scale of the proposed loss terms. Similarly, I could imagine challenges of training with the contrastive loss when the number of tasks is large. Any thoughts/analysis that authors have on this end?\n\nQUESTIONS:\n1. In the Figure 6 analysis, when the token-wise MoE count is changed, how is the task-wise MoE count handled? I think they could be related.\n2. Also, in the Figure 6 analysis, how is the decision to activate one-third of experts in token-wise MoE and always 2 experts in task-wise MoE reached? \n3. Since the task variations seem to be parameterized in the experiments, do you also use the task variation parameterization in the context?\n\nLIMITATIONS:\nYes. It is given in the appendix.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-02"}

comments: SUMMARY:
The paper proposes improving in context RL performance by introducing explicit MoE layers that operate at token vs task level to ensure that transformers can represent the multi-modality due to state, actions and rewards efficiently. The paper demonstrates the efficacy of this architecture on environments like DarkRoom, MuJoCo, Meta-World and demonstrates faster convergence and better results than baseline methods.

STRENGTHS AND WEAKNESSES:
Strengths:
1. The paper addresses the fact that introducing MoE leads to imbalance and tries to address them via explicit losses to ensure all experts can be utilized.
2. The insight of pooling trajectory to ensure interference of gradient is reduced is interesting.
3. Results demonstrate faster convergence and better rewards compared to baseline methods.

Weakness:
1. Adding two MoE blocks per layer adds a lot of complexity for both training and inference. Balancing the experts requires explicit loss, introduces even more hyper parameters and I am not completely convinced that additional complexity via explicit architectural bias for tasks such as in context RL is required.
2. The paper performs experiments on small sized data with relatively small modes and I am not completely sure if this scales well and if standard large models even require explicit architectural bias.

QUESTIONS:
1. In the ablation experiments, its not very clear that when you remove MoE do you keep the parameter count of the model same as with MoE to ensure fair representation capacity?

LIMITATIONS:
Yes
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-07-03