Mixture-of-Experts Meets In-Context Reinforcement Learning

---

Q1. Empirical evidence of gradient conflicts in ICRL pre-training.

A1. In the original manuscript, we have visualized the comparison of gradient conflicts between T2MIR-AD (MoE) and AD (MLP) as shown in Fig. 7 (L297, Page 9). It shows that AD suffers from significant gradient conflict between opposing tasks and fails to discriminate tasks. In contrast, T2MIR-AD maintains nearly orthogonal gradients across diverse tasks, especially for opposing ones. This comparison verifies T2MIR’s advantage in both gradient conflict mitigation and task discrimination.

Further, we also provide quantification of the gradient conflict issue. The following table presents the proportion of task pairs with negative correlations among all pairs of tasks, which demonstrates T2MIR’s significant superiority in mitigating gradient conflicts.

Thank you for your insightful comment. In addition to Section 5 Experiments, we will include empirical evidence of gradient conflicts and corresponding analysis in Section 1 Introduction, to further justify our motivation for the task-wise MoE.

---

Q2. Isolate and quantify the multi-modality issue of state-action-reward tokens.

A2. In the original manuscript, we can observe the performance gain when addressing the multi-modality issue by token-wise MoE in Section 5.2 Ablation Study (Figure 4, Page 8). The AD and DPT baselines obtain relatively poor performance, and incorporating the token-wise MoE induces a 6$\sim$36% return increase for AD and a 2$\sim$12% increase for DPT.

Further, motivated by your insightful comment, we completely isolate the multi-modality issue and try to quantify the challenge incurred by this semantic discrepancy. We manually design an MoE structure where state-action-reward tokens are routed to three experts by a heuristic gating scheme, i.e., one expert for one modality. The following table shows that the heuristic gating mechanism can improve ICRL performance by a 3$\sim$24% increase in received return. It again demonstrates the significant challenge induced by the semantic discrepancy across multi-modal tokens in the state-action-reward sequence, and the significant superiority of our token-wise MoE design.

Thank you for your insightful comment. In addition to Section 5 Experiments, we will include the above analysis in Section 1 Introduction, to further justify our motivation for the token-wise MoE.

---

Q3. Evaluate T2MIR on more challenging benchmarks such as the full ML1 or ML10.

A3. Following the main practice in literature [1-3], we adopt various benchmarks that are popular to evaluate in-context RL algorithms. Following your advice, we have conducted new evaluations on the challenging Ant-Dir environment in MuJoCo and the ML10 environment in Meta-World to demonstrate whether T2MIR can scale to these harder suites. The following table shows T2MIR’s promising scalability and consistent superiority over various baselines in these challenging benchmarks.

[1] In-context RL with algorithm distillation, ICLR 2023.

[2] Supervised pretraining can learn in-context RL, NeurIPS 2023.

[3] Distilling RL algorithms for in-context model-based planning, ICLR 2025.

---

Summary Response

Thank you for your valuable review comments, which help us gain more critical insights and further enhance our empirical evaluation. We are honored to have your positive comments on our motivation ("making the architectural choices easy to follow") and empirical evaluation ("the ablation study and t-SNE visualizations shed light on…").

Also, we are grateful that other reviewers (we refer to wTXA as R1, yZc9 as R2, and sbb8 as R3) make positive comments to our overall contributions, including the novelty and motivation (”the insight of…is interesting” by R1, "the idea of…sounds promising" by R2, and "propose a novel architecture" by R3), the empirical evaluation (”results demonstrate faster convergence and better rewards”, by R1, "a good set of experiments addressing various aspects of the work, which I appreciate" by R2, and "the experiments…are comprehensive, the ablations and analysis are insightful" by R3), and the writing (”the paper is fairly well-written” by R2). We hope that our contributions can be adequately justified.

We summarize that your main concerns may include the evidence of the two ICRL challenges and the evaluation on more challenging benchmarks. We have extended a number of justifications and experimental analyses to address your concerns. Please let us know if we have addressed your concerns. We are more than delighted to have further discussions and improve our manuscript. If our response has addressed your concerns, we would be grateful if you could re-evaluate our work.

---

Q1. Evaluation on other benchmarks like Ant-Dir and Humanoid, or text-based benchmarks like webshop.

