Mixture-of-Experts Meets In-Context Reinforcement Learning
Authors: Wenhao Wu
Fuhong Liu
Haoru Li
Zican Hu
Daoyi Dong
Chunlin Chen
Zhi Wang
Venue: NeurIPS 2025 poster ()
Keywords: In-context reinforcement learning
mixture-of-experts
offline reinforcement learning
contrastive learning
transformers
Primary area: reinforcement_learning
TL;DR: We propose a token- and task-wise MoE structure to in-context RL models, harnessing architectural advances of MoE to unleash RL's in-context learning capacity.

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).

Mixture-of-Experts Meets In-Context Reinforcement Learning
Wenhao Wu' Fuhong Liu' Haoru Li'
Zican Hu' Daoyi Dong² Chunlin Chen' Zhi Wang
Nanjing University University of Technology Sydney (wenhaowu, chika, haoruli, zicanhu)Jemail. .nju edu cn daoyi. www.outs..eee.oooomooo edu.au
{clchen, zhiwang) @nju. edu. edu. cn
Abstract
In-context reinforcement learning (ICRL) has emerged as a promising paradigm for adapting RL agents to downstream tasks through prompt conditioning. How- ever, two notable challenges remain in fully harnessing in-context learning within RL domains: the intrinsic multi-modality of the vate-action-rewar 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-expeee -feeeeerss (MoE) into transformer-based decision models. T2MIR substitutes the feedforward layer with two parallel layers: 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, introduce contrastive learning method that maximizes the mutual information between the task and its router representa- tion, enabling 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 MoE ICRL,, offering simple and scalable architectural enhancement advance ICRL one step closer toward achievements in language and vision communities. Our code is available at https://ithubbcom///U//////
precise capture of task-relevant information. The outputs of
Introduction
Reinforcement learning (RL) is emerging a powerful mechanism for training autonomous agents to solve complex tasks in interactive environments 2]. unleashing its potential across frontier challenges including preference optimization 3. training diffusion models [4) and reasoning [5] 6
such as OpenAl-ol harness the in-context learning capabilities of the transformer architecture achieve substantial improvements in RL's adaptability to downstream tasks through prompt conditioning without any model updates, i.e., in-context RL (ICRL) [TO]. Current research in offline settings encompasses two primary branches, algorithm distillation (AD) [1 and decision-pretraiing
g and DeepSeek-RI ( Recent studies have been actively exploring how to
transformer (DPT) [2] owing their simplicity and generality. They share a common structure that uses across-episodic transitions as few-shot prompts to transformer policy, and following-up up studies continue to increase the in-context learning capacity by hierarchical decomposition
[3], noise distillation [4] model-
based planning
[5]. dynamic programming updates [16] 17 etc. 18
19 20
*Correspondence Zhi Wang <zhiwang nej..du.nu...........
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Despite these efforts. two notable challenges remain in fully harnessing in-context learning within decision domains, as RL notably more dynamic and complex than supervised learning. The first is the intrinsic multi-modality of datasets. In language or vision communities, the inputs to transformers are atomic words or pixels with consistent semantics in the representation space [21. In ICRI., the prompt inputs are typically transition samples containing three modalities of state, action, and reward with large semantic discrepancies. States like physical quantities are usually continuous in nature in RL, actions like joint torques tend to be more high-freequecy and less smooth, and rewards are simple scalars that can be sparse over long sequences [22]. The second is task diversity and heterogeneity, Compared supervised learning, RL models more prone to overfit the training set and struggle with generalization across diverse tasks [23]. The learning efficiency can be hindered by intrinsic gradient conflicts in challenging scenarios where tasks vary significantly (24] The above limitations raise a key question: Can we design a scalable ICRL. framework to tackle the multimodaality and diversified task distribution within single transformer, advancing RI''s in-context learning capacities one step closer to achievements in language and vision communities?
In the era of large language models (LLMs), the mixture-of-expert (MoE) architecture 22 has shown remarkable scaling properties and high transferability while managing computational costs, such as in Gemini (26), Llama-MoE [27,, and DeepSeekMoE (28] The architectural advancement
also extends to various domains such as computer vision (29]. image generation [30), and RL [31 32),
highlighting the significant potential and promise of MoE models. Intuitively, MoE architectures can serve as an encouraging remedy to tackle the two aforementioned bottlenecks in ICRL. First, MoE enables different experts to process tokens with distinet semantics more effectively B3 This is natural match for processing the multiple modalities within the state-action-reward sequence. Second, MoE can alleviate gradient conflicts by dynamically allocating gradients to specialized experts for each input through the sparse routing mechanism (34. This property is promising for handling a broad task distribution with significant diversity and heterogeneous complexity.
Inspired by this, we propose an innovative frame- work T2MIR (Token- and Task-wise MoE for In-context RL), which for the first time harnesses architectural advances of MoE to develop more scalable and competent ICRL algorithms. We sub- stitute the feedforward layer in transformer blocks with two parallel ones: a token-wise MoE and a task-wise MoE. First, the token-wise MoE layer is responsible for automatically capturing distinct semantic features of input tokens within the multi- modal state-action-reward sequence. We include load-balancti loss and an importance loss avoid tokens from all modalities collapsing onto minority experts. Second, the task-wise MoE layer is designed to assign diverse tasks to specialized experts with sparse routing, effectively managing broad task distribution while alleviating gradient conflicts. To facilitate task-wise routing, we intro- duce contrastive learning method to maximize the mutual information between the task and its router representation in the MoE layer, enabling more precise capture of task-relevant information. Finally, the outputs of the two parallel MoE components are concatenated and fed into the next layer.
stabe
Figure 1: t-SNE visualization of expert assign- ments on Cheetah- Vel where tasks differ in target velocities. Left: token-wise MoE enables differ- ent experts to process tokens with distinct seman- tics. Right: task-wise MoE effectively manages a broad task distribution, where the difference between expert assignments is positively related to the difference between tasks.
In summary, our main contributions are threefold:
We unleash RL's in-context learning capacities with a simple and scalable architectural enhance- ment. To our knowledge, we are the first to bring the potential and promise of MoE to ICRL.
We design a token-wise MoE to facilitate processing the distinct semantics within multi-modal inputs, and task-wise MoE to tackle a diversified task distribution with reduced gradient conflicts.
We build our method upon AD and DPT, and conduct extensive experiments on various benchmarks to show superiority over competitive baselines and visualize deep insights into performance gains.
For example, given two navigation tasks where the goals are in adverse directions, the single RL model ought to execute completely opposite decisions under the same states for the two tasks.

---

2 Related Work
In-Context RL. Tackling task generalization is long-standing challenge in RL. Early meth- ods address this via the lens of meta-RL, including the memory-based (e.g.. RL [5] and LLIRL [36 37 the optimization-based (e.g.. MAML
[38 and MACAW (3).. and the context-
based ge. VariBAD [40], Meta-DDT
etc. [42
43) With the emergence of transformers that
show remarkable in-context learning abilities 44 the community has been exploring its potential to
enhance RL's generalization via prompt conditioning without model updates [10].
Many ICRL methods have emerged, each differing in how to train and organize the context. In online settings, AMAGO 45 46 trains long-sequence transformers over entire rollouts with actor-critic learning, and 47 leverages the S4 (structured state space sequence) model to handle long-range sequences for ICRL tasks. Classical studies offline settings include AD [T] and DPT [1]. AD trains a causal transformer to predict actions given preceding histories as context, and DPT predicts the optimal action given a query state and a prompt of interactions. Follow-up studies enhance in-context learning from various algorithmic aspects [1, 18
20 16, 17], such as IDT with hierarchical 15 Complementary to these algorithmic studies, we introduce MoE to advance ICRL's development from an architectural perspective.
19
structure 133 and DICP with model-based planning
Mixture-of-Expertt The concept of MoE first proposed in 488] which employs different expert networks plus gating network that decides which experts should be used for each training case (49).
Then, it applied to tasks of language modeling and machine translation (25), gradually showing prominent performance in scaling up transformer model size while managing computational costs [50]
ST. Up to now, MoE architectures have become a standard component for advanced LLMs, such as Gemini (26). Mama-MoFF [27] and DeepSeekMoE [28]. Recent efforts have explored extending MoE advancements to RI. domains.
31 incorporates soft MoE modules into value-based RI. agents,
and performance improvements scale with the number of experts used. (2 uses an MoF backbone with a CNN encoder to process visual inputs and improves the policy learning ability to handle complex robotic environments. (34] strengthens decision transformers with MoE to reduce task load on parameter subsets, with a three-stage training mechanism to stabilize MoE training. Accompanied by these encouraging efforts, we aim to harness the potential of MoE for ICRL.
3 Preliminaries
3.1 In-Context Reinforcement Learning
We consider multi-task offline RL setting, where tasks follow a distribution M, = (S,A, R,, 7) ~ P(M). Tasks share the same state space and action space A, while differing in the reward function R or transition dynamics T. An offline dataset D== D"=((s)) [(x]] [ af, Tr 8,"" "))) is collected by arbitrary behavior policies for each task out of a total of N training tasks. The agent can only access the offline datasets [D") to train an in-context policy as क prompt that encodes task-relevant information. At test time, the trained policy is evaluated on unseen tasks sampled from P(M) by directly interacting with the environment. The prompt is initially empty and gradually constructed from history interactions. The objective is to learn an in-context policy to maximize the expected episodic return over test tasks as ((T) = EMr- (T)).
T ) where T is a trajectory
3.2 Mixture-of-Experte
standard sparse MoE layer consists of K experts {FI,..,, Eg and a router G. The router predicts an assignment distribution over the experts given the input Jr Following the common practice [2] 32 we only activate the top-k experts to process the inputs. general, the number of activated experts is fixed and much smaller than the total number of experts, thus scaling up model parameters and significantly reducing computational cost. Formally, the output of the MoE layer can be written as
w x) = softmax (Gz)) == Σx 11 w(t;))E
where topk(.) selects top k experts based on the router output G(x). softmax(-) normalizes top-k values into weight distribution w(-). Probabilities of non-selected experts are set to 0.
3

---

T2MIR FrameWork
To(a 8: Tpro - (())
expert,
expert
expert Kg
Action Head
19aa
expert]
expert,
experty,
Lecotesstive
Router
Routertaken
Token-MoE
Task-MoE
Layer Norm
PT
Countrattiv
(b) Token-wise MoE
Multi-Head Attention
Concatenation
Layer Norm
Router
Momentum Router rack
+
Residual connection
Embedding
Higher gate weights
(5o, ao, ST. or, or,
M
M
& M₁
Lower gate weights
(a) Overall pipeline
Task-wise MoE
The overview of T2MIR. (a) Overall pipeline: we substitute the feedforward layer in causal transformer blocks with two parallel MoE layers and concatenate their outputs feed into the next layer. (b) Token-wise MoE: it automatically captures distinct semantic features within the multi-mode modall (s, a, inputs, and uses Lealance as regularization loss avoid tokens from all modalities collapsing onto minority experts. (c) Task-wise MoE: it assigns diverse tasks to specialized experts, and includes a contrastive learning loss Lcontrastive to enhance task-wise routing via more precise capture of task-relevant information, where т; is the query and Ti /Ty are positive/negative keys.
Figure 2:
Method
In this section, we present T2MIR (Token- and Task-wise MoE for In-context RL), an innovative framework that harnesses architectural advancements of MoE to tackle ICRL. challenges of the multi-modality and diversified task distribution. Figure @illustrates the method overview, followed by detailed implementations. Algorithm pseudocodes are presented in Appendix A
4.1 Token-wise MoE
By interacting with the outer environment, we gather the data represented as a sequence state- action-reward transitions in the form of (so, ao; To, 8T, dr,, rr), which often serves as the prompt input to ICRL models. The state is usually continuous in nature, which can contain physical quantities of the agent (e.g., position, velocity, and acceleration) or an image in visual RL (32). The action tends to be more high-frequency and less smooth, such as joint torques or discrete commands. The reward is simple scalar that often sparse over long horizons. In language or vision communities, the inputs to transformers are atomic words or pixels with consistent semantics in the representation space 21). In contrast, ICRL encounters new challenge of processing the prompt data that encompasses three modalities with significant semantic discrepancies.
Inspired by the surprising discoveries that the MoE structures effectively capture distinct semantic features input tokens [50] 33,, we introduce a cosen-wise MoE layer to tackle the multiple modalities within the state-actions seeeeacction-eeaaan sequence. As shown in Figure (b), each element in the multi-modal sequence corresponds to a token. The MoE consists of K experts Ens and a router G and the router learns assign each token specific experts with noisy top-k gating. Let l denote the hidden state of a given token after self-attention calculation. The token router Gook computes a distribution of logits (Grok (i) h)) K for expert selection. The top-k experts will be selected to process this token, and their outputs will be weighted to produce the final output yhok as
Week (i; fi) softmax (topk (topk(Giokl (i]h))) [i],
Vlok >>'
Wook (i: h) Evo((i/h)
(2)
The router tends to converge on producing large weights for the same few experts, as the favored experts are trained more rapidly and thus are sclected even more by the router. In order to avoid

---

tokens from all modalities collapsing onto minority experts, we incorporate a regularization term to balance expert utilization that consists of an importance loss and a balancing loss [25] as
Lbllance = Wimp CV (Imp(h))² + Whond CV (L))
(3)
where CV(-) denotes the coefficient of variation, and Wimp and Wlond are hand-tuned scaling factors. The first item encourages all experts to have equal importance as defined by Imp(h) = Σi, Guak (-14).
The second item encourages experts to receive roughly equal numbers of training examples, where Load(-) is smooth estimator of the number of examples assigned to each expert for a batch of inputs. more detailed description of the balance loss can found in Appendix
4.2 Task-wise MoE
In typical ICRL settings, a single policy model trained across multiple tasks. The learning efficiency can be impeded by intrinsic gradient conflicts in challenging scenarios with significant task variation. For instance, given two navigation tasks where the goals are in adverse directions, the policy ought to make contrary decisions under the same states. Given the same state-action trajectories, the two tasks can produce exactly opposite policy gradients, leading to a sum of zero gradient during training.
The MoE structure originally proposed to use different parts of model, called experts, to specialize in different tasks or different aspects of a task 48 49 Drawing from this natural inspiration, we introduce task-wise MoE layer handle broad task distribution with significant diversity and heterogeneous complexity. leveraging modular expert learning to alleviate gradient conflicts. As shown in Figure (c), the MoE consists of K₂ experts Eusk and router Gusk. and the router learns to assign a sequence of tokens to specialized experts at the task level. Given a trajectory
7- 0o, To,... .... 8T, aT, rT) from some task M, we use l to denote the average hidden state of all three-modality tokens after self-attention calculation. The task router G computes a distribution of logits [Gusk h)) [==1 for expert selection. The top-k experts will be selected to process these tokens,
and their outputs will be weighted produce the final output Ylask as
Waisk (i; h) softmax (oopk(GG.ss [.]
Mlask Σ
K2 Di=1
Wrisk (i; h)) Emss (ilh).
Task-wise Routing via Contrastive Learning. The task-wise router can be viewed as an encoder for extracting task representations from input tokens, based the intuition that similar tasks are preferred to the same expert and tasks with notable differences are allocated to disparate experts. An ideally discriminative task representation should accurately capture task-relevant information from offline datasets, while remaining invariant to behavior policies or other unrelated factors. To achieve this goal, we propose to maximize the mutual information between the task and its router representation, enhancing task-wa wise routing that minimally preserves iisk-irrelevant information.
Formally, the mutual information 7(.) aims to quantify the uncertainty reduction of one random vari- able when the other observed, measuring their mutual dependency. Mathematically, we formalize the router Guskk probabilistic encoder ~~~ Gusk (-Ih), where denotes the task representation. The distribution is determined by its task M, where M ~ P(M). The objective for the router is:
p(M[z)]
max (2; M) E2,M
log
p(M)
Drawing inspiration from noise contrastive estimation (InfoNCE) [52 in the contrastive learning literature [53 (4], we derive lower bound of Eq. (5) using the following theorem.
Theorem 1. Let M denote a set of tasks following the task distribution P(M), and |M| = N.
M M is a given task. Let h- 2~ Grakk (-Ih,, and e(h, 2) = p(=la)
p(z) where T is a trajectory
from task M and h is average hidden state of all tokens after the self-attention calculation function f(). Let ht' denote the average hidden state generated by any task M' EM, then we have
e(h,2)
M2;; 2 log N E,, log
(6)
EMvEme e(h',
Appendix C presents the complete proof. Since direct evaluation of p(2) or p(2[h) is intractable, we employ NCE and importance sampling techniques that compare the target value with randomly
5

---

sampled negative ones. Consequently, we approximate using the exponential of a score function, similarity metric between the latent codes of two examples. Given a batch of m trajectories (Ti..., Tm), we denote router representation ž from тi as the query q. and representations of other trajectories as the keys (=)...... 2m 24 Points from the same task with the query, i M,, are set as positive key {k+, and those from different tasks are set as negative {k_ } IK [k+]- We adopt bilinear products [53] as the score function, with similarities between the query and keys computed as Wk, where W is learnable parameter matrix. Then, we formulate sampling-basee version
of the tractable lower bound as the InfoNCE loss for the contrastive router as
exp(q T W
Er exp(2, 22-))
Ccontrattvve = log exp(q +ex)) Wk-)
log i*EM;
Eexp( W2.) E Eexp(2, W2y) (7) P'MM
optimizes an N-way classification loss classify the positive pair among all pairs, equivalent to maximizing lower bound on mutual information in Eq. (6) with this bound tightening as N gets larger. Following MoCo (55), we maintain a query router G and key router G and use a momentum update with coefficient / E (0,1) to ensure the consistency of the key representations as
G + 3G task + (1 G)
(8)
where only query parameters Gosk are updated through backpropagation.
4.3 Scalable Implementations
We develop scalable implementations of T2MIR using two mainstream ICRL backbones that offer
promising simplicity and generality: AD [ and DPT 12 yielding two T2MIR variants as follows.
T2MIR-AD. It trains an MoE Mog-auggeet causal transformer policy Ty to autoregressively predict actions given preceding learning histories his, the prompt:
his; == (8o: 0o, ro. 4 St, az, r²) = (ggt;
(9)
During training, we minimize negative log-likelihood loss over N individual tasks
[( = Σ N
log my (a = ai sp; -po); where Tpro = his
(10)
which can be derived as a cross entropy loss for discrete actions, and a mean-square error (MSE) loss when the action space is continuous. Akin to AMAGO [45) we adopt Flash Attention (56) to enable long context lengths on single GPU. For the context, we use the same position embedding for 8t. ag and rt to maintain semanties of temporal consistency.
T2MIR-DPT. It trains a transformer policy to predict an optimal action given query state and prompt of interaction transitions. In practice, we randomly sample trajectory Tpro from the dataset D as the task prompt for each query-label pair (st, ai), and train the policy to minimize the loss as
(()) - N
T
log T (a a, s
T where pm D
(11)
Different from AD in Eq. (T). DPT requires labeling optimal actions for query states to construct the offline dataset Σn (s( ", inheriting principles from pure imitation learning.
5 Experiments
We comprehensively evaluate our methods on popular benchmarking domains using various qualities of offline datasets. In general, aim to answer the following research questions:
Can T2MIR demonstrate consistent superiority on in-context learning capacity when tested in unseen tasks, compared to various strong baselines? (Sec.5.)
How do the token-wise and the task-wise MoE layers affect T2MIR's performance, respectively We also gain deep insights into each component via visualizations in Sec. 5.4
(Sec. 5.2 ?
How robust is T2MIR across diverse settings? We construct offline datasets with varying qualities to extensively evaluate T2MIR and baselines. Also, we investigate T2MIR's sensitivity to various hyperparameters, such the expert configuration and the loss ratio. (Sec. 5.3 and Appendix G

---

Environments. We evaluate T2MIR on four benchmarks that are widely used in multi-task settings: i) the discrete environment DarkRoom, which is grid-world environment with multiple goals: ii) the 2D navigation environment Point-Robor which aims to navigate to target position in 2D) plane; iii) the multi-task MuloCo locomotion control environment, containing Cheetah-Vel and Walker-Param, iv) the Meta- Meta-World manipulation platform, including Reach and Push. We construct three datasets with different qualities: Mixed, Medium -Expert,, and Medium. Appendix presents more details about environments and dataset construction.
Baselines. We compare T2MIR five competitive baselines, including four ICRL methods AD [[]. DPT [ IDT [3), DICP [1], and context- based offline meta-RL. approach UNICORN [42]. More details about baselines are given in Appendix [E] results among them in our experiments.
As DICP three implementations, we use the best
For all experiments, we conduct evaluations across five random seeds. The mean of the obtained return is plotted as the bold line, with 95% bootstrapped confidence intervals of the mean results indicated by the lighter shaded regions. Appendix gives implementation details of T2MIR-AD and T2MRR-PP.. Appendix G presents comprehensive hyperparameter analysis, including expert selection in token-wise and task-wise MoE, InfoNCE loss ratio, and the positioning of MoE layers.
5.1 Main Results
compare our method against various baselines under an aligned evaluation setting, where prompts or contexts are generated from online environment interactions to infer task beliefs during testing. Figure and Table T present test return curves and numerical results the best performance of T2MIR and baselines using Mixed datasets. In these diverse environments with varying reward functions or transition dynamics, both T2MIR-AD and T2MIR--DP achieve significantly superior performance regarding the in-context learning speed and final asymptotic results compared to the strong baselines. In most cases, our two implementations take the top two rankings for the best and best performance, showcasing comparably strong generalization abilities unseen tasks. It highlights the effectiveness of introducing MoE architectures to the two mainstream ICRL backbones. Notably, DICP exhibits faster learning speed and comparable asymptotic performance compared to our T2MIR in Dork In this environment with a small state-action space, DICP naturally enables efficient policy search via look-ahead model- planning. Then, DICP struggles when facing more complicated environments like Reach and Push, as it is more difficult to find high-return trajectories when planning in large space. contrast, the superiority of T2MIR is more pronounced in harder problems, underscoring its potential to tackle challenging ICRL tasks.
T2MIR-AD (Ours)
T2MIR-DPT DPT (Ours)
AD
DPT
DICP
IDT
UNICORN
DarkRoom
Cheetah-Vel
Reach
850
80
-50
800
Returs
60
Returns
-100
Returns
750
40
20
-150
700
50
100
-
200
10
20
650
30
In-Context Episodes
In-Context Episodes
In-Context Episodes
Point-Robot
Walker-Param
Push
700
-5
500
650
Returns
-6
Returns
Returns
600
550
-7
300
-8
450
-9
200
400
10
20
30
0
10
20
30
40
10
20
30
In-Context Episodes
In-Context Episodes
In-Context Episodes
Figure 3: Test return curves of two T2MIR implementations against baselines using Mixed datasets.
7

---

Table 1: Test return of T2MIR against baselines using Mixed datasets, i.e., numerical results of the best performance from Figure
Best result in bold and second best underline.
Environments
T2MIR-AD
T2MIR-DPT
DPT
DICP
IDT
UNICORN
DarkRoom
90,9+000
78.3+14
57.5=5.6
72.5=13.1
90,9+00
49.2=28
Point-Robot
-5.2+03
-5.0=0.
-5.6+044
-5.8+03
-5.3402
-5.5+00
-6.6+006
Cheetah Vel
-68.9+21
-43.2=08
-119.24:30.
-56.1+10
-74.9+32
-82.8+30
88288
Walker-Param
435.7=72
492.8
395.217.1
425.919.7
403.7:154
429.4158
372.81246
Reach
823.9.7.1
764.6193
790.4120.2
754.91180
775.311.
808.5136
748.8401
Push
665.811139
633.2186
604.7163
615.4160
560.0164.2
620.624.
547.41612
5.2
Ablation Study
We compare T2MIR to three ablations: 1) w/o Task-MoE, it only retains the token-wise MoE layer and regularizes the router with balance loss in Eq. (3); 2) w/o Token-Mon MMOE, it only retains the task- wise MoE layer and enhances the router with contrastive loss in Eq. 7 and 3) w/o all, degrades to vanilla AD and DPT. We keep the same amount of activated parameters in our T2MIR and three ablations for fair comparison. Figure 4 and Table of ablation study on T2MIR-AD and 12MIR-DPT.
2 present test return curves and numerical results
First, both T2MIR architectures show decreased performance when any MoE is removed, with the worst occurring when both MoE layers are excluded. It demonstrates that the two kinds MoE designs essential for T2MIR's capability and they complement each other. Second, with the AD backbone, ablating the token MoE incurs a more significant performance drop than ablating the task MoE, indicating that the token MoE plays training histories as context, which can contain redundant trajectories for identifying the underlying task, akin the memory-based meta-RL. approach RL,2 (35 Hence, using the token MoE to handle multi-modality within the long-horizon context may yield greater benefits. Third, the situation is reversed with the DPT backbone, the task MoE. is more essential for T2MIR-DPT. DPT takes small number of transitions as the task prompt, placing greater emphasis on accurately identifying task-relevant information from limited data.
vital role for T2MIR-AD. AD takes long
T2MIR
w/b Task MoE
w/a Token-MoE
w/o all
Walker-Param
Push
-60
700
440
675
420
Returns
-100
Returns
Returns
-120
400
-140
380
600
-160
360
575
550
10
20
30
10
30
30
In-Context Episodes
In-Context Episodes
In-Context Episodes
(a) T2MIR-AD
Cheetah-Vel
Walker-Param
Push
-40
500
640
Returns
Returns
Returns
620
400
600
-70
350
580
In Context Episodes
In Context Episodes
In Context Episodes
(b) T2MIR-DPT
Figure 4: Ablation results of both T2MIR-AD AD and T2MIR-DPT architectures using Mixed datasets.

---

Table 2: Numerical results of ablation study both T2MIR-AD and T2MIR-DPT using Mixed datasets, i.., the best performance from Figure Best result in bold.
T2MIR-AD
T2MIR
w/o Task-MoE
w/o Token-MoE
Cheetah Vel
-68.942.1
-36.1=3.7
101.1+2997 410.8#t.11
319.2+3304
Warker-Param
435.7=722
418.6+3.9
395.2+7.1
Push
665.8+13.9
656.8+178
634.0+40
604.7+6.3
T2MIR-DPT
T2MIR
w/o Task-MoE
w/o Token- Moee
DPT
Cheetah- Vel
-3..2.08
-49.5334
-47 1.6
56.1 11D
Walker-Param Param
492.8+293
++6.
475.3+32.5
425.95 225.9+9.7
Push
633.2
627.2+400
628.0+7.
615.4+6.0
T2MIR-AD (Ours)
T2MIR-DPT (Ours)
AD
DPT
DICP
IDT
UNICORN
Cheetah-Vel
Push
Cheetah-Vel
Push
50
-50
600
600
Returns
-100
Returns
Returns
-100
Returns
400
400
-150
-150
-200
0
10
20
30
200
10
20
30
-20
0
10
20
30
200
10
20
30
In-Context Episodes
In-Context Episodes
In-Context Episodes
In-Context Episodes
(a) Medium-Exper
(b) Medium
Figure 5: Test return curves of T2MIR against baselines using Medium-Expper and Medium datasets.
Table 3: Numerical results of T2MIR against baselines using Medium-Expert (top) and Medium (bottom) datasets, i.e., numerical results of best performance from Figure 5 second best underline.
Best result in bold and
Medium Medium-Exppert 12MIR-AD T2MIR-DPT
DPT
DICP
IDT
UNICORN
Cheetah-Vel
-5
-56.0168
-130.31242
-6..3106
-73.4116 184.11144 -80.61228
Push
654.4
586.3195
595.4135
573.6185
485.5129.2
632.917.1
513.7151
Medium
T2MIR-AD T2MIR-DPT
AD
DPT
DICP
IDT
UNICORN
Cheetah- Vel Push
76..6+114
-63.3++40
-135.8+192
-67,9+322
75.6+25
101.3++4
105.9+37
600.5168
530.0167
577.74168
514.2+127
416.94309
554.3498
402.441629
5.3 Robustness Study
Robustness Dataset Qualities. We evaluate T2MIR on Medium-Expert and Medium datasets. and Table both T2MIR-AD and T2MIR-DPT exhibit superior performance
As shown in Figure and over baselines, validating their robustness when learning from varying qualities of datasets. Notably, most baselines suffer significant performance drop when trained on Medium datasets, while T2MIR maintains satisfactory performance despite the lower data quality. highlights T2MIR's 's appealing applicability in real-world scenarios where agents often learn from suboptimal data.
Robustness to Expert Configurationss We vary expert configurations in both MoE layers. For token-wise MoE that manages tokens of different modalities, we always activate one-third of the experts there three modalities and vary the total number of experts. Results in Figure 6 (a) and Figure (a) show that a moderate configuration (2/6) yields the best performance. For tisk MoE that manages task assignments, of experts, following the common configuration in popular works. When testing the impact of expert configuration one MoE structure, we keep the expert configuration of the other one unchanged. Results in Figure (b) and Figure (b) show that the performance slightly increases with more experts in total, and tends saturate soon.
always activate two experts while varying the total number
5.4 Visualization Insights into MoE Structure

---

1/3
2/6
4/12
2/4
2/6
2/8
Point Robot
Cheetah Vel
Point Robat
Cheetah Vel
-50
-50
-75
Returns
-6
Returns
Returns
-
Returs
-100
100
-8
88
-125
-125
-100
-150
1
10
-150
In-Context Episades
Episodes
Episodes
In-Context Episodes
(a) Token-wise MoF
Task-wise MoF
Figure 6: Analysis results of the number of selected experts against the total in token- and task-wise MoE on T2MIR-AD. For example, 2/6 denotes selecting the top 2 from 6 experts.
1/3
2/6
4/12
2/4
2/6
2/8
Point-Robot
Cheetah-Vel
Point-Robot
Cheetah-Vel
-50
-50
Return
-6
Returns
Returns
-6
Returns
-75
-75
-8
-8
-100
1000
-100
110
In-Context Episodes
In-Context Episodes
In-Context Episodes
In-Context Episodes
(a) Token-wisc MoE
Task-wisc MoE
Figure Analysis results of the number of selected experts against total experts in token- and task-wise MoE on T2MIR-DPP For example, 2/6 denotes selecting the top 2 from 6 experts.
Modality Clustering in Token-wise MoE. In Figure (left, the spatial proximity indicates similarity in expert assignments. It ex- hibits a clear clustering pattern that tokens from different modalities routed to distinct experts, highlighting the successful utilization of the MoE structure to process tokens with distinct semantics.
IV
02
03
00
Task Clustering in Task-wise MoE. In Figure (right), points of similar color come from similar tasks. It forms clear pattern in the router representation space, where similar tasks cluster closely and distinct tasks are widely separated. This confirms the effective use the MoE structure to manage broad task distribution. Appendix J gives more visualization on the two MoE layers.
m
-0
Figure 8: Cosine similarity of gradients between Point-
Robot tasks in four quadrants (I-IV), comparing T2MIR-AD (MoE) with AD (MLP).
Gradient Conflict Mitigation. In Figure g AD (lower triangle) shows significant gradient conflict between opposing tasks (e.g.. vs. III, deep blue with a large negative cosine value) and fails to discriminate tasks (e.g., vs. IV, deep red with large positive cosine value). In contrast, T2MIR-AD (upper triangle) maintains nearly orthogonal gradients across diverse tasks, especially for opposing ones ge.g., I vs. or vs. IV, cosine value near 0). This comparison verifies T2MIR's advantage in both gradient conflict mitigation and task discrimination.
Conclusions, Limitations, and Future Work
In the paper, we introduce architectural advances MoE to address input multi-modality and task diversity for ICRL. We propose scalable framework, where token-wise MoE processes distinct semantic inputs and a task-wise MoE handles broad task distribution with contrastive routing. Improvements in generalization capacities highlight the potential impact of our method, and deep insights via visualization validate that our MoE design effectively addresses associated challenges.
Though, our method is evaluated on widely adopted benchmarks in ICRL, with relatively lightweight datasets compared to popular large models, An urgent improvement
evaluate on more complex [57 10 with huge datasets, unlocking the scaling properties langu-lagagggaatttnat (VLA) that naturally involve more complex input multi-modality and task diversity.
environments such XLand-MiniGin of MoE ICRL domains. Another step to deploy our method
tasks [58 (59
10

---

Acknowledgements
This work was supported part by the National Natural Science Foundation of China (Nos. 62376122 and 72394363), and in part by the AI & Al for Science Project of Nanjing University.
References
[1] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, et al. Human- level control through deep reinforcement learning. Nature, 518(7540):529-533 2015.
[2] Alhussein Fawzi, Matej Balog, Aja Huang. Thomas Hubert, Bernardino Romera-Paredess al. Discovering faster matrix multiplication algorithms with reinforcement learning, Nature, 610(7930):47-53, 2022.
[3] Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. In Advances in Neural Information Processing Systems, volume 36, pages 53728-53741, 2023.
[4] Kevin Black, Michael Janner, Yilun Du, llya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. In Proceedings of International Conference on Learning Representations, 2024.
[5] Ganqu Cui, Yuchen 7hang, Jiacheng Chen, Lifan Yuan, Zhi Wang, Yuxin Zuo, Haozhan Li, Yuchen Fan, Huayu Chen, Weize Chen, et al. The entropy mechanism of reinforcement learning for reasoning language models. arXiv preprint MXiv:2555.226117 2025.
[6] Jianhao Yan, Yafu Li, Zican Hu, Zhi Wang, Ganqu Cui, Xiaoye Qu, Yu Cheng, and Yue Zhang. Learning to reason under off- off-policy guidance. Advances in Neural Information Processing Systems, 2025.
[7] Zican Hu, Wei Liu, Xiaoye Qu, Xiangyu Yue, Chunlin Chen, Zhi Wang, and Yu Cheng. Divide and conquer: Grounding LLMs as efficient decision-making agents via offline hierarchical reinforcement learning. In Proceedings of International Conference on Machine Learning. 2025.
[8] Aaron Jaech, Adam Kalai, Adam Leter, Adam Richardson, Ahmed El-Kishky, et al. OpenAI ol system card. arXiv preprint arXiv:24/2.16720, 2024.
[9] Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song. Ruoyu Zhang, et al. DeepSeek- R1: Incentivizing reasoning capability lims via reinforcement learning. arXiv preprint arXiv:2501. 2025.
[10] Alexander Nikulin, Ilya Zisman, Alexey Zemtsov, Viacheslav Sinii, Vladislav Kurenkov, and Sergey Kolesnikov. XLand- 1000: A large-scale multi-task dataset for in-context reinforcement learning. In Proceedings International Conference Learning Representations, 2025.
[1I] Michael Laskin, Luyu Wang, Junbyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steiger- wald, DJ Strouse, Steven Stenberg Hansen, Angelos Filos, Ethan Brooks, et al. In-context reinforcement learning with algorithm distillation. In Proceedings of International Conference on Learning Representations, 2023.
[12] Jonathan Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, and Emma Brunskill. Supervised pretraining can Iearn in-context reinforcement learning. In Advances Neural Information Processing Systems, volume 36, pages 43057-43083, 2023.
[13] Sili Huang, Jifeng Hu, Hechang Chen, Lichao Sun, and Bo Yang. In-context decision trans- former: Reinforcement learning via hierarchical chain-of-thought. In Proceedings of Interna- tional Conference on Machine Learning, 2024.
[14] llya Zisman, Vladislav Kurenkov, Alexander Nikulin, Viacheslav Sinii, and Sergey Kolesnikov. Emergence of in-context reinforcement learning from noise distillation. In Proceedings of International Conference Machine Learning, pages 32832-62846 2024.

---

[15] Jachyeon Son, Soochan Lee, and Gunhee Kim. Distilling reinforeement learning algorithms for in-context model-based planning. In Proceedings of International Conference on Learning Representations, 2025.
[16] Denis Tarasov, Alexander Nikulin, Ilya Zisman, Albina Klepach, Andrei Polubarov, Lyubaykin Nikita, Alexander Derevyagin, Igor Kiselev, and Vladislav Kurenkov. Yes, Q-learning helps offline in-context RI.. In ICLR 2025 Workshop on Scalable Optimization for Efficient and Adaptive Foundation Models, 2025.
[17] Jinmei Liu, Fuhong Liu, Jianye Hao, Bo Wang. Huaxiong Li, Chunlin Chen, and Zhi Wang. Scalable in-context Q-learning arXiv preprint (rXiv:2306001299 2025.
[18) Viacheslav Sinii, Alexander Nikulin, Vladislav Kurenkov, Ilya Zisman, and Sergey Kolesnikov. In-context reinforcement learning for variable action spaces. In Proceedings of International Conference on Machine Learning, pages 45773-45793, 2024.
[19] Juncheng Dong, Moyang Guo, Ethan X Fang, Zhuoran Yang, and Vahid Tarokh. In-context reinforcement learning without optimal action labels. In ICML 2024 Workshop on In-Context Learning, 2024.
[20] Zhenwen Dai, Federico Tomasi, and Sina Ghiassian. In-context exploration-explootaaion for re- inforcement learning. In Proceedings of International Conference on Learning Representations, 2024.
[21] loannis Kakogeorgiou, Spyros Gidaris, Bill Psomas, Yannis Avrithis, Andrei Bursuc, Konstanti- nos Karantzalos, and Nikos Komodakis. What hide from your students: iveettionguided masked image modeling. In Proceedings of European Conference on Computer Vision, pages 300-318, 2022.
[22] Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-making? In Proceedings of International Conference on Learning Representations, 2023.
[23] Robert Kirk, Amy Zhang. Edward Grefenstette, and Tim Rocktäschel. A survey of zero-shot generalisation in deep reinforcement learning. Journal of Artificial Intelligence Research, 76:201 264,, 2023.
[24] Siao Liu, Zhaoyu Chen, Yang Liu, Yuzheng Wang, Dingkang Yang, et al. Improving generaliza- tion in visual reinforcement learning via conffict-aware gradient agreement augmentation. In Proceedings of IEEE/CVF International Conference on Computer Vision, pages 2023.
[25] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely- peateely-gated mixture-du mixture-of-experts of-exts layer. In Proceedings of International Conference on Learning Representations, 2017.
[26] Gemini Team, Petko Georgiev, Ving Ian Lei, Ryan Burell, Libin Bai, Anmol Gulati, Garrett Tanzer, Damien Vincent, Zhufeng Pan, Shibo Wang, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint 20X5v:2403.05550, 2024.
[27] Tong Zhu, Xiaoye Qu, Daize Dong, Jiacheng Ruan, Jingqi Tong, Conghui He, and Yu Cheng. LLaMA-MoE: Building mixture-of-expert from llama with continual pre-training. In Proceed- ings of Conference on Empirical Methods in Natural Language Processing, pages 15913-15923, 2024.
[28) Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, et al. DeepSeek Towards ultimate expert specialization in mixture-of-experts language models. arXiv preprint arXiv:2401.06066, 2024.
[29] Zhiwen Fan, Rishov Sarkar, Ziyu Jiang, Tianlong Chen, Kai Zou, et al. MDVIT: Mixture-of- experts vision transformer for efficient multi-task learning with model-accelerator co-design. In Advances in Neural Information Processing Systems, volume 35, pages 28441-28457, 2022.
12

---

[30] Zeyue Xue, Guanglu Song, Qiushan Guo, Boxiao Liu, Zhuofan Zong, Yu Liu, and Ping Luo. RAPHAEL: Text-to-image generation via large mixture of diffusion paths. In Advances in Neural Information Processing Systems, volume 36, pages 41693-41706, 2023.
[31] Johan Samir Obando Ceron, Ghada Sokar, Timon Willi, Clare Lyle, Jesse Farebrother, Jakob Nicolaus Foerster, Gintare Karolina Dziugaite, Doina Precup, and Pablo Samuel Castro. Mixtures of experts unlock parameter scaling for deep RI.. In Proceedings of International Conference Machine Learning, pages 38520-38540, 2024.
[32] Suning Huang. Zheyu Zhang, Tianhai Liang. Yihan Xu, Zhehao Kou, Chenhao Lu, Guowei Xu, Zhengrong Xue, and Huazhe Xu. MENTOR: Mixture-of -f-peertss network with cask-triented perturbation for visual reinforcement learning. arXiv preprint arXiv:2410. /4972, 2024.
[33] Jihai Zhang, Xiaoye Qu, Tong Zhu, and Yu Cheng. CLIP-MoE: Towards building mixture of experts for clip with diversified multiplet upcycling. arXiv preprint arXiv:2409, 292.19999..1291,, 2024.
[34] Yilun Kong, Guozheng Ma, Qi Zhao, Haoyu Wang, Li Shen, Xueqian Wang, and Dacheng Tao. Mastering massive multi-task reinforcement learning via mixture-of-expent decision transformer. In ICLR Workshop on Modularity for Collaborative, Decentralized, and Continual Deep Learning, 2025.
[35] Yan John Schulman, Xi Chen, Peter I. Bartlett, llya Sutskever, and Pieter Abbeel. RI.2: Fast reinforcement learning via slow reinforcement learning. arXiv preprint 1rivv:611.002779 2016.
John
[36] Zhi Wang, Chunlin Chen, and Daoyi Dong. Lifelong incremental reinforcement learning with online Bayesian inference. IEEE Transactions on Neural Networks and Learning Systems, 330))44003-4016 2022.
[37] Zhi Wang, Chunlin Chen, and Daoyi Dong. A Dirichlet process mixture of robust task models for scalable lifelong reinforcement learning. IEEE Transactions Cybernetics, 53(12):7509-752) 2023.
[38] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-leaming for fast adapta- tion of deep networks. In Proceedings of International Conference on Machine Learning, pages 1 1126-1 1135, 2017.
[39] Eric Mitchell, Rafael Rafailov, Xue Bin Peng, Sergey Levine, and Chelsea Finn. Offline meta- reinforcement learning with advantage weighting. In Proceedings of International Conference on Machine Learning, pages 7780-7791, 2021.
[40] Luisa Zintgraf. Sebastian Schulze, Cong Lu, Leo Feng, Maximilian Igl. Kyriacos Shiarlis, Yarin Gal, Katja Hofmann, and Shimon Whiteson. VariBAD: Variational Bayes-adaptive deep RL via meta-learning. The Journal of Machine Learning Research, 22119:13198 1121))3361113666 2021.
[41] Zhi Wang, Li Zhang, Wenhao Wu, Yuanheng Zhu, Dongbin Zhao, and Chunlin Chen. Meta- DT: Offline meta-rl as conditional sequence modeling with world model disentanglement, In Advances in Neural Information Processing Systems, volume 37, pages 44845-44870, 2024.
[42] Lanqing Li, Hai Zhang, Xinyu Zhang, Shatong Zhu, Yang YU, Junqiao Zhao, and Pheng-Ame Heng. Towards an information theoretic framework context-based offline meta-reinforeemmnt learning. In Advances in Neural Information Processing Systems, volume 37, pages 75642- 75667, 2024.
[43] Shilin Zhang, Zican Hu, Wenhao Wu, Xinyi Xie, Jianxiang Tang, Chunlin Chen, Daoyi Dong, Yu Cheng, Zhenhong Sun, and Zhi Wang. Text-to-decision agent: Offline mena-reinnoonemenn learning from natural language supervision. In Advances in Neural Information Processing Systems, 2025.
[44] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, et al. Language models are few-shot learners. In Advances in Neural Information Processing Systems, volume 33, pages 1877-1901, 2020.
13

---

[45] Jake Grigsby, Linxi Fan, and Yuke Zhu. AMAGO: Scalable in-context reinforcement learning for adaptive agents. In Proceedings of International Conference on Learning Representations, 2024.
[46] Jake Grigsby, Justin Sasek, Samyak Parajuli, Daniel Adebi, Amy Zhang, and Yuke Zhu. AMAGO-2: Breaking the multi-task barrier in meta-reinforcemmnt learning with transformers. In Advances in Neural Information Processing Systems, volume 37, pages 87473-87508, 2024.
[47] Chris Lu, Yannick Schroecker, Albert Gu, Emilio Parisotto, Jakob Foerster, Satinder Singh, and Feryal Behbahani. Structured state space models for in-context reinforcement learning. In Advances in Neural Information Processing Systems, volume 36, pages 47016-47031, 2023.
[48] Robert A Jacobs, Michael Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural Computation, 3(1):79-87, 1991.
[49] Michael Jordan and Robert A Jacobs. Hierarchical mixtures of experts and the EM algorithm. Neural Computation, 6(2):181-2) 1994.
[50] Damai Dai, Li Dong, Shuming Ma, Bo Zheng, Zhifang Sui, Baobao Chang, and Furu Wei. StableMoE: Stable routing strategy for mixture of experts. In Proceedings of Annual Meeting of the Association for Computational Linguistics, pages 7085-7095, 2022.
[51] Joan Puigcerver, Carlos Riquelme, Basil Mustafa, and Neil Houlsby. From sparse to soft mixtures of experts. In Proceedings of International Conference on Learning Representations, 2024.
[52] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint 1rXivv1807.03744 2018
[53] Michael Laskin, Aravind Srinivas, and Picter Abbeel. CURL: Contrastive unsupervised repre- sentations for reinforcement learning. In Proceedings of International Conference on Machine Learning, pages 5639-56366 (6509,550.2000 2020.
[54] Zican Hu, Zongzhang Zhang, Huaxiong Li, Chunlin Chen, Hongyu Ding, and Zhi Wang. Attention -uidedd contrastive role representations for multi-agent reinforcement learning. In Proceedings of International Conference on Learning Representations, 2024.
[55] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 9729-9738. 2020.
[56] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with 1O-awareness. In Advances in Neural Information Processing Systems, volume 35, pages 16344-16359, 2022.
[57] Alexander Nikulin, Vladislav Kurenkov, Ilya Zisman, Artem Agarkov, Viacheslav Sinii, and Sergey Kolesnikov. Ainnd--ini Scalable meta-reinforcement learning environments in JAX. In Advances in Neural Information Processing Systems, volume 37, pages 43809-43835 2024.
[58] Moo Jin Kim, Karl Pertsch, Siddharth Karamcheti, Ted Xiao, Ashwin Balakrishna, et al. OpenVLA: An open-source vision- language-action model. In Proceedings of Conference on Robot Learning, 2024.
[59] Physical Intelligence, Kevin Black, Noah Brown, James Darpinian, Karan Dhabalia, et al. To.5: vision-language-acioon model with open-world generalization. arXiv preprint arXiv:2504.1600: 2025.
[60] Emanuel Todorov, Tom Erez, and Yuval Tassa. MuJoCo: A physics engine for model-based control. In Proceedings of IEEE/RSJ International Conference on Intelligent Robots and Systems, pages 5026-5033 2012.
[61] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta- World: benchmark and evaluation for multi-task and meta reinforcement learning. In Proceedings of Conference Robot Learning, pages 1094-1 100, 2020.
14

---

[62] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In Proceedings of International Conference on Machine Learning, pages 1861-1870, 2018.
[63] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:170..06347, 2017
[64] Antonin Rafin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann.
Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268):1-8, 2021.
[65] Hao Liu and Pieter Abbeel. Emergent agentic transformer from chain of hindsight experience. In Proceedings of International Conference on Machine Learning, volume 202, pages 21362- 21374, 2023.
[66] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (GELUS), arXiv preprint arXiiv: 1606. 15, 2016.
[67] Zhendong Chu, Renqin Cai, and Hongning Wang. Meiifreiiforeement learning via exploratory task clustering. In Proceedings of the AAAI Conference Artificial Intelligence, pages 11633-
1 11641, 2024.
15

---

Appendix
Algorithm Pscudocodes
17
B Regularization Loss for Balancing Expert Utilization
19
Contrastive Learning for Task-wise MoE Router
20
D The Details of Environments and Dataset Construction
21
D.1 The Details of Environments
D.2 The Details of Datasets
22
Baseline Methods
23
Network Architecture and Hyperparameters of T2MIR
23
Analysis of Hyperparameters and T2MIR Architecture
24
Hyperparameter Analysis
G.2
Architecture Analysis
24
H
Robustness to OOD Setups
26
Robustness to OOD test tasks
26
Robustness to OOD Offline Prompt Data
26
More Experiments on Ant-Dir and Meta- World
26
More Visualization Insights
27
Empirical Evidence of Gradient Conflicts)
27
Limitations
28
16

---

A Algorithm Pseudocodes
Based on the implementations in Sec. 4 this section gives the brief procedures T2MIR, Algorithm ! and Algorithm 3 show the pipline of training T2MIR-AD and T2MIR- respectively. We train all components including token-wise MoE, task-wise MoE and their routers together with the main I show the evaluation
causal transformer network end to end. Then, Algorithm 2 and Algorithm phase, where the agent can improve performance on test tasks by interacting with the environments without any parameter updates.
Algorithm 1: Model Training of T2MIR-AD
Input:
Training tasks Tuain and corresponding offline datasets Duain
Causal transformer Te: Trajectory number L; Prompt length Batch size B;
Contrastive loss weight técon
for each iteration do
Initialize query batch B = 0, key batch 5' (
for h= 1,2, do
Sample task M, Timin and obtain the corresponding dataset D, from Drain Sample L trajectories from D;
Sort L trajectories by return and concatenate to т; {so, dg To..... 87, @T, FT) Sample another trajectories from D; to obtain -. as positíve key
Add 7; and T/ and B' respectively
end
Get query batch B [ Thm and key batch B' {T)) Autoregressively predict the actions with and compute loss in Eq. 10 using B
Jh
11
Compute balance loss D'laance in Eq. 33)
13
Compute contrastive loss Lecarastive with key batch B in Eq.
Update Te by minimize loss (()) + Lisalance + Weon Lcontrstive
15
Update momentum router in task wise MoE in Eq. 8
16
end
Algorithm 2: Model Evaluation of T2MIR-AD
Input:
Test tasks Tust
Causal transformer me;
Evaluation episodes N
History buffer R: History episode length
for each task M, TeS do
Initialize history buffer R 0
for 1.... N do
Reset env. and get feedback 80 for each timestep do
Predict action a₂ = h (st, R)
Step env. and get feedback 8 and r.
Add $,, a,, rt to buffer R end
Sort R. by episode return nn > K then
11
12
Drop the first episode in
13
end
14 end
15
end
17

---

Algorithm 3: Model Training of T2MIR-DPT
Input:
Training tasks Tuaann and corresponding offline datasets Dwain Causal transformer Te: Prompt length T; Expert policy π Batch size B:
Contrastive loss weight Wcon
for each task M, E Trriin do
Obtain dataset D; from Drain
Get optimal action for states in D; using expert policy π.
Obtain query dataset D.
5 end
Obtain query datasets D for each iteration do Initialize query batch B = (}. key batch {
train
for b 1.2,.... .... B do
10
Sample task M, ~ Timin
11
Obtain the corresponding dataset D; from Pirain and query dataset DV from D Sample state-action pairs (s,a*) a²) and (s', aa') from D. Sample trajectory from D; and obtain т; {su, du To:..., ST, ar, rr)
train
12
14
Sample another trajectory from D; to obtain T- as positive key
Add ((s,a"), T.) and ((s', ",,,,a"), to B and 5' respectively
15
end 16
Get query batch B {(s, )) j== and key batch = ((s', (o5,a") "a)),'))) Autoregressively predict the optimal action with TT and compute loss in Eq.
17
1==1
18
using B
Compute balance loss Lialunce in Eq. (3)
20
Compute contrastive loss Levntrativee with key batch 13' in Eq.
Update Te by minimize loss L == ((0) + Lbalance Weon Lcontrastive
22
Update momentum router in task wise MoE in Eq.
end
Algorithm 4: Model Evaluation of T2MIR-DPT
Input:
Test tasks TTest: Causal transformer no:
Evaluation episodes
History buffer R; Episode history buffer
for each task M; E Tee do
Initialize episode history buffer 0
for n = 1,
do
Initialize history buffer R R= 0 Reset env. and get feedback for each timestep do
Predict action
n (st, T)
Step env. and get feedback 81+1 and rs
Add Pt to buffer R
10
end
11
end
end
18

---

B Regularization Loss for Balancing Expert Utilization
In this section, we provide detailed information about the regularization loss mentioned in Sec. 4.1 To prevent the gating network from converging to a state where it consistently assigns large weights to the same few experts, we adopt both importance loss and load- boad-balancing loss as proposed by Shazeer al. 225
The importance loss aims to encourage all experts to have equal significance within the model, and the importance of an expert relative to batch of training examples is defined the sum of the gate values assigned to that expert:
Imp(h)
Σ Gink ZEX
(12)
where Giok (a)(h) represents the gate value for expert h given input x. An additional loss term Lamportunce is then added to the overall loss function, which is calculated as the square of the coefficient of variation of the set of importance values multiplied by hand-tuned scaling factor Wmp:
Limportance
Wimp (Imphh)).
(13)
Despite ensuring equal importance among experts, discrepancies may still arise in the number of input tokens each expert receives due to the top-k activation mechanism. One expert might be assigned large weights for only a few tokens while another expert could receive small weights across many tokens but fail to activate. To address this issue, load-balancing loss is introduced to ensure that each expert handles approximately the same number of input tokens.
To achieve this, smooth estimator Load(h) of the number of examples assigned to each expert. This estimator allows gradients to propagate through backpropagation, by utilizing the noise term in the gating function, which we omit in Eq. 2 Given trainable weight of noise matrix W Wnoises Eq. is rewrite as
Whok (i; h) (o]tm))(HH(((b)]]]
(14)
Hiih) = topk (Gue (ilh) + StandardNormal() Softplus(h W Wnoise)i).
Let h)) denote the probability that wook (i; h) is non-zero, given a new random noise choice for element i, while keeping the already sampled noises for other elements constant. Note that, Work (1: h)
is non-zero if and only if (ilh) is greater than the htth -greates element of (-Ihh) excluding itself. Thus, we have
P(i:h) Pr (ilhok Standarddoodalllll -ootooormmmmiSSSSSllsssssh/u/// Softplus(h- Wnohe-Wnooosee Wn kth exc.lldin((HH (H (-11), kr i)),
(15)
where kth excluding(())) (-Ih), k, i) means the Rtth -dh--tettest element of H (h), excluding i-th element. Based on this, we can compute:
P(i:h) (oo((t)h) - (((lecccccdiig((((ii))..... k,
(16)
Softplus((h: Wnosse)s)
where & is the cumulative distribution function (CDF) of the standard normal distribution. Then, the estimated load for expert is given by:
Load hb)
Σ h(i;h).
(17)
Finally, the load-balancing loss is defined as the square of the coefficient of variation of the load vector, scaled by hand-tuned parameter Whodd
Lion = Whead (Lodd)(h))
(18)
Combining both losses yields the final regularization loss showed in Eq. 3 used in our model which helps maintain balanced utilization of experts during training, preventing any single expert from dominating the computation.
19

---

C Contrastive Learning for Task-wise MoE Router
In this section, we give the proof of Theorem | based on lemma as follows.
Lemma 1. Give a task from the task distribution M P(M), let f(r) as the average of hidden state after self-attention calculation of task M, Gtask
Then, we have
p(M
p>=hh)
(19)
P(M)
p(2)
Proof.
p(M[2)
p(2(MM
(MM)
p(z)
J. p(al))((hMM) dh
(20)
p(2)
p(2hh)
p(z)
The proof is completed.
Theorem 1. Let M denote set of tasks following the task distribution P(M), and |M| = N.
M MEMin is a given task. Let h= f(T), ~~ Grask (-1), and e(h, 2) = r(=la)
p(2) where t is a trajectory
from task M and h is average hidden state of all tokens after the self- eeffaatentionn calculation function f(). Let h'' denote the average hidden state generated by any task M' eM, then we have
[(2; log N + E,,,, log e(h,2)
(21)
ΣMEM e(h',
Proof. Using Lemma and Jensen's 's inequality, we have
p(z
2(h,2)
E,,z, log
E B.M.,h
log
p(=)
EMrEMette e(h',2)
p(ulh
EM'MMA p(=)
p(ul))
p(=)
1 p(2)
p(=[T')
E,
log
E
p(=[h) M'MMM
p(z)
1 p(2) (N - 1)BMEMM
p(2[h") p(2)
})]
BM.ah
log
p(2[h)
1 + p(z)
(N -1))
=EMash
log
p(z[h)
1
=EM,s,h
log
1 p(z)
(N 1)
p[2[h]
< BMM..h
log
ple) M(=lh)
N
p(=|h) p(2)
< EMM=
log E
log N
log p( p(M12)]
E,M,
log N
p(M)
[(2; M)
log N.
(22)
Thus, complete the proof.
O
20

---

D The Details of Environments and Dataset Construction
In this section, we show details of the evaluation environments over a variety of benchmarks, as well as the collection of offline datasets on these environments.
D.I The Details of Environments
We evaluate T2MIR and all the baselines on classical benchmarks including discrete environment
commonly used in ICRL [[2] the 2D navigation, the multi-task MuJoCo control [60] and the
Meta-Worke [6T. More specifically, we evaluate all tested methods on the following environments as
DarkRoom: a discrete environment where the agent navigates in 10 x 10 grid to find the goal. The observation is the 2D coordinate of the agent, and the actions are left, right, up, down, and stay. The agent is started from [o, 0] to find the goal which is uniformly sampled from the grid. The reward is sparse and only r = 1 when the agent is at the goal, and r = otherwise. It provides 100 tasks, from which we randomly sample 80 tasks as training tasks and hold out the remaining 20 for evaluation. The maximal episode step is set to 100.
Point-Robot: a problem of a point agent navigating to a given goal position in the 2D space. The observation is the 21) coordinate of the robot. The action space is [-0.1, 0.1]2 with each dimension corresponding to the moving distance in the horizontal and vertical directions. The reward function is defined the negative distance between the point agent and the goal location. Tasks differ in goal positions that are uniformly distributed in a unit square, resulting in the variation of the reward functions. We randomly sample 45 goals as training tasks, and another goals for evaluation. The start position is sampled uniformly from [-0.1, 0.1]² for each learning episode and the maximal episode step is set to 20.
Cheetah- :eee: a multi-task MuJoCo continuous control environment in which the reward function differs across tasks. It requires planar cheetah robot to run at particular velocity in the positive z-direction. The observation space is TR which comprises the position and velocity of the cheetah, the angle and angular velocity of each joint. The action space is [-1, with each dimension corresponding to the torque of each joint. The reward function is negatively correlated with the absolute value between the current velocity of the agent and the goal, plus the control cost. The goal velocities are uniformly sampled from the distribution U [0.0755 3.0]. We randomly sample 45 goals as training tasks, and another 5 goals for evaluation. The maximal episode step is set to 200.
Ant-Dir: a multi-task MuJoCo continuous control environment in which the reward function differs across tasks. It requires quadruped ant robot run in target direction. The observation space is R22T, which comprises the position and velocity of the ant robot, together with the angle and angular velocity of 8 joints. The action space is [-1, 1]s with each dimension corresponding to the torque of each joint. The reward function is positively correlated with the velocity of the ant robot the target direction, plus the control cost. The target directions are uniformly sampled from U10, 2T). We randomly sample 45 target directions as training tasks, and another 5 for evaluation. The maximal episode step is set to 200.
Walker-Param: a multi-task MuJoCo benchmark where tasks differ state transition dynamics. It need to control a two-legged walker robot to run as fast possible with varying environment dynamics. The observation space is 1R17 and the action space is [-1, 1]6. The reward function is proportional to the rumning velocity in the positive z-direction, which remains consistent for different tasks. The physical parameters of body mass, inertia, damping, and friction are randomized across tasks. We randomly sample 45 physical parameters as training tasks, and another for evaluation. The maximal episode step is set to 200.
Reach, Push: two typical environments from the robotic manipulation benchmark Meta-World. Reach and Push control robotic arm to reach goal location in 3D space and push the puck to a goal, respectively. The observation space is R39, which contains current state, previous state, and the goal. We only use the current state vector, thus the observation space is modified to 18 without goal information. The action space is 1-1, Tasks differ in goal positions which are uniformly sampled, resulting in the variation of reward functions. The initial position of object is fixed across all tasks. We randomly sample 45 goals as training tasks, and another 5 goals for evaluation. The maximal episode step is set 100.
1I.10 contains 10 different robotics manipulation tasks for training and another 5 different tasks for evaluation. Specifically, randomly sample 5 goals for each training task and goal for
21

---

Table 4: Hyperparameters of SAC used collect multi-task datasets.
Training
Warmup steps
Save frequency
Learning rate
Soft update
Discount factor
Entropy ratio
Environments
steps
Point-Robot
2000
100
20
3e-4
0.005
0.99
0.2
Cheetah- Vel
250000
2000
10000
3e-4
0.005
0.99
0.2
Cheetah-Vei
400000
2000
10000
3e-4
0.005
0.99
0.2
3 Clusterr
Walker-Param
1000000
2000
10000
3e-4
0.005
0.99
0.2
Table 5: Hyperparameters of PPO used to collect multi-task datasets.
Environments
total timesteps
n steps
learning rate
batch size
n e_epphss
Reach
400000
2048
3e-4
64
10
0.99
Push
1000000
2048
3e-4
64
10
0.99
each evaluation task, resulting in 50 goals for training and 5 goals for evaluation. We use the setting for each robotic manipulation task in Reach and Push. The maximal episode step is set to 100.
Furthermore, we give the details of the heterogeneous version Cheetah-Vel -eel we used Appendix J
-llstt a heterogeneous version of Sheetah-Vel,, where the goal velocities are different. We sample the goal velocities from three Gaussian distributions: (0.5, 0.15²). (1.5,, 0.15²), 1.15%),N(( (2,., 0.15²). From each Gaussian distribution, we sample 14 tasks for training and tasks for evaluation, resulting in 42 training tasks and test tasks.
Cheetah-Vel-t Cluster:
Note that we construct different tasks by setting different parameters for the environment, we cannot access these parameters goal velocities Cheetah- during training following common in-context RL settings.
D.2 The Details of Datasets
For discrete environment DarkRoom, we use the expert policy to collect datasets by progressively reducing the noise, as in [14). For Point-Robot and MuJoCo environments, we employ the soft actor-critic (SAC) algorithm to train policy independently for each task, and the detailed hyperparameters are shown in Table 4 For Meta-Wort environments, we use the Proximal Policy Optimization (PPO)(33) algorithm implementation provided by Stable Baselines 3[64], and the detailed hyperparameters are shown in Table 5 During training, we periodically save the policy checkpoints and use them to generate various qualities of offline datasets as
Mixed: using all policy checkpoints to generate datasets. We use each checkpoint to generate same number of episodes, e.g., using one checkpoint to generate one episode.
Medium- Expert: using the policy checkpoints whose performance is below 80% the final achieved level to generate datasets. The max performance in Medium-Expert datasets is about 80% of that in Mixed datasets.
Medium: using the policy checkpoints whose performance is below 50% of the final achieved level to generate datasets. The max performance in Medium datasets is about 50% of that in Mixed datasets.
For query state-action pairs used in DPT and T2MIR-DPT, we use the expert policies that achieve 100%, 80% and 50% performance to provide actions for states in Mixed, Medium- Expertt and Medium datasets, respectively. As we find the actions provided by expert policies contain too many values outside the boundary of the action space in Meta- World environments, we just use the state-action pairs in offline datasets query datasets for Reach and Push.
22

---

E
Baseline Methods
This section gives the details of the five representative baselines, including four ICRL approaches and one context-based offline meta-RL. method. These baselines are thoughtfully selected as they are representative in distinctive categories. Furthermore, since our proposed T2MIR method belongs to the ICRL category, we incorporate more methods from this class as baselines for comprehensive comparison. The detailed descriptions of these baselines are as follows:
AD M. is the first method to achieve ICRL through sequential modeling of offline historical data using an imitation loss. By employing causal transformer with sufficiently large across-episodic contexts to imitate gradient- baseed RL algorithms, AD learns improved policies for new tasks entirely in context without requiring external parameters updating.
DPT [1), is an ICRL method that models contextual trajectories via supervised learning to predict optimal actions. After pretraining on offline datasets, DPT provides contextual learning capabilities for new tasks, enabling online exploration and offline gecision-making given query states and contextual prompts.
IDT 13,. is an ICRL method that simulates high-level trial-and-error processes. It introduces an innovative architecture comprising three modules: Making Decisions, Decisions to Go, and Reviewing Decisions. These modules generate high-level decisions in an autoregressive manner to guide low-level action selection. Built upon transformer models, IDT can address complex decision-making in long-horizon tasks while reducing computational costs.
combines model-based reinforcement learning with previous ICRL algorithms like AD and IDT. By jointly learning environment dynamics and policy improvements in context, DICP employs a transformer architecture to perform model-basee planning and decision-making without updating model parameters Furthermore, DICP generates actions by simulating potential future states and predicting long-term returns.
DICP 15
UNICORN 42), is context-basee offline meta-RL RL algorithm grounded in information theory. proposes a unified cnormatioo-theorette framework, which focuses on optimizing different approximation bounds of the mutual information objective between task variables and their latent representations. Leveraging the information bottleneck principle, it derives a novel gen- eral and unified task representation learning objective, enabling the acquisition of robust task representations.
To ensure a fair comparison, all baselines are adjusted to an aligned setting, where test datasets are not available for policy evaluation. We also standardize the size and quality of the offline datasets for all baselines.
In our experimental setup, AD faces difficulties in distilling effective policy improvement operators due to the scarcity of available offline historical data. Thus, we use an enhanced version of AD that incorporates a trajectory-rankin mechanism based on cumulative rewards, inspired by AT 65
F
Network Architecture and Hyperparameters of T2MIR
This section gives details of the architecture and hyperparameters of T2MIR implementations as follows.
Router Network. We implement the router network using multi-layer perceptron (MLP) archi- tecture without bias. Specifically, the router contains two linear layers with Tanh as the activation function between them. The first linear maps hidden state h to a m__eeee _experts-dim vector along with the Tanh activation, and the second linear maps the n www.eeeettsteee -eip vector to another eerperts-dim vector as the output of router.
Architecture of T2MIR. We implement the T2MIR-AD based on causal transformer akin to AM- AGO 455 with Flash Attention [56]. For T2MIR-DPT, we build our T2MIR framework upon the transformer architecture in DPT [12), flatten the trajectories into Mate-action-rewart sequence as input. Specifically, we employ separate embeddings for states, action, and rewards, adding learnable positional embedding based on timesteps. Each block employs a multi-head self-attention module followed by a feedforward network or a MoE. layer with GELU activation (66). For action prediction, employ linear to map the output of transformer blocks to an action vector with Tanh activation.
23

---

Table 6: Hyperparameters in training process of T2MIR-AD using Mixed datasets.
Hyperparameters
DarkRoom
Point-Robot
Cheetah-Va Vel
Walker-Param
Reach
Push
training steps
300000
100000
100000
100000
100000
100000
learning rate
3c-4
3e-4
3e-4
3e-4
5c-5
3e-4
Prompt length
400
800
800
400
400
token experts K
6
6
6
6
6
task experts K2
12
8
8
8
8
8
InfoNCE weight
0.01
0.01
0.01
0.01
0.01
0.01
MoE layer position
top
top
top
top
top
top
Momentum ratio B
0.995
0.995
0.995
0.995
0.995
0.995
Table Hyperparameters in training process of T2MIR-DPP using Mixed datasets.
Hyperparameters
DarkRoom
Point-Robot
Cheetah- Vel
Walker-Param
Reach
Push
training steps
300000
100000
100000
100000
100000
100000
learning rate
3e-4
3e-4
3e-4
3e-4
5e-5
3e-4
Prompt length
100
20
200
200
100
100
token experts K 6
6
6
6
task experts K₂
8
8
8
8
8
8
InfoNCE weight
0.001
0.001
0.01
0.01
0.01
0.001
MoE layer position
top
top
top
top
top
top
Momentum ratio 3 0.995
0.995
0.995
0.995
0.995
0.995
Table 6 and Table 7 show the detailed hyperparameters used for T2MIR-AD and T2MIR-DPT using Mixed datasets, respectively.
Compute. We train our models on one Nvidia RTX4080 GPU with the Intel Core i9- 10900X CPU and 256G RAM. The training process takes about 0.5-3 hours, depending the complexity of the environments. Compared with T2MIR-free backbones, T2MIR takes a lightly higher computational cost during the training process, but it requires less computational resources during task inference.
G A Analysis of Hyperparameters and T2MIR Architecture
G.1 Hyperparameter Analysis
Analysis of Expert Configurations on T2MIR-DPT. We also investigate different expert configura- tions on T2MIR- -DP, shown Figure Table and Table give the numerical results of T2MIR with different expert configurations in token- and task-wise MoE, respectively. For loken-wise MoE, as same T2MIR-AD, the results show that moderate configuration (2/6) outperforms both smaller (1/3) and larger (4/12) settings. For task-win wise MoE, the performance improves with the increase number of total experts, and tends be stable.
Weight of InfoNCE Loss. We investigate the influence of the loss weight of InfoNCE loss in task-wise MoE, the results are shown in Figure and Table 10 The performance of T2MIR is not sensitive to small weights. But when the InfoNCE loss weight is set to 0.1, it gains an obvious decrease. We suppose this is because the final loss of Loontrastive in Eq. (7) has same magnitude with imitation loss C()) in Eq. (0) and Eq. [1] and it disturbs the overall learning process.
G.2 Architecture Analysis
Position of MoE Layer Further more, we explore how the placement of the MoE layer affects performance. We substitute the feedforward network in transformer block with MoE layer and

---

Table 8: Numerical results of T2MIR with varying expert numbers in token-wise MoE. Best in bold.
T2MIR-AD
1/3
2/6
4/12
T2MIR-DPT
1/3
2/6
4/12
Point-Robot
6.6++7
-5.2++0.3
-6 -6,1+0.5
Point-Robot
-5.6+022
-5.0
-5.3+0.3
Cheetah- Vel
86.1+22
-68.9+2.1
-80.1+0.
Cheetah Vel
-53.3+13
-43.2
-50.7+199
Table 9: Numerical results of T2MIR with varying expert numbers in task-wise MoE. Best in bold.
T2MIR-AD
2/4
2/6
2/8
T2MIR DPT
2/4
2/6
2/8
Point-Robot
5.3110
-3.2103
-5.2.103
Point-Rohot
-..2104
-5.2.102
-5.0
Cheetah- Vel
-91.7+279
-77.3+1.7
-68.9+2.1
Cheetah-Ve Vel
-51.5=74
-45 -45.9221
-43.2+008
0.01
0.001
Point-Rabct
Cheetah-Vil
Foint-Robct
Cheetah-Vil
-40
-60
Reture
Reture
20
-120
10
In-Conteot Epiendes
In-Conteat Spisades
In-Conteat Epicodas
In-Contest Episodies
()) T2MIIRR-
(b) T2MIR-DDP
Figure 9: Test return curves of T2MIR with different contrastive loss weights.
Table 10: Numerical results of T2MIR with different loss weights of InfoNCE loss, i.e., numerical results of best performance from Figure Best result bold.
T2MIR-AD
0.1
0.01
0.001
T2MIR-DPP
0.1
0.01
0.001
Point-Robot
1..9116
-0.2
55.2103
Point-Robot
-..1104
-5.410.1
-5.0
Cheetah Vel
-150
-68.9+2.1
-74.1+10
Cheetah- Vel
-9..441.
-43.2+08
-54.9+1.3
study three settings of different positions: bottom, middle and top, where the MoE layer is used in the first, -tth and last transformer block, respectively. Figure 10 and Table 11 present the results, where substituting the feedforward network in the last transformer block with MoE layer (top) gets more stable performance.
Importance of Balance and InfoNCE Loss. We conduct further ablation studies to investigate the respective impacts of additional losses used to balance the experts. We compare T2MIR-AD to three ablations: 1) w/o balance_loss omitting balance loss (Eq. (3) in token-wise MoE, while retaining InfoNCE loss in task wise MoE; 2) w/o anfonce_losee omitting InfoNCE loss (Eq. in task-wise MoE, while retaining balance loss in token-wise MoE; and 3) w/o aux_ aux_loss, omitting both balance loss and InfoNCE loss. For the ablation studies, we just omit the additional losses while keeping the token- and task MoE architectures. The numerical results in Table 12 show that both balance loss and InfoNCE loss play significant roles in T2MIR's superior performance. Ablating the balance loss will cause a bit more performance degradation than ablating the InfoNCE loss.
Architecture of Token-wise MoE. To quantify the challenge incurred by the semantic discrepancy among states, actions, and rewards, we completely isolate the multi-modality issue and manually design MoE structure where state-action-reward tokens are routed to three experts by heuristic gating scheme, i.e., one expert for one modality. The results in Table 13 show that the heuristic gating mechanism contributes to performance improvement, while T2MIR achieves more pronounced enhancement by automatically routing state-action-rewann transitioning from heuristic routing to automatic routing may stem from more flexible approach to processing the integrated information generated after the attention mechanism. The results again demonstrate the significant challenge induced by the semantic discrepancy across multi-modal tokens in the date-actton-ee actio--rwaard sequence, and the significant superiority of token-wise MoE architecture.
tokens. The enhancement achieved by
25

---

bottom
modile
top
Cheetah-Val
Choctah
-40
-60
Roture
-80
Rotur
-00
-120
-140
-10
-10
-120
T2MBR-AD
T2MIR-DFT
Figure 10: Test return curves of T2MIR with different positions of MoE layer.
Table Numerical results T2MIR with different positions of MoE layer, i.e, numerical results
of best performance from Figure 10
Best result in bold.
T2MIR-AD
bottom
middle
top
T2MIR-DPT
bottom
middle
top
Point-Robot -6.0
-5.8403
-5.2+03
Point-Robar
-5.5203
-..720.
-5.0=0.5 -43.2
Cheetah- Vel
1.6
-70.245.4
-
Cheetah- Vel
-52.9240
-43.842.
Table 12: Numerical results of ablation study on the importance of balance and InfoNCE loss T2MIR-AD using Mixed datasets. Best result in bold.
Environments
T2MIR-AD
w/o infonce loss
w/o balance loss
w/o aux loss
Cheetah-VV Vel
-68.922.1
-76.7=35 417.0126
-9..8+1. 416.0
-112.4130.4 393.21167
Walker-Param
435.7
Table 13: Numerical results of different architectures of token-wise MoE T2MIR-AD using Mixed datasets. Best result in bold.
Environments
T2MTR-AD
heuristic routing
AD
Cheetah-Ve Vel Walker-Param
-68.9
-91.14304 407.1:10. 10.5
-119.2=90.4 395.2:17.1
435.7
H Robustness to OOD Setups
H.1 Robustness to OOD test tasks
We conduct experiments to evaluate T2MIR's capability handle OOD test tasks Cheetah- Vel against baselines. The goal velocities of training tasks are uniformly sampled from U(0.075, 3), but during testing sample OOD tasks from 733.1,3.9 The results in Table 14 demonstrate T2MIR's consistent superiority even when dealing with OOD test tasks.
H.2 Robustness to OOD Offline Prompt Data
In our main experiments in Sec. 5.1 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 infer task beliefs during testing. In this section, we have conducted new experiments to evaluate T2MIR-DPT's robustness to OOD offline data as the test prompt. Specifically, we employ a random behavior policy to interact with the environment and collect one trajectory the offline prompt. The results in Table demonstrate T2MIR's s consistent superiority even when prompted with OOD offline data.
I More Experiments on Ant-Dir and Meta-World
In addition, we have conducted new evaluations on the challenging Ant-Dir in MuJoCo and MI.10 in Meta- World to demonstrate T2MIR's ability to these harder suites. Details about the two environments
26

---

Table 14: Numerical results of T2MIR against baselines on OOD test tasks using Mixed datasets. Best result in bold and second best underline.
Environment T2MIR-AD T2MIR-DPT
DPT
DICP
IDT
UNICORN
Cheetah-Vel -117.3+44
-71.0+07
-78.5+26
-109.1+21
-1.
88.6+2.1
Table 15: Numerical results of T2MIR-DPT against DPT with OOD offline prompt data using Mixed datasets.
Environments
T2MIR-DPT
DPT
Point-Robot
-5.0402
-5.940.1
Cheetah- Vel
-35.0229
63.2244
Walker-Param
424.1+17.3
419.0#11.5
Table 16: Numerical results of T2MIR on Ant-Dir MuJoCo and ML10 deta-World. Best result in bold and second best underline.
Environments
T2MIR-AD 12MIR-DPT
AD
DPT
DICP
IDT
UNICORN
Ant-Dir
681.522 144.4
708.61245
432. 432.21425
634.1 46.7
643.7=34.5
450.9=102 193.6-10
MI.10
318.1
992.99219
277.01
230.6225.1
215.92112
23,
can be found in Appendix [ The numerical results in Table 16 show the promising scalability and the consistent superiority of T2MIR over various baselines.
J More Visualization Insights
t-SNE Visualization on Push. We also visualize the multi- modal property and task clustering in T2MIR-AD on Push, as shown in Figure 12 The visualization results consis- tently exhibit the ability of token-wise MoE to process token from different modalities with different experts, and task-wise MoE to distribute trajectories from different tasks to different experts.
Cheetah-Vel-3 Cluster
-50
Returns
-100
T2MBBAC (Dun)
Heterogeneous Setting. For further study of the task clustering ability of T2MIR, we construct a heteroge- neous version Cheetah- -el inspired by MILE, [67]. We split the velocities of tasks into three intervals, construct- ing Cheetah-Vel-33 wew---ell-CCusteerrrrree Cluster environment. Figure 11 presents the test return curves of T2MIR and baselines using Mixed datasets. We demonstrate the probability of task assignments to some experts, as shown in Figure [3 the velocities from three different distributions are divided by dashed lines. The expert- dominates at lower speeds, the expert-2 dominates at higher speeds, and it exhibits a mixture of different experts for medium speeds, with expert-1 I being dominant. Moreover, there are switched among different experts at the boundaries between different velocity distributions. The results indicate that distinct experts dominate different tasks sampled from the three distributions, which further validates our motivation to employ token-wiig wise MoE.
T2MIR DPT 10urs)
DCP
AD
Dr
-20
In-Context Episodes
Figure 11:
Test return curves of
T2MIR against baselines on Cheetah- J0l-3_Clustering using Mixed datasets.
K Empirical Evidence of Gradient Conflicts
In addition to the results in Figure g we also provide quantification the gradient conflict issue. The results in Table 17 present the proportion of task pairs with negative correlations (cosine similarity < <0.055 among all pairs of tasks, which demonstrate T2MIR's significant superiority in mitigating gradient conflicts.
27

---

rerigh)
Figure 12: t-SNE visualization of expert assignments on Push. Left: token-wiss MoE enables different experts to process mtoee-modality tokens. Right: task-wise MoE effectively manages task distribution, where trajectories from same task are prone be distributed to same experts.
expert-1
expert-2
expert-3
1.00
0.75
Probability
0.50
0.25
0.00
0.5
1.5
2.5
Velocity
Figure 13: Probability of task assignments some experts. The goal velocities are sampled from three distributions and divided by dashed lines.
Table 17: Quantification of the gradient conflict.
Methods
Point Robot
Cheetah- Vel
Walker-Param
Reach
Push
AD
31.1% 7.4%
33.9% 1.2%
23.3% 1.1%
20.2% 15.6%
12.8%
T2MIR-AD
1.8%
Limitations
We discuss the limitations and computational efficiency of our method in this section. Constrained by limited computing resources, our method trained on lightweight datasets, e.g., the Mixed datasets of Push contain 500 episodes per task and 45 training tasks, which have totally 2.25 million transitions. Although the size of datasets are small, we find it's enough to gain high performance. But the size of datasets limits the capacity of MoE in extending the size of model parameters. Future work may evaluate their study on more complex environments and larger datasets such as XLand- MiniGrid (57 10] The efficiency of contrastive loss in task-wise MoE when facing massive number of tasks is not thoroughly explored in this work. It interesting to explore whether the contrastive loss effective for managing task assignments in scenarios involving a large number tasks. Integrating the T2MIR framework incurs slightly higher computational cost during the training process, but requires less computational resources during task inference.
28

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Please refer the Abstract and Section |
Introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Please refer to Section 6 Conclusions, Limitations, and Future Work, and Appendix L
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]

---

Justification:
Please refer to Appendix C for the proof of Theorem !
Guidelines:
The answer NA means that the paper does not include theoretical results,
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced
statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Please refer to Appendix A for algorithm pseudocodes. Appendix D for the details of environments and dataset construction, and Appendix for the implementa- tion details and hyperparameters of our method. We also provide our source code in the supplementary materials.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
30

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Please refer Appendix Appendix D and Appendix for details of our method and dataset construction. We also provide our source code in the supplementary materials.
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.com cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
y /nips. /nips.cc///pp http:/ccccccccguueeeeeeeeeeeeeee gwalcc//uuess//ee httpbliiigguuieesoodduuubiiiiiie
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question:
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: Please refer to Section 5 Experiments, Appendix D Appendix G provides analysis of hyperparameters.
Appendix F
And
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Please refer Section Experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
31

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Please refer to Appendix
F
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https::/neuripp: httpp//rwwrrprrcpporree///aaaaaa cccppublic/Ehhiiciiieeelleeeeeee
Answer: [Yes]
Justification: Please refer to Appendix A Appendix D and Appendix F for details of implementation and dataset construction. We also provide our source code in the supplementary materials.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., s if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Please refer Section 6 Conclusions, Limitations, and Future Work.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
32

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification: We construct the datasets used in this work ourselves, and not use existing assets, please refer to Appendix D for details.
Guidelines:
The answer NA means that the paper does
existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
33

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssi
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Our work does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/codeemoodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our work does not involve crowdsourcing and research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: Our work does involve crowdsoureing and research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
34

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in our work does not involve L.MSs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM for what should or should not be described.
35


Reviews:
[{"review_id": "2n9vBbIs2j", "reviewer": "Reviewer_wTXA", "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.", "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", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. 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.\n2. The insight of pooling trajectory to ensure interference of gradient is reduced is interesting.\n3. Results demonstrate faster convergence and better rewards compared to baseline methods.\n\nWeakness:\n1. 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.\n2. 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.", "was_revised": true, "final_justification": "Thanks to the authors for addressing the review comments and providing more empirical data! I have revised my ratings for the paper.", "rebuttal": "***\n\n**`Q1. The necessity of leveraging explicit architectural bias for in-context RL.`**\n\nA1. 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**.\n\nThe 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**.\n\nFurther, 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**.\n\n[1] Mixtures-of-experts unlock parameter scaling for deep RL, ICML 2024.\n\n[2] Don't flatten, tokenize! Unlocking the key to SoftMoE's efficacy in deep RL, ICLR 2025.\n\n***\n\n**`Q2. The scalability of our method, with experiments on small-sized data with relatively small models.`**\n\nA2. 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**.\n\n| T2MIR-AD | 1.8M transitions | 3.6M transitions | 5.4M transitions |\n| --- | --- | --- | --- |\n| **4.8M parameters** | −68.9 $\\pm$ 2.1 | -65.3 $\\pm$ 8.6 | -54.8 $\\pm$ 0.9 |\n| **5.8M parameters** | -69.0 $\\pm$ 1.6 | -58.6 $\\pm$ 0.8 | -51.1 $\\pm$ 0.6 |\n| **7.2M parameters** | -61.2 $\\pm$ 1.0 | -57.3 $\\pm$ 0.7 | -49.6 $\\pm$ 0.7 |\n\nFurther, 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**.\n\n| **Method** | **UNICORN** | **IDT** | **DICP** | **DPT** | **AD** | **T2MIR-DPT** | **T2MIR-AD** |\n| --- | --- | --- | --- | --- | --- | --- | --- |\n| Ant-Dir | 450.9 $\\pm$ 10.2 | 643.7 $\\pm$ 34.5 | 690.3 $\\pm$ 20.4 | 634.1 $\\pm$ 46.7 | 432.2 $\\pm$ 42.5 | **708.6** $\\pm$ 24.5 | 681.5 $\\pm$ 30.4 |\n| ML10 | 193.6 $\\pm$ 10.7 | 223.3 $\\pm$ 13.3 | 215.9 $\\pm$ 11.2 | 230.6 $\\pm$ 25.1 | 277.0 $\\pm$ 39.0 | 291.3 $\\pm$ 21.9 | **318.1** $\\pm$ 11.3 |\n\n \n\nFollowing 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.\n\n[3] Network Sparsity Unlocks the Scaling Potential of Deep RL, ICML 2025.\n\n[4] In-context RL with algorithm distillation, ICLR 2023.\n\n[5] Supervised pretraining can learn in-context RL, NeurIPS 2023.\n\n[6] Distilling RL algorithms for in-context model-based planning, ICLR 2025.\n\n***\n\n**`Q3. Complexity of adding two MoE blocks per layer for both training and inference.`**\n\nA3. 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.**\n\n| Training/Inference Time | T2MIR-AD (MoE) | AD (FFN) | T2MIR-DPT (MoE) | DPT (FFN) |\n| --- | --- | --- | --- | --- |\n| Cheetah-Vel | 4h18min / 78.0s | 3h57min / 67.9s | 2h46min / 59.6s | 2h27min / 46.8s |\n| Walker-Param | 4h22min / 79.8s | 3h53min / 69.7s | 2h47min / 61.8s | 2h26min / 48.5s |\n\n***\n\n**`Q4. Hyperparameters and additional loss for balancing the experts.`**\n\nA4. 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**.\n\nFollowing 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:\n\n- **w/o balance_loss**: omitting balance loss (Eq.3) in token-wise MoE, while retaining contrastive loss in task-wise MoE.\n- **w/o contrastive_loss**: omitting contrastive loss (Eq.7) in task-wise MoE, while retaining balance loss in token-wise MoE.\n- **w/o aux_loss**: omitting both balance loss and contrastive loss.\n\nFor 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.\n\n| **Ablation** | **w/o aux_loss** | **w/o balance_loss** | **w/o contrastive_loss** | **T2MIR** |\n| --- | --- | --- | --- | --- |\n| Cheetah-Vel | -112.4 $\\pm$ 30.4 | -92.8 $\\pm$ 1.6 | -76.7 $\\pm$ 3.5 | **-68.9** $\\pm$ 2.1 |\n| Walker-Param | 393.2 $\\pm$ 6.7 | 416.0 $\\pm$ 5.3 | 417.0 $\\pm$ 2.6 | **435.7** $\\pm$ 7.2 |\n\nTo 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**.\n\n| **Ablations** | **w/o balance_loss** | **w/o contrastive_loss** |\n| --- | --- | --- |\n| Rank 1 | (1,2): 83% | (1,2): 100% |\n| Rank 2 | (1,3): 5.6% | 0% |\n| Rank 3 | (2,4): 5.1% | 0% |\n\n***\n\n**`Q5.  Maintaining fairness regarding the parameter count in ablation studies.`**\n\nA5. 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.\n\n***\n\n**`Q6. Formatting Concerns about Fig. 2.`**\n\nA6. Thanks a lot for your detailed suggestion, and we will improve the diagram and the text underneath.\n\n***\n\n### **`Summary Response`**\n\nThank 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”).\n\nAlso, 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.\n\nWe 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**."}, {"review_id": "7oLVur6AIX", "reviewer": "Reviewer_yZc9", "summary": "This 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.", "questions": "1. 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?", "limitations": "Yes. It is given in the appendix.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "The 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?", "was_revised": true, "final_justification": "The proposed Mixture of Experts-based architecture seems promising, and the authors have done a good set of experiments to show the value of the method. While the techniques used in this work are not necessarily novel, putting them together in use for in-context RL is interesting. The rebuttal was responsive to all my major points, and I think this work is of value to the community. Therefore, I lean towards an accepting score, considering the novelty of the idea, the techniques used, the rigour of experiments conducted, and the final evaluation results.", "rebuttal": "***\n\n**`Q1. Clear formalism and specification of task variation, and the use of task variation parameterization in the context.`**\n\nA1. 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.\n\n***\n\n**`Q2. The necessity of learning the gating mechanism in MoEs, and its superiority over heuristic gating approaches.`**\n\nA2. 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.\n\n| **Ablations** | **AD** | **T2MIR-AD (heuristic gate in token-wise MoE)** | **T2MIR-AD** |\n| --- | --- | --- | --- |\n| Cheetah-vel | −119.2 $\\pm$ 30.4 | -91.1 $\\pm$ 30.4 | **−68.9** $\\pm$ 2.1 |\n| Walker-Param | 395.2 $\\pm$ 7.1 | 407.1 $\\pm$ 10.5 | **435.7** $\\pm$ 7.2 |\n\nFurther, 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.\n\n***\n\n**`Q3. Ablations to show the impact of balancing loss and contrastive loss.`**\n\nA3. 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:\n\n- **w/o balance_loss**: omitting balance loss (Eq.3) in token-wise MoE, while retaining contrastive loss in task-wise MoE.\n- **w/o contrastive_loss**: omitting contrastive loss (Eq.7) in task-wise MoE, while retaining balance loss in token-wise MoE.\n- **w/o aux_loss**: omitting both balance loss and contrastive loss.\n\nFor 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.\n\n| **Ablation** | **w/o aux_loss** | **w/o balance_loss** | **w/o contrastive_loss** | **T2MIR** |\n| --- | --- | --- | --- | --- |\n| Cheetah-Vel | -112.4 $\\pm$ 30.4 | -92.8 $\\pm$ 1.6 | -76.7 $\\pm$ 3.5 | **-68.9** $\\pm$ 2.1 |\n| Walker-Param | 393.2 $\\pm$ 6.7 | 416.0 $\\pm$ 5.3 | 417.0 $\\pm$ 2.6 | **435.7** $\\pm$ 7.2 |\n\nTo 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**.\n\n| **Ablations** | **w/o balance_loss** | **w/o contrastive_loss** |\n| --- | --- | --- |\n| Rank 1 | (1,2): 83% | (1,2): 100% |\n| Rank 2 | (1,3): 5.6% | 0% |\n| Rank 3 | (2,4): 5.1% | 0% |\n\n***\n\n**`Q4. Training the contrastive loss when the number of tasks is large.`**\n\nA4. 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.\n\nTo 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.\n\n| **# training tasks** | **10-train/5-eval** | **20-train/5-eval** | **45-train/5-eval** |\n| --- | --- | --- | --- |\n| T2MIR (w/o Token-MoE) | -158.2 $\\pm$ 25.5 | -120.3 $\\pm$ 26.1 | −101.1 $\\pm$ 29.7 |\n\n***\n\n**`Q5. When the token-wise MoE count is changed, how is the task-wise MoE count handled?`**\n\nA5. 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.\n\n**`Q6. How is the decision to activate one-third of experts in token-wise MoE and always 2 experts in task-wise MoE reached?`**\n\nA6. 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.\n\n[1] GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding, ICRL 2021.\n\n[2] BASE Layers: Simplifying Training of Large, Sparse Models, ICML 2021.\n\n[3] MoE++: Accelerating Mixture-of-Experts Methods with Zero-Computation Experts, ICLR 2025 Oral.\n\n***\n\n### **`Summary Response`**\n\nThank 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”).\n\nAlso, 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. \n\nWe 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**."}, {"review_id": "EFC7XOAuyR", "reviewer": "Reviewer_sbb8", "summary": "This 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.", "questions": "1. 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.", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\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).", "was_revised": true, "final_justification": "The extensive experiments have addressed my concerns well and the OOD results significantly extend the scope of the work, showing the robustness to OOD data in offline RL. I have updated my score", "rebuttal": "***\n\n**`Q1. Evaluation on other benchmarks like Ant-Dir and Humanoid, or text-based benchmarks like webshop.`**\n\nA1. 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**. \n\n| **Method** | **UNICORN** | **IDT** | **DICP** | **DPT** | **AD** | **T2MIR-DPT** | **T2MIR-AD** |\n| --- | --- | --- | --- | --- | --- | --- | --- |\n| Ant-Dir | 450.9 $\\pm$ 10.2 | 643.7 $\\pm$ 34.5 | 690.3 $\\pm$ 20.4 | 634.1 $\\pm$ 46.7 | 432.2 $\\pm$ 42.5 | **708.6** $\\pm$ 24.5 | 681.5 $\\pm$ 30.4 |\n| ML10 | 193.6 $\\pm$ 10.7 | 223.3 $\\pm$ 13.3 | 215.9 $\\pm$ 11.2 | 230.6 $\\pm$ 25.1 | 277.0 $\\pm$ 39.0 | 291.3 $\\pm$ 21.9 | **318.1** $\\pm$ 11.3 |\n\n[1] In-context RL with algorithm distillation, ICLR 2023.\n\n[2] Supervised pretraining can learn in-context RL, NeurIPS 2023.\n\n[3] Distilling RL algorithms for in-context model-based planning, ICLR 2025.\n\n***\n\n**`Q2. Analysis of how the behavior policy and task-related information contribute to the task representation, respectively.`**\n\nA2. 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**.\n\nFurther, **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.\n\n***\n\n**`Q3. Clarify the different modalities within the state-action-reward sequence in token-wise MoE.`**\n\nA3. 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.\n\n***\n\n**`Q4. Demonstrate T2MIR’s robustness by evaluating on OOD offline data.`**\n\nA4. 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**.\n\n| **Methods** | Point-Robot | Cheetah-Vel | Walker-Param |\n| --- | --- | --- | --- |\n| **DPT** | -5.9 $\\pm$ 0.1 | -63.2 $\\pm$ 4.4 | 419.0 $\\pm$ 11.5 |\n| **T2MIR-DPT** | **-5.0** $\\pm$ 0.2 | **-35.0** $\\pm$ 2.9 | **424.1** $\\pm$ 17.3 |\n\nFurthermore, **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**.\n\n| **OOD Testing** | **UNICORN** | **IDT** | **DICP** | **DPT** | **AD** | **T2MIR-DPT** | **T2MIR-AD** |\n| --- | --- | --- | --- | --- | --- | --- | --- |\n| Cheetah-Vel  | -88.6 $\\pm$ 2.1 | -114.3 $\\pm$ 2.3 | -109.1 $\\pm$ 2.1 | -78.5 $\\pm$ 2.6 | -158.8 $\\pm$ 47.1 | **-71.0** $\\pm$ 0.7 | -117.3 $\\pm$ 4.4 |\n\n***\n\n### **`Summary Response`**\n\nThank 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\").\n\nAlso, 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.\n\nWe 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**."}, {"review_id": "HqlqGSDrYJ", "reviewer": "Reviewer_ptNc", "summary": "The 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.", "questions": "1. 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?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "The authors have adequately addressed my concerns regarding empirical evidence and benchmarks. I have raised my score accordingly.", "rebuttal": "***\n\n**`Q1. Empirical evidence of gradient conflicts in ICRL pre-training.`**\n\nA1. 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.**\n\nFurther, 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.**\n\n| **Gradient Conflict** | Point-Robot | Cheetah-Vel | Walker-Param | Reach | Push |\n| --- | --- | --- | --- | --- | --- |\n| **AD** | 31.1% | 33.9% | 23.3% | 20.2% | 12.8% |\n| **T2MIR-AD** | **7.4%** | **1.2%** | **1.1%** | **15.6%** | **1.8%** |\n\nThank 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**.\n\n***\n\n**`Q2. Isolate and quantify the multi-modality issue of state-action-reward tokens.`**\n\nA2. 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**.\n\nFurther, 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.\n\n| **Ablation** | **AD** | **T2MIR-AD (heuristic gate in token-wise MoE)** | **T2MIR-AD** |\n| --- | --- | --- | --- |\n| Cheetah-vel | −119.2 $\\pm$ 30.4 | -91.1 $\\pm$ 30.4 | **−68.9** $\\pm$ 2.1 |\n| Walker-Param | 395.2 $\\pm$ 7.1 | 407.1 $\\pm$ 10.5 | **435.7** $\\pm$ 7.2 |\n\nThank 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**.\n\n***\n\n**`Q3. Evaluate T2MIR on more challenging benchmarks such as the full ML1 or ML10.`**\n\nA3. 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**. \n\n| **Method** | **UNICORN** | **IDT** | **DICP** | **DPT** | **AD** | **T2MIR-DPT** | **T2MIR-AD** |\n| --- | --- | --- | --- | --- | --- | --- | --- |\n| Ant-Dir | 450.9 $\\pm$ 10.2 | 643.7 $\\pm$ 34.5 | 690.3 $\\pm$ 20.4 | 634.1 $\\pm$ 46.7 | 432.2 $\\pm$ 42.5 | **708.6** $\\pm$ 24.5 | 681.5 $\\pm$ 30.4 |\n| ML10 | 193.6 $\\pm$ 10.7 | 223.3 $\\pm$ 13.3 | 215.9 $\\pm$ 11.2 | 230.6 $\\pm$ 25.1 | 277.0 $\\pm$ 39.0 | 291.3 $\\pm$ 21.9 | **318.1** $\\pm$ 11.3 |\n\n \n\n[1] In-context RL with algorithm distillation, ICLR 2023.\n\n[2] Supervised pretraining can learn in-context RL, NeurIPS 2023.\n\n[3] Distilling RL algorithms for in-context model-based planning, ICLR 2025.\n\n***\n### **`Summary Response`**\n\nThank 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…\").\n\nAlso, 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.\n\nWe 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**."}]

Author rebuttal:

Meta-review: This paper introduces a framework for in-context reinforcement learning that integrates Mixture-of-Experts (MoE) architecture into a transformer backbone. The core contribution is using a token-wise MoE to handle the distinct semantics of state, action, and reward data, and a parallel task-wise MoE to route diverse tasks to specialized experts. Extensive experiments demonstrate that the proposed architecture outperforms strong baselines. 

Reviewers expressed concerns regarding the method's novelty and scalability, the added complexity of the architecture, and the empirical justification for its design. In particular about the initial evaluation, asking for tests on more challenging environments (wTXA, sbb8, ptNc), analysis of robustness to out-of-distribution data (sbb8), and stronger empirical evidence for the claimed benefits of the dual MoE structure (yZc9, ptNc). During the rebuttal many of these potential weaknesses were addressed including evidence showing the method scales to more difficult benchmarks (Ant-Dir, ML10), limited computational overhead, and justification for architectural choices.
Decision: Accept (poster)