A1. Following the main practice in literature [1-3], we adopt various benchmarks that are popular to evaluate in-context RL algorithms. Following your advice, we have conducted new evaluations on Ant-Dir in the MuJoCo benchmark and ML10 in the Meta-World benchmark to demonstrate whether T2MIR can scale to these harder suites. The following table shows T2MIR’s promising scalability and consistent superiority over various baselines.

[1] In-context RL with algorithm distillation, ICLR 2023.

[2] Supervised pretraining can learn in-context RL, NeurIPS 2023.

[3] Distilling RL algorithms for in-context model-based planning, ICLR 2025.

---

Q2. Analysis of how the behavior policy and task-related information contribute to the task representation, respectively.

A2. In our settings, we generate datasets by employing various checkpoints from different training steps, where one checkpoint (i.e., one behavior policy) may collect several trajectories. In order to extract the high-level representation from the gating output in task-wise MoE, we utilize contrastive learning to pull the trajectory representations from the same task closer while pushing those from different tasks farther apart. Intuitively, as the trajectories from the same task may be collected by different behavior policies, pulling these representations closer means forcing the representations to exclude information of behavior policy and focus on task-related information in the trajectory.

Further, we also conduct new experiments to assess T2MIR’s capacity to handle OOD offline data, with the results presented in A4 to Q4. Specifically, we employ a random policy to interact with the environment and collect one trajectory to serve as a prompt for both T2MIR-DPT and DPT. Despite the significantly low return of the random trajectory, T2MIR-DPT still gains impressive improvement than DPT.

---

Q3. Clarify the different modalities within the state-action-reward sequence in token-wise MoE.

A3. As stated in the paper (L139-142), the physical information contained within the state-action-reward sequence is different. The state is typically continuous in nature, indicating that $s_{t+1}$ evolves based on $s_t$, and contains some physical quantities of the agent (e.g., position, velocity, and acceleration). The action $a_t$ is less smooth (i.e., $a_{t+1}$ may be very different from $a_t$, although $s_{t+1}$ is similar to $s_t$), and may contain some mechanical control quantities, such as joint torques. The reward is a simple scalar that depends on state and action. Consequently, there exists a substantial semantic gap among these three elements in the sequence. Hence, we consider these three elements as different modalities and introduce a token-wise MoE to adaptively handle their distinct semantics. We will provide clearer clarifications on these separate modalities.

---

Q4. Demonstrate T2MIR’s robustness by evaluating on OOD offline data.

A4. In our experiments, we evaluate T2MIR and baselines via direct online interaction with the test environment, without access to any offline data. That is, the prompts or contexts are generated from online environment interactions to infer task beliefs during testing. Following your suggestion, we have conducted new experiments to evaluate T2MIR’s robustness to OOD offline data as the test prompt. Similar to the DPT setting where offline datasets are used as prompts, we employ a random behavior policy to interact with the environment and collect one trajectory as the offline prompt. The table below demonstrates T2MIR’s consistent superiority even when prompted with OOD offline data.

Furthermore, we conduct new experiments to evaluate T2MIR’s capability to handle OOD test tasks in Cheetah-Vel. The target velocities in training tasks are uniformly sampled from $U[0.075,3]$, but during testing we sample OOD tasks from $U[3.1,3.3]$. The result presented below demonstrates T2MIR’s consistent superiority even when dealing with OOD test tasks.

---

Summary Response

Thank you for your valuable review comments, which help us gain more critical insights and further enhance our empirical evaluation. We are honored to have your recognition on our method, especially on its novelty ("propose a novel architecture ") and empirical evaluation ("the experiments…are comprehensive, the ablations and analysis are insightful").

Also, we are grateful that other reviewers (we refer to wTXA as R1, yZc9 as R2, ptNc as R4) make positive comments to our overall contributions, including the novelty and motivation ("the insight…is interesting" by R1, "the idea of…sounds promising" by R2, and "making the architectural choices easy to follow" by R4), the empirical evaluation ("results demonstrate faster convergence and better rewards" by R1, "a good set of experiments addressing various aspects of the work, which I appreciate" by R2, and "the ablation study and t-SNE visualizations shed light on…" by R4), and the writing (”the paper is fairly well-written” by R2). We hope that our contributions can be adequately justified.

We summarize that your main concerns may include: 1) the evaluation on harder benchmarks, 2) deeper analysis on task representation learning, and 3) the robustness to OOD testing. We have extended a number of justifications and experimental analyses to address your concerns. Please let us know if we have addressed your concerns. We are more than delighted to have further discussions and improve our manuscript. If our response has addressed your concerns, we would be grateful if you could re-evaluate our work.

---

Q1. Clear formalism and specification of task variation, and the use of task variation parameterization in the context.

A1. We follow the general setting in in-context RL domains, where tasks are sampled from a distribution but task variations cannot be parameterized. This is a more realistic setting for practical applications, as we cannot access some oracle information about the underlying task (i.e., the goal position in a navigation environment). Appendix D.1 introduces the details about how to construct the multi-task environments and corresponding datasets. Taking Cheetah-Vel as an example, tasks differ in the goal velocity, and we uniformly sample the goal velocity from the distribution $U[0.075, 3.0]$ to construct the multi-task dataset. In Figure 1, we visualize the gating representations of tasks with different goal velocities to gain deep insights into the relationship between the learned gating mechanism and task similarity. But, we cannot access the oracle task information (i.e., goal velocity) during in-context training. Following your advice, we will make a clear formalism on task variation.

---

Q2. The necessity of learning the gating mechanism in MoEs, and its superiority over heuristic gating approaches.

A2. As mentioned in our A1 to Q1, since we cannot access the oracle task information, it is infeasible to manually route the task-wise MoE based on heuristics (e.g., splitting the goal velocity range). In the token-wise MoE, we can manually assign tokens of three modalities to three experts, i.e., one expert for one modality. We have conducted additional ablations to gain insights into the heuristic gating mechanism. The result in the following table shows the significant superiority of our gating mechanism over heuristic gating in the token-wise MoE. We conjecture that the improvement can come from the flexibility and adaptability of our end-to-end gating mechanism.

Further, we are striving to design more universal, end-to-end architectures, with minimal requirements on domain knowledge or manual interventions. This principle aligns with Sutton’s The Bitter Lesson: While task-specific designs may offer immediate gains, scalable approaches provide more sustainable long-term benefits.

---

Q3. Ablations to show the impact of balancing loss and contrastive loss.

A3. Following your suggestion, we have conducted new ablation studies to investigate the respective impacts of the two loss terms. We compare T2MIR to three new ablations:

• w/o balance_loss: omitting balance loss (Eq.3) in token-wise MoE, while retaining contrastive loss in task-wise MoE.
• w/o contrastive_loss: omitting contrastive loss (Eq.7) in task-wise MoE, while retaining balance loss in token-wise MoE.
• w/o aux_loss: omitting both balance loss and contrastive loss.

For all ablations, we keep all other hyperparameters the same as in the main experiments. The result in the following table shows that both the balance loss and contrastive loss play significant roles in T2MIR’s superior performance. Ablating the balance loss will cause a bit more performance degradation than ablating the contrastive one.

To further investigate the impact, we sample a batch of data and record the top-2 expert selections at the last training step. The percentages of experts being selected are shown in the following table. When ablating balance loss, 83% tokens activate experts 1 and 2; when ablating contrastive loss, all tokens activate experts 1 and 2. The result shows the homogeneous expert assignments without these losses. In contrast, in Figure 1, we can observe that T2MIR’s expert assignments are much more heterogeneous, demonstrating the effectiveness of the balance and contrastive losses for MoE training.

---

Q4. Training the contrastive loss when the number of tasks is large.

A4. According to the literature, a key ingredient in the success of many contrastive learning methods is using a larger batch size with sufficient negative samples. Only when the batch size is big enough, the loss function can cover a diverse enough collection of negative samples, challenging enough for the model to learn meaningful representation to distinguish different examples.

To gain more insights into the contrastive loss in T2MIR, we omit the token-wise MoE and investigate T2MIR’s performance with different numbers of training tasks (10, 20, and 45). The result in the following table shows that T2MIR’s performance increases as the number of training tasks increases. In the task-wise MoE, we treat samples from the same task as positive pairs and samples from other tasks as negatives. When the number of tasks increases, we can obtain a larger batch with more negative samples, contributing to more reliable contrastive learning.

---

Q5. When the token-wise MoE count is changed, how is the task-wise MoE count handled?

A5. The token-wise MoE and task-wise MoE are two parallel layers. When changing one side (e.g., the token-wise MoE count), we strictly keep the other side invariant (e.g., the task-wise MoE count or other hyperparameters) to investigate the robustness of each component.

Q6. How is the decision to activate one-third of experts in token-wise MoE and always 2 experts in task-wise MoE reached?

A6. Since there are three modalities of the state-action-reward tokens, we always activate one-third of experts for each specific token. For the simplest case, we have three experts and assign one expert to each modality. Hence, we will only activate one expert for routing each token. For task-wise MoE, we follow the common practice in popular works, such as Gshard [1], BASE layers [2], and MoE++[3], which usually use top-2 routing. We will provide clearer clarification about the MoE design.

[1] GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding, ICRL 2021.

[2] BASE Layers: Simplifying Training of Large, Sparse Models, ICML 2021.

[3] MoE++: Accelerating Mixture-of-Experts Methods with Zero-Computation Experts, ICLR 2025 Oral.

---

Summary Response

Thank you for your valuable review comments, which help us gain more critical insights and further enhance our empirical evaluation. We are honored to have your recognition on our method, especially on its novelty ("the idea of…sounds promising"), empirical evaluation ("a good set of experiments addressing various aspects of the work, which I appreciate"), and the writing (”the paper is fairly well-written”).

Also, we are grateful that other reviewers (we refer to wTXA as R1, sbb8 as R3, ptNc as R4) make positive comments to our overall contributions, including the novelty and motivation ("the insight…is interesting" by R1, "propose a novel architecture" by R3, and "making the architectural choices easy to follow" by R4), and the empirical evaluation ("results demonstrate faster convergence and better rewards" by R1, "the experiments…are comprehensive, the ablations and analysis are insightful" by R3, and "the ablation study and t-SNE visualizations shed light on…" by R4). We hope that our contributions can be adequately justified.

We summarize that your main concerns may include: 1) the task variation parameterization, 2) the necessity of learning the gating mechanism, and 3) the ablation on regularization losses. We have extended a number of justifications and experimental analyses to address your concerns. Please let us know if we have addressed your concerns. We are more than delighted to have further discussions and improve our manuscript. If our response has addressed your concerns, we would be grateful if you could re-evaluate our work.

---

Q1. The necessity of leveraging explicit architectural bias for in-context RL.

A1. In the era of large language models (LLMs), MoE is a promising architecture for managing computational costs when scaling up model parameters, such as in Gemini, Llama-MoE, and DeepSeekMoE. The architectural advancement also extends to various domains such as computer vision and image generation. Particularly, recent studies [1,2] show that incorporating MoE leads to substantial performance increases for DRL, providing strong empirical evidence towards developing scaling laws (a model’s performance scales proportionally to its size) for RL.

The success of LLMs comes with training huge models on a diversity of datasets, exhibiting the power and generalization of in-context learning. In-context RL aims to harness this potential to improve RL’s generalization to downstream tasks. Compared to single-task RL, it is more urgent for in-context RL to leverage explicit architecture bias to establish possible scaling laws that enable training of large, generalizable foundation models.

Further, in our A3 to Q3, we demonstrate that adding two MoE blocks per layer yields $11\sim 42$% performance increase at the cost of $10\sim 14$% running time increase, with the same number of total activated parameters as baseline models. In our A4 to Q4, we show that T2MIR is insensitive to hyperparameters related to MoE layers, incurring a minimum cost for hyperparameter tuning. In summary, we hope that the necessity and superiority of leveraging architectural bias for in-context RL could be sufficiently justified.

[1] Mixtures-of-experts unlock parameter scaling for deep RL, ICML 2024.

[2] Don't flatten, tokenize! Unlocking the key to SoftMoE's efficacy in deep RL, ICLR 2025.

---

Q2. The scalability of our method, with experiments on small-sized data with relatively small models.

A2. Following the common approach in literature to evaluate model scalability [1-3], we evaluate the scalability of T2MIR-AD, i.e., the performance with respect to the model size and the dataset size. The following table shows that introducing MoE while increasing model/dataset size can unlock further scaling potential of T2MIR.

Further, we have conducted new evaluations on the challenging Ant-Dir in MuJoCo and ML10 in Meta-World to demonstrate that our method can scale to these harder suites. The following table shows T2MIR’s promising scalability and consistent superiority over various baselines.

Following the main practice in literature [4-6], we adopt various benchmarks that are popular to evaluate in-context RL algorithms. Further, as stated in our limitations and future work, our method is evaluated on widely adopted benchmarks in ICRL, with relatively lightweight datasets compared to popular large models. An urgent improvement is to evaluate on more complex environments such as XLand-MiniGrid and XLand-100B with huge datasets, unlocking the scaling properties of MoE in ICRL domains.

[3] Network Sparsity Unlocks the Scaling Potential of Deep RL, ICML 2025.

[4] In-context RL with algorithm distillation, ICLR 2023.

[5] Supervised pretraining can learn in-context RL, NeurIPS 2023.

[6] Distilling RL algorithms for in-context model-based planning, ICLR 2025.

---

Q3. Complexity of adding two MoE blocks per layer for both training and inference.

A3. For the two MoE blocks, we keep the total activated parameters the same as those in a normal feed-forward network (FFN) block. During training, only the activated experts need to compute gradients, so the cost is approximately the same as FFN. During inference, the additional cost of MoE mainly includes computing the expert distribution by the router. The following table shows that introducing the MoE architecture only incurs a 10$\sim$14% increase in training time and a 14$\sim$27% increase in inference time, which is minor.

---

Q4. Hyperparameters and additional loss for balancing the experts.

A4. We have conducted hyperparameter analysis on MoE configurations in Section 5.3 and Appendix G.1. The results show that a moderate configuration is more suitable for token-wise MoE, and the performance slightly increases with more experts in task-wise MoE.

Following your suggestion, we have conducted new ablation studies to investigate the respective impacts of additional losses for balancing the experts. We compare T2MIR to three new ablations:

• w/o balance_loss: omitting balance loss (Eq.3) in token-wise MoE, while retaining contrastive loss in task-wise MoE.
• w/o contrastive_loss: omitting contrastive loss (Eq.7) in task-wise MoE, while retaining balance loss in token-wise MoE.
• w/o aux_loss: omitting both balance loss and contrastive loss.

For all ablations, we keep all other hyperparameters the same as in the main experiments. The result in the following table shows that both the balance loss and contrastive loss play significant roles in T2MIR’s superior performance. Ablating the balance loss will cause a bit more performance degradation than ablating the contrastive one.

To further investigate the impact, we sample a batch of data and record the top-2 expert selections at the last training step. The percentages of experts being selected are shown in the following table. When ablating balance loss, 83% tokens activate experts 1 and 2; when ablating contrastive loss, all tokens activate experts 1 and 2. The result shows the homogeneous expert assignments without these losses. In contrast, in Figure 1, we can observe that T2MIR’s expert assignments are much more heterogeneous, demonstrating the effectiveness of the balance and contrastive losses for MoE training.

---

Q5. Maintaining fairness regarding the parameter count in ablation studies.

A5. In ablation studies, we strictly keep the same number of activated parameters for the MoE and FFN architectures during training and inference, for a fair comparison.

---

Q6. Formatting Concerns about Fig. 2.

A6. Thanks a lot for your detailed suggestion, and we will improve the diagram and the text underneath.

---

Summary Response

Thank you for your valuable review comments, which help us gain more critical insights and further enhance our empirical evaluation. We are honored to have your positive comments on our motivation (”the insight of…is interesting”) and the empirical evaluation (”results demonstrate faster convergence and better rewards”).

Also, we are grateful that other reviewers (we refer to yZc9 as R2, sbb8 as R3, ptNc as R4) make positive comments to our overall contributions, including the novelty and motivation ("the idea of…sounds promising" by R2, "propose a novel architecture" by R3, and "making the architectural choices easy to follow" by R4), the empirical evaluation ("a good set of experiments addressing various aspects of the work, which I appreciate" by R2, "the experiments…are comprehensive, the ablations and analysis are insightful" by R3, and "the ablation study and t-SNE visualizations shed light on…" by R4), and the writing (”the paper is fairly well-written” by R2). We hope that our contributions can be adequately justified.

We summarize that your main concerns may include the necessity of leveraging architectural bias for ICRL, the model scalability, and the model complexity of our method. We have extended a number of justifications and experimental analyses to address your concerns. Please let us know if we have addressed your concerns. We are more than delighted to have further discussions and improve our manuscript. If our response has addressed your concerns, we would be grateful if you could re-evaluate our work.