Decision Mixer: Integrating Long-term and Local Dependencies via Dynamic Token Selection for Decision-Making
Authors: Hongling Zheng
Li Shen
Yong Luo
Deheng Ye
Bo Du
Jialie Shen
Dacheng Tao
Venue: ICML 2025 poster (2025)
Keywords: Offline RL
Conditional Sequence Modeling
Decision Transformer
Primary area: reinforcement_learning->batchoffline
TL;DR: In this paper, we propose Decision Mixer (DM), which addresses the conflict between features of different scales in the modeling process from the perspective of dynamic integration.

Abstract: The Conditional Sequence Modeling (CSM) paradigm, benefiting from the transformer's powerful distribution modeling capabilities, has demonstrated considerable promise in offline Reinforcement Learning (RL) tasks. Depending on the task's nature, it is crucial to carefully balance the interplay between inherent local features and long-term dependencies in Markov decision trajectories to mitigate potential performance degradation and unnecessary computational overhead. In this paper, we propose Decision Mixer (DM), which addresses the conflict between features of different scales in the modeling process from the perspective of dynamic integration. Drawing inspiration from conditional computation, we design a plug-and-play dynamic token selection mechanism to ensure the model can effectively allocate attention to different features based on task characteristics. Additionally, we employ an auxiliary predictor to alleviate the short-sightedness issue in the autoregressive sampling process. DM achieves state-of-the-art performance on various standard RL benchmarks while requiring significantly fewer computational resources, offering a viable solution for building efficient and scalable RL foundation models. Code is available at here.

Decision Mixer: Integrating Long-term and Local Dependencies via Dynamic Token Selection for Decision-Making
Hongling Zheng'
Li Shen Yong Luo
Deheng Ye
Bo Du
Jialie Shen
Dacheng Tao
Abstract
Standard Markov Environment
Non-lander landard Markov Environment
bosbrrrr Dause
Dasteon
The Conditional Sequence Modeling (CSM) paradigm, benefiting from the transformer's powerful distribution modeling capabilities, has demonstrated considerable promise in offline Re- inforcement Learning (RL) tasks. Depending on the task's nature, it is crucial carefully balance the interplay between inherent local features and long-term dependencies in Markov decision tra- jectories to mitigate potential performance degra- dation and unnecessary computational overhead. In this paper, we propose Decision Mixer (DM), which addresses the conflict between features different scales in the modeling process from the perspective of dynamic integration. Draw- ing inspiration from conditional computation, we design a plug-and-play dynamic token selection mechanism to ensure the model can effectively al- locate attention to different features based on task characteristics. Additionally, we employ an aux- iliary predictor to alleviate the sightenigtss issue in the autoregressive sampling process. DM achieves state-of-the performance various standard RL benchmarks while requiring signifi- cantly fewer computational resources, offering a viable solution for building efficient and scalable RL foundation models. Code is available at here.
Dasius Man
79
Figure I. The evaluation results for DT, DC, and DM in tasks with standard Markov properties and non-standard Markov properties. DM consistently secures the top performance across all tested environments, showcasing its superiority.
term temporal dependencies, becoming a foundational ar- chitecture in fields such as Natural Language Processing (NLP) (Brown et al., 2020; Achiam et al., 2023; Xu et al., 2021) and Computer Vision (CV) (Liu et al., 2021; Si et al.,
2022). Inspired by this success, Decision Transformer (DT) (Chen et al., 2021; Li et al., 2023) and its variants (Lee et al., 2022; Xie et al., 2023) introduce the transformer to the field of offline Reinforcement Learning (RL), demon- strating its powerful capabilities in Conditional Sequence Modeling (CSM). Specifically, DT integrates cumulative rewards, states, and actions into tuple and trains of- fline datasets autoregressively to output appropriate actions. DT's limitation lies in overlooking the inherent local as- sociations between adjacent timestep tokens in offline RL data, which are crucial for the model's learning transition and reward functions. Decision Convformer (DC) (Kim et al., 2024) employs causal convolution filters instead of attention modules for data modelling to address this limi- tation and outperforms DT on specific tasks. However, its overemphasis on capturing local features leads to subop- timal performance on long-term sequence modeling tasks. Such tasks typically do not adhere to the standard Markov property indicated by DC but rather follow a non-standard Markov process (Ching al., 2013), where the current state may depend multiple previous time steps. Consequently, fixed convolution range may fail to adapt to the specific characteristics of different tasks (Figure 1). Furthermore, existing works rarely consider the presence of suboptimal segments within trajectories, and trajectory selection and concatenation would substantially improve the efficiency of
1. Introduction
Transformer (Vaswani, 2017) is widely regarded for its capacity to capture complex data distributions and long-
'School of Computer Science, National Engineering Research Center for Multimedia Software and Hubei Key Laboratory of Multimedia and Network Communication Engincering, Wuhan University, Wuhan, China "School of Cyher Science and Tech- nology, Shenzhen Campus of Sun Yat-sen University, Shenzhen, China 'Tencett Inc., Shenzhen, China 'Department of Computer Science, City St George's. ss University of London, London, UK "Nanyang Technological University, Singapore, Singapore. Corre- spondence to: Li Shen <shenli6 www.illssuuuuuuc... Yong Luo <luoyong @ whu.edu.en
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Decision Mixer: Integrating Long-term and Local Dependencies via Dynamic Token Selection for Decision-Making
both attention and convolution operations.
transformer, providing a direction for exploring the scaling law in offline RI.
Although increasing the number of parameters in foundation models has demonstrated success across a broad range of tasks (Bommasani et al., 2021; Muennighoff et al., 2023), this scaling law (Kaplan et al., 2020; Ye et al., 2020a:b) remains insufficiently explored in the context of deep RL. Blindly expanding the architecture of DT not only escalates computational demands but may also lead decline in performance. Conditional computation algorithms (Ainslie et al., 2023; Lei et al., 2023), which statically or dynami- cally allocate the number of parameters used during model training, appear promising for achieving computational ef- ficiency while enhancing model performance. However, whether such architectures can better capture the inherent Markov patterns in offline RL datasets while preserving the ability to learn long-term temporal dependencies remains unclear.
We demonstrate the effectiveness of DM through inten- sive experiments on broad spectrum of benchmarks, highlighting its competitive performance in Offline RL scenarios.
2. Preliminary
2.1. Conditional Computation
The transformer architecture has become a comeestone in driving the artificial intelligence revolution, but its high com- putational cost has sparked significant interest in improving efficiency. The concept of conditional computation (Bengio et al., 2015) has been proposed as promising approach to address the abovementioned issues, with learned mecha- nisms determining when and how computation is expended. Mixture of Experts (MOE) (Zhou et al., 2022) is a represen- tative method of conditional computation, where tokens are routed to one of several experts. This sparse activation archi- tecture maintains a constant total computation cost while ex- panding the number of parameters (Dai et al., 2024; Raposo et al., 2024; Zheng et al., 2025). Some other work (Schuster et al., 2022; Ainslie et al., 2023) focuses on the design of an early exit mechanism within the transformer, where the model learns to decide when to terminate the computation for given token, allowing the token to skip any remain- ing transformer layers once the exit decision is made. Our approach can be viewed as token-level sparse activation strategy. In this strategy, a routing mechanism similar to that in MoE, along with an hypernetorr,, determines whether a specific token should pass through the transformer layers or skip them.
In this paper, we propose Decision Mixer (DM), a low- complexity architecture that dynamically balances long- term dependency features and local Markov features through a token selection mechanism. For tokens in given se- quence, DM dynamically selects whether token should participate in the current layer's 's attention calculation or be passed directly to the next layer. Specifically, drawing in- spiration from the Mixture of Experts (MoE) (Zhou et al., 2022), DM utilizes a router to assign weights to each to- ken in the sequence. At the same time, a hypernetwork determines the number of tokens data features. The final selection is achieved by combining these two components with weight-based top-k mecha- nism. This dynamic token selection mechanism ensures that the model can allocate attention to long-term dependencies and local features based on the characteristies of the task. After being selected, tokens are fed into the attention layer in their original order, which can be viewed as a concatena- tion mechanism and allows the model learn features of optimal trajectories from suboptimal ones. The unselected tokens are passed through the residual link and then com- bined with the tokens processed through attention. During testing, since the entire sequence cannot be obtained at once in autoregressive sampling, we have designed an auxiliary network that directly predicts at the token level whether the token should participate in the attention computation.
be selected based on
Our approach also related to token dropping, which was initially proposed to reduce BERT inference costs (Press et al., 2022; Wang et al., 2021) and later adapted to improve training efficiency (Hou et al., 2022). Random-LTD (Yao et al., 2022) further advanced this idea by introducing random-layer token dropping combined with learning rate scheduling. While prior work focuses on static efficiency strategies for vision and language tasks (Zhong et al., 2023; Liu et al., 2024), they often lack dynamic adaptation, which can lead to semantic disruption. In contrast, DM dynami- cally selects tokens using a router and hypemneworr, align- ing with the Markovian nature of offline RL for improved performance. Its plug-and-play design and synchronized training offer greater flexibility than existing methods.
The main contributions of this work
follows:
We reconsidered the trade-off between capturing long- term dependencies and extracting local Markov fea- tures in CSM methods from an experimental perspec- tive, proposing a selection, stitching, and computation mechanism.
2.2. Conditional Sequence Modeling for Offline RI.
Drawing inspiration from MOE architecture, we signif- icantly reduced the computational burden by dynam- ically selecting important tokens at each layer of the
In contrast to online RL, offline RL (Agarwal et al., 2020; Wei et al., 2021; Qu et al., 2023) focuses on training models and performing trial-and-error using offline data without
2

---

Decision Mixer: Integrating Long-term and Local Dependencies via Dynamic Token Selection for Decision-Making
environmental interaction arrive appropriate strategies. Recently, conditional sequence modeling for RL (Hu al., 2024b; Branddonbrener et al., 2022), represented by the transformer architecture, has further demonstrated the ad- vantages of data-driven policy learning. DT (Chen et al., 2021) is trained on an offline dataset of triplets encapsu- lating return-to-go Vt, state 8t, and action úe, and outputs the optimal action. The rt token quantifies the cumulative reward from the current time step the end of the episode. During training, DT processes a trajectory sequence in auto-regressive manner, which encompasses the most recent K-step historical context:
As shown in Figure 2, compared the halfcheetah-mediing the attention matrices obtained from the maze2d-umaze dataset exhibited strong correlations between tokens that are far apart. This observation aligns with our intuition that, in non-standard Markov tasks, the relationships between tokens are not strictly local. Additionally, we perform a visualization analysis on datasets of varying quality from the same task and discover that the relationships between tokens in lower-quality data sequences more chaotic and unevenly distributed. Therefore, feasible concatenation mechanism could help the model capture compelling local or long-term features from suboptimal data.
T (fr- K+1, 01-K1, fr, 81, ax) at)
Our proposed Decision Mixer aims to use a dynamic token selection mechanism concatenate high-quality sequences from suboptimal trajectory sequences. The sequence's 's qual- ity and length that enter the attention layer reflect balance between local and long-term features, and the model can decide how to select and concatenate tokens based on the task's specific characteristics. In cases where the sequence length is extremely short (only two adjacent time-step to- kens are selected and concatenated), DM degenerates into a DC-like architecture that utilizes the attention mechanism capture more extreme local features than DC. When the sequence length matches the original context length, DM degenerates into DT.
The prediction head associated with state token 8 is trained to predict the corresponding action dz Subsequent work has made various improvements to DT, including prompt tuning (Xu et al., 2022; Zheng et al., 2024), tra- jectory concatenation (Wu et al., 2024), and value regular- ization (Chebotar et al., 2023; Hu al., 2024a). These approaches often involve more complex modifications of DT to adapt specific tasks.
The most relevant work to ours is Decision Con- vformer (Kim et al., 2024), which solely utilizes convolu- tional filters to capture local patterns but shows limitations in tasks requiring long-term dependencies. Additionally, the fixed and finite convolution range limits its performance on tasks with non-standard Markov features. Our approach's distinction lies in designing innovative token selection mechanism, which dynamically balances the trade-off be- tween -org-term and local features by retaining the attention mechanism. This innovation addresses the inherent limita- tions of DT and DC, is orthogonal to previous methods, and can be seamlessly integrated into existing architectures.
maze2d-medium
halfcheetah-medium
10
20
20
20-
20
1:
32
hallcheetah-meduu-eeppprrr
halrchhetta--nedll--eellaaa
2.3. Rethinking the Trade-offs of Features in CSM
RL tasks with the standard Markov property imply that the current state contains all the information needed to predict future states. Therefore, using fixed convolutional kernels capture features across adjacent time steps has performed well in such tasks. However, in scenarios where the Markov property is standard, the current state depends not only on the previous state but also on multiple past states or ac- tions. In such cases, focusing on the long-term dependencies the data is crucial for ensuring the model's performance. LSDT (Wang al.. 2025) experimentally integrates self- attention and dynamic convolution via a branch design for decision- decision-making, further demonstrating the effectiveness of combining features at different scales. We visualized the attention scores of the DT's first layer on two tasks: Gym HalfCheetah (with standard Markov properties) and Maze2D (with non-standard Markov properties) to show the degree of token associations across sequences.
10
33
20
20
50
Figure 2. Visualization of attention scores in the first layer of the model for Maze2D and HalfCheetah tasks. We visualize the atten- tion scores for the top 30 tokens, with color depth representing the correlation between tokens (return-to-go, state, and action). These scores serve as alignment measures, indicating the strength of as- sociation between each target token and source token. The Query index (or Key index j) is ordered such that i corresponds ft- K41, = 2 corresponds to state 8t- K+I, i = 3 corresponds to action a:- +1. and so on, continuing until the end. Given the application of causality, the attention matrix is lower-triangllar
3

---

Decision Mixer: Integrating Long-term and Local Dependencies via Dynamic Token Selection for Decision-Making
3. Methodology
selected tokens are arranged in the original order to form X which is then passed through the first subblock, con- sisting of layer normalization LN2, and an attention layer Aut, yielding Zí. Z is subsequently passed through the second subblock, comprising layer normalization LN2 and ferd-forward network FFN', to produce Z in Equation 3.
This section provides a comprehensive description of the proposed Decision Mixer. We analyze how CSM methods learn from historical trajectories from the perspective of reweighting. We integrate our dynamic token selection mechanism into the mixer layer to reduce computational burden while ensuring reasonable trade-off between long- term dependencies and local features. Considering the short- sightedness during sampling, we also design an auxiliary predictor trained in parallel to compensate for this limitation.
z Aut (LN(
7/2 = (LN, (3)
In the specific implementation, we multiply Z/j by the router weights and perform a generalized residual connection with the unselected tokens from X preserving the original order. By placing the router weights along the gradient path, we enable the router to receive feedback from gradient descent during training, which allows for updates. The token se- lection mechanism allows the DM to use the subsequent attention module to further mine the relationships between relatively important tokens, avoiding interference from irrel- evant or suboptimal tokens. Meanwhile, the re-connection of Z with the unselected tokens in X ensures that usable information continues be passed to the next layer, fur- ther guarantecing the retention of important context. We integrate the dynamic selection mechanism, attention mech- anism, and forward network into a single layer, referred to as the mixer layer, and show the data flow through it in Equation 4.
3.1. A Reweighting Perspective
We let the et be the set of the history of states, actions and rewards up to step & along with 8t. When the trajectory is evident from the context, we use to refer a generic set, with generic action. The essence of CSM is to learn distribution (a e). where 3 is the behavior policy that generated the data, and P3 refers to the joint distribution over states, actions, and rewards induced by 3. Here, f(e) is the conditioning function used during sampling to adjust the policy. It often chosen to be constant at the initial state and decrease with the observed reward along the trajectory. By factoring this distribution, we can express the optimal CSM for a specific conditioning function f(e) as: policy of
x!+ = (u//22))-((u E top-k)+ + (x--)) -lw, & top-k), (4) where Z/ represents the token at the i-th position in Zzz.
(a (al e) Ps(a e, f(e)) Pala ([(e) c.a) Ps(f(e) (e) ((( c.a).
Jf
Pa(a e) Ps(f(e)
From computational burden perspective, if reduce the length of the token sequence entering the attention layer to half of the original, the FLOP consumption during the key and value matrix multiplication process is reduced to 25% of the original. Similar calculations can determine the FLOP savings for the MLP.
CSM thus can be viewed as reweighting the behavior based on the distribution of future returns. DT directly uses the attention mechanism to learn the internal weighting rela- tionships in the data, However, the absence of explicit constraints means that suboptimal data can still disrupt the model's final decision. Filtering the data before the attention mechanism could better guide the model's learning, which motivates the design of the token selection mechanism.
3.3. Sampling Process
The token selection mechanism significantly reduces com- putational costs while enabling dynamic feature balancing. However, the generation of threshold i depends on the in- formation from the entire sequence, meaning that both past and future tokens influence the selection of specific token. While this is straightforward during training, the autore- gressive sampling, which generates tokens step by step, ob- structs the implementation of the token selection mechanism. To overcome this issue, we introduce an auxiliary predic- tor 0' that receives a single token as input and directly predicts whether the current token will be selected for the attention mechanism layer at the token level. Specifically, the auxiliary predictor outputs logits D: E R2, represent- ing the probabilities for token selection and non-sclection, respectively. The auxiliary predictor is trained separately and in parallel with the primary model training, using the token selection results from each round in real time. The
3.2. Mixer Layer
We draw inspiration from the expert routing mechanism in the MOE architecture. Given the input X' of length S = 33 at the current mixer layer 1, the router R₂ assigns a weight wf = (x+) to each token xf in XL. We also design a hyperneworrk Hi, which takes X as input and generates a threshold k- = (X'). We use simple MLPs to construct R₁ and H,, with the detailed architecture shown in Table B. Rú and H; enable the model to adaptively de- sign the threshold based on the data characteristics. This allows the model to select the top-k tokens for the attention mechanism based on the weights output by the router while skipping the unselected tokens in the current layer. The
4

---

Decision Mixer: Integrating Long-term and Local Dependencies via Dynamic Token Selection for Decision-Making
Decision Mixer
Mixer Layer
Output
X
X
Weights
Linear Predictor
Router R
FFN
Mixer Layer
Transformer Layer
Normalize
Weights
Training
Sampling
Mixer Layer
Self-attention
Hypernetwork
H1
Auxiliary e
Predictor mux
Transformer Layer
Normalize
Top-
Mask
Embedding Layer
X
X
Input
Figure 3. The architecture of Decision Mixer. During the training phase, the input sequence x .X!...., undergoes token selection based on the router Ri and the hypemetwork Hi for a specific mixer layer During sampling, DM selects tokens from the X based on the auxiliary predictor 0' Each mixer layer independently trains its own R., Hi, and Oun..
loss function for training is cross-entropy loss, applied to the predicted logits and the corresponding binary labels:
4. Experiment
In this section, we extensively evaluate our proposed De- cision Mixer using the widely recognized D4RL bench- mark (Fu et al., 2020). Our main objective is assess the effectiveness of DM across various domains. Additionally, we execute empirical ablation studies to dissect and under- stand the individual contributions of the core components of our method.
1 ΣE l log (is)) + (1 - 2;) log (1 (ú:)))
Laux
5
(5)
where (Wi)) ) is the sigmoid function that converts the logits d: into probabilities. z; is the binary label for each token in the mask, where 2, = indicates the selected token and 2, 0 indicates the unselected token. We concisely outline the pipeline of DM in Alg 2.
Datasets. We consider five different domains of tasks in the widely used D4RL benchmark: Gym, Adroit, Kitchen, AntMaze and Maze2D. A detailed introduction these five environments is presented in the appendix D.
3.4. Model Architecture
We adopt the same data format as DT Equation maxi- mize the algorithm's feasibility and highlight the advantages of our designed components. The input data is processed by structure formed by alternating mixer layers and stan- dard transformer layers, with the final layer incorporating a state-based prediction head to generate feasible actions. The purpose alternating layers is to ensure model sta- bility. Our ablation experiments observe that using only mixer layers as the main structure led significant fluc- tuations during training despite achieving similar overall performance. A potential explanation is that the dynamic to- ken selection mechanism prunes the entire sequence, which may cause parameter distribution to fluctuate. Intermittently deploying attention mechanisms over the entire sequence helps the model retain comprehensive input information and accelerates convergence. The loss function LDM is shown in Equation 6.
Baselines. We compare our approach with representative offline RL algorithms from value-based and CSM meth- ods. Each algorithm excels specific domains but per- forms sub-optimal in others. For value-based methods, including BEAR (Kumar al., 2019), BCQ (Fujimoto et al., 2019), CQL (Kumar et al., 2020), IQL (Kostrikov et al., 2022), TD3+BC (Fujimoto & Gu, 2021), MoRel (Ki- dambi et al., 2020), O-RL (Brandfonbrener et al., 2021) and COMBO (Yu et al., 2021). For CSM methods, in- cluding DT, DC, DD (Ajay et al., 2023), EDAC (An et al., 2021), D-QL (Wang al., 2023). MPPI (Pravitra al., 2020), StAR (Shang et al., 2022), GDT (Hu al., 2023). LSDT (Wang al., 2025) and CGDT (Wang et al., 2024). The performance scores for these baseline methods are sourced from the best results published in their respective papers or from our runs, ensuring fair comparison.
Implementation details. All experiments are carried out on server with 8 NVIDIA 3090 GPUs, each with 24GB of
LDM = EnDDly K
Σ (at (To (TT()))
(6)
imt K+1
5

---

Decision Mixer: Integrating Long-term and Local Dependencies via Dynamic Token Selection for Decision-Making
Table 1. The performance of DM and SOTA baselines on D4RT. Gym, Adroit, Kitchen, Maze2D, and AntMaze tasks. Results for DM
correspond to the mean and standard errors of normalized scores over 30 random rollouts (3 independently trained models and 10 trajectories per model) for tasks, which generally exhibit low variance performance. Our method outperforms all prior methods by a clear margin in almost all domains, including the conventional value-based and CSM methods.
Dataset
Value-Based Methods
Conditional Sequence Modeling Methods
Gym Tasks
BEAR
BCQ
CQL
IQL
ID3-BC
MoRel
StAR
GDT
CGDT
LSDT
DC
DM
halfchoctall-modee fitehhotll-mcduuu--paaa--2
38.6
34.8
37,5
44.1
44.6
40.2
36.6
36.8
40,5
40.4
42.9
41.3
39.6,02
hopper medium mediuu-reepay-
33.7
31.1
95.0
92.1
60.9
93.6
82.7
29.2
85,3
93.4
93.9
94.2
95.4.n
walleerd-mediun--rlla------
19.2
13.7
77.2
73.7
81.8
49.8
79.4
39.8
77.5
78.1
74.7
76.6
85.5
halfcheetah-medium-22
41.7
41.5
44.0
47.4
48.3
42,1
42.6
42,9
42,9
43.0
43.6
43.0
43.5
hopper-mediun-v2
52.1
65.1
58.5
63.8
59.3
95.4
67.6
59.5
77.1
96.9
87.2
92.5
98. 98.1.s6
walker2d medium v2
59.1
52.0
72.5
79.9
83.7
77.8
74.0
73.8
76.5
79.1
81.0
79.2
83.8
halfcheeluh mediumm expert-v2
$3.4
69.6
91.6
86.7
90.7
53.3
86.8
93.7
93.2
93.6
93.2
93.0
93.9_n1
hopper- noper--mdium--xxpprrt222
96.3
109.1
105.4
91.5
98.0
108.7
107.6
111.1
111.1
107.6
111.7
110.4
111
walker2d masi:rr---mmmuu--peeee--
40.1
67.3
108.8
109.6
110.1
95.6
108.1
109.0
107.7
109.3
109.8
109.6
112.7.1.3
Average
48.2
53.8
77.6
76.5
75.3
72.9
76.2
66.2
79.1
82.4
82.0
82.2
84.7
Adroit Tasks
BEAR
BCQ
CQL
IQL
O-RL
MoRel
EDAC
BC
DT
D-QL
StAR
GDT
DM
pen-human-vi
-1.0
66.9
37,5
71.5
90.7
-3.2
52.1
63.9
79.5
72.8
77.9
92.5
125.4
hammer-human--l
4.4
1.2
3.7
5.5
6.1.a1
door. murrhhmaan-v
2.2
0.05
9.9
4.3
-0.3
2.3
10.7
2.0
14,8
0.0
1.5
20.6
24.0.18
pen-eloned-vl
-0.2
50.9
39.2
37.3
60.0
-0.2
68.2
37.0
75.8
57.3
33.1
86.2
117.0.3.
lammer-ccknede
2.3
0.4
2.1
2.1
2.0
2.3
0.3
0.6
3.0
3.1
0.3
8.9
9.5,2x
door-cloned-vi
2.3
0.01
0.4
1.6
-0.1
2,3
9.6
0.0
16.3
0.0
0.0
19.8
23.7
Average
19.8
15.6
19.7
25.5
23.6
17.5
32.2
22.2
19.4
38.9
51.0
Kitchen Tasks
BEAR
BCQ
CQL
IQL.
O-RL
TD3+BC
BC
DT
DD
SIAR
GDT
DC
DM
kitchen complete-x0
43.8
62.5
50.8
40.8
43.8
40.9
65.3.a1
kitchen-pariia--00
13.1
18.9
49.8
46.3
35.5
0.0
33,8
57.9
57.0
12.3
73.3
66.8
75.0
Average
6.6
13.5
46.8
54.4
18.8
0.0
51.5
54.4
61.0
26.6
58.6
58.7
70.2
Maze2D Tasks
BEAR
BCQ
CQL
1QL.
TD3+BC
COMBO
BC
MPPI
DT
QDT
GDT
DC
DM
mazeld-umaze-vl
65.7
49.1
86.7
42.1
14.8
76.4
85.7
33.2
31.0
57.3
50.4
20.1
86.9.19
mauz2d-mddim--1
41.8
34.9
10.2
13.3
38.2
95.2g27
Average
45,35
33.1
64.3
38.5
38.5
72.5
63.6
21.7
19.6
35.3
29.1
57.6
91.1
AntMaze Tasks
BEAR
BCQ
CQL
IQL
TD3+BC
O-RL
BC
DT
RvS
StAR
GDT
DC
DM
antmaze-mmize-xo
73.0
78.9
74.0
87.1
78.6
64.3
54.6
59.2
65.4
51.3
76.0
85.0
100.0
uatmaze-umaaa-divrre--00
61.0
55.0
84.0
64.4
71.4
60.7
45.6
66.2
60.9
45.6
69.0
78.5
100.0.a.s
memuaz-medium-ddiee miieeiii---iivre--i
8.0
0.0
53.7
70.0
0.0
0.0
0.0
7.5
67.3
0.0
0.0
0.0
60.0_1.3
Average
47.3
44.6
70.6
73.8
50.0
41.7
33.4
44.3
75.0
32.3
48.3
54.5
86.7
memory. The experimental hyperparameter configurations of DM are shown in Appendix A.
datasets, demonstrating the effectiveness of our architecture. Although most baseline models demonstrate proficiency in the Gym environment, DM still exhibits outstanding perfor- mance on almost all tasks. DM's slightly lower performance compared to TD3+BC on certain datasets is primarily due to CSM's ' inability to adapt well to datasets with low quality or insufficient coverage of the state space in scenarios where the value function is completely absent. On the other hand, the Adroit environment is characterized by limited scope of human demonstrations, which leads to extrapolation er- rors that particularly challenge offline RL. is precisely for this reason that DM's excellent performance across all Adroit tasks can be attributed to its high expressiveness and more effective token selection mechanism. Both DC and DM demonstrate that capturing local information can sig- nificantly enhance DT's performance in standard Markov
4.1. Main Results
We compare our DM with the baselines on five domains of tasks and report the results in Table 1. To ensure fair comparisons, we normalize the scores according to the pro- tocol established in D4RL (Fu et al., 2020). where a score of 100 corresponds to an expert policy. We analyze the perfor- mance of DM on standard Markov tasks and non-standard Markov tasks separately.
Standard Markov Environment. We evaluate DM in Gym and Adroit environments with dense rewards. Our DM con- sistently achieves or approaches sota performance in all
6

---

Decision Mixer: Integrating Long-term and Local Dependencies via Dynamic Token Selection for Decision-Making
Table Ablation study on model components. For simplicity, we have removed the version numbers after each task, which does affect understanding All experiments are repeated three times, and the average value is taken.
Dynamic selection
Auxiliary predictor hopper-medium walker2d maikker2d-mmdduu-eexperrr pen-cloned maze2d-medium caanaazz-uazze
67.6
108.1
75.8
32.1
33,6
70.9
75.3
29.9
59.6
93.7
109.5
105.4
85.0
88.9
98.1
112.7
117.0
95.2
100.0
environments.
with our intuition: the fixed number of token selections pre- vents the model from adapting well to the varying nature of tasks and the complex distribution of sample quality,
Non-standard Markov Environment. We use the Kitchen, Maze2d, and more complex AntMaze environments to eval- uate DM's stitching and long-term credit assignment ca- pabilities on non-standard Markov datasets. For Kitchen tasks requiring generalization to unseen states and long- term value optimization, the DM outperforms the CSM and Value-Based Methods. These results demonstrate that DM can learn useful data features from offline trajectories, enhancing generalization and stability. For the Maze2d envi- ronment, which serves as benchmark to evaluate the capac- ity of offline RL algorithms to stitch segments of disparate trajectories effectively, the performance of DM significantly outperforms other methods, demonstrating the advantage of the token selection mechanism in stitching high-quality trajectories. The AntMaze environment is characterized by sparse rewards and many suboptimal trajectories, which presents an even more significant challenge. In this context, DC performs poorly due to its neglect of long-term infor- mation and an overreliance on data quality based on prior assumptions. The performance results of DM demonstrate the effectiveness and generalizability of the architecture we designed, particularly in antmaze-umaze-divent tasks.
Selected
Unselected
(a) hopper-medium
(b) walker 22d wwdimm--dddueeeeee expert
(c) mooze2d-mediung
antmaze-umaze
Figure 4. The visualization results of the token selection mecha- nism. The mixer and standard transformer layers arranged alternately, from bottom to top, corresponding to the first to the sixth layer. The input sequence length is S=60, with tokens marked in blue representing the selected tokens and those marked in beige representing the unselected tokens.
4.2. Ablation Study
Role of Different Components. As shown in Table 2, we observe significant decline in the model's performance when we do not use the auxiliary predictor. The main reason is that the autoregressive testing process leads to insufficient predictive scope for the router, causing biases that are par- ticularly noticeable in the early stages of testing when the number of visible tokens is limited. Additionally, increasing the sequence length of the training data is essential for en- abling the CSM method to exhibit performance comparable to that of language vision foundation models in RL tasks. The auxiliary predictor compensates for the performance degradation caused by the shorter data length available dur- ing sampling compared to the training data length. In the second row, we remove the dynamic token selection mecha- nism and fix the number of tokens model can select, setting 50% of the training trajectory length. This change al- lows the auxiliary predictor to function normally. However, we observe a significant performance drop, which aligns
Visualization of the token selection mechanism. To eval- uate the performance of the token selection mechanism across tasks with different properties, we visualized the to- ken selection at each layer for four tasks: hopper-medium and maaker-medium-expie ties), and maze2d-mediun and antmaze-umaze (with non- standard Markov properties). The visualization results in Figure 4 reveal that the selected tokens exhibit a certain degree of positional proximity in standard Markov tasks. In contrast, the positional proximity of selected tokens is lower in non-standard Markov tasks, which aligns with our moti- vation, as the spatial relationships between relevant tokens follow a discrete distribution in these tasks. Interestingly, selected tokens tend be picked in multiple mixer layers
(with standard Markov proper-

---

Decision Mixer: Integrating Long-term and Local Dependencies via Dynamic Token Selection for Decision-Making
hogper-medlum
hallchesähhmmddum-epperrt
walkeer2--medlu---ellay
maze2d-umaze
04
DT
DT
DT
DT
65
DC
DC
04
DC
DC
DM DM-ALL
DM
03
DM
DM
6.4
DM-ALL
DM-ALL
DM-ALL
0.3
02
02
011
61
00
0.0
20
48
60
90
100
43
60
90
100
100
153
300
20
32
Training Epochs
Training Epochs
Training Epochs
Training Epochs
Figure 5. Visualization of loss curves for DT, DC, DM, and DM with all mixer layers (DM-ALLL) across four tasks.
across all tasks, while some tokens tend to bypass all layers. We hypothesize that this phenomenon is closely related to the varying information quality within the sequence. We present the average selected token numbers for each mixer layer across all tasks in Table 10 and Figure 6. The mixer layer exhibits a disparity in the number of tokens selected for data of different qualities within the same task, which validates the potential of our method for high-quality trajec- tory selection.
of the environments and the training/testing task division is provided in Appendix D. The results in Table 4 affirm DM's comprehensive enhancements across test tasks. Further- more, the transformer architecture ensures that DM exhibits more substantial generalization capabilities as its parameter scale increases. We also conducted an ablation study on the context length K used in DM (Table 11), which shows performance improvements with an appropriate increase in context length. However, increasing K beyond a cer- tain threshold, such K 120, may lead to performance degradation due to the prevalence of suboptimal trajectories.
Computational complexity. Table 3 shows the memory usage, parameter size, and FLOPs for several methods in the medeer- task. The FLOPs of DM are nearly halved compared to the original DT. Although DM has more pa- rameters than DC, its unique token selection mechanism results in fewer FLOPs during training. Furthermore, due to the better alignment of transformers with the scaling law, DM outperforms the convolution- DC regarding scala- bility and generalization. When we replace all transformer layers with mixer layers (DM-ALL), memory usage and FLOPs are further reduced. DM adopts a more suitable approach for offline RL data characteristics, achieving im- proved performance while reducing computational resource requirements. Figure 5 presents the training loss curves for several methods. We observe that DM-ALL exhibits significant fluctuations during training, which inspires to choose alternating layers as the core architecture for DM to stabilize the training process.
Table 4. Ablation on the zero-shot generalization ability. We de- note the model size as having x layers, y attention heads, and an embedding dimension of z, represented as (1, y, z) in the table.
Model Size
Game
BC
DT
DC
DM
Cheetah-vel
147.41
-138.05
-154.20
-135.62
Ant-dir
171.91
169.64
173.32
166.91
(3,1,256)
MetaWorld 10
350.36
352.67
339.09
373.53
MetaWorld 50
275.93
280.57
257.91
305.28
Cheetah-vel
-161.84
-148.24
-146.67
-139.18
Ant-dir
168.33
170.66
142.09
167.73
(6,4,256)
MetaWorld 10
337.40
343.16
255.21
345.32
MetaWorld 50
297.14
279.70
290.73
319.36
Cheetah-vel
140.59
-137.54
-142.30
-134.67
Ant-dir
165.92
169.50
169.73
177.24
(12,12,768)
MetaWorld 10
327.48
330.25
257.86
365.04
Table 3. Ablation on the computational complexity.
MetaWorld 50
291.07
313.82
308.80
347.53
Complexity
DT
DC
DM-ALL
DM
A %
Memory +
7128M
4830 M
3004M
4474M
4 37.2%
Params t
43.3M
28.5M
43.6M
43.5M
0.5%
Combination with Other Methods. DM addresses the in- herent trade-off problem in CSM methods, enhancing DT's performance from foundational perspective. It is worth noting that DM is orthogonal to most previous works rather than confticting with them and can improve the performance of past methods in plug-and-play manner. We select two methods to demonstrate this. QT (Hu et al., 2024a) intro- duces Q-value regularization to optimize action selection on top of DT and excels in handling long time horizons and sparse reward tasks. We refer to the QT enhanced with
Flops
752.5 G
436.9G
330.1G
398.5 G
J 47.0%
Ability to unseen tasks. The model's performance on un- seen tasks is key criterion for evaluating its generalization ability and potential scalability. Inspired by Prompt-DT (Xu et al., 2022), we conduct experiments in three meta-lcarning environments: Cheetah-vel, Antdir, and Meta World. After training on the training tasks, we test the model on tasks that do not overlap with the training tasks. description
8

---

Decision Mixer: Integrating Long-term and Local Dependencies via Dynamic Token Selection for Decision-Making
DM as QDM. ODT (Zheng et al., 2022) combines offline pretraining with online fine-tuning by introducing entropy regularization and hindsight experience replay mechanisms, optimizing the policy to adapt to tasks. We refer to the ODT enhanced with DM as ODM. As shown in Table 5, the meth- ods enhanced with DM outperform the original methods on nearly all tasks, indicating that the DM architecture can serve as an efficient and feasible auxiliary mechanism to further improve model performance without compromising the advantages of the original model.
iliary predictor is updated iteratively using predictions from the hypernetwork and router, effectively mitigating potential distribution shifts and ensuring stable performance.
5. Conclusion
In this study, we propose Decision Mixer (DM), a novel approach that addresses the challenge of modelling both long-term dependencies and local Markov properties in offline RL tasks. DM optimizes attention allocation and significantly reduces computational complexity by dynami- cally selecting and concatenating tokens at each layer of the transformer architecture. Extensive experiments on standard RL benchmarks demonstrate that DM outperforms existing methods and highlights its ability to effectively balance fea- ture conflicts, providing a viable path forward for scaling RL models with reduced computational overhead.
Table 5. Ablation on the model portability. DT architecture in QT and ODT replaced with DM architecture to examine whether the DM method can enhance other methods in non-confiicting manner. For simplicity, we have removed the version numbers after each task, which does not affect understanding.
Datasets
DT
DM
QT
QDM
ODT
ODM
halfcheetah medium replu
36.6
39.1
48.9
49.3
40.4
41.1
hogper-medium-rellyy
82.7
95.4
102.0
104.5
83.9
95.9
Further enhancing DM's robustness, particu- larly in incomplete or noisy data scenarios, could improve its adaptability to diverse situations.
Limitation.
walker2d 2d-redium--mepppa
79.4
85.5
98.5
90.6
76.9
80.3
balfcheeluh medium
42.6
43.5
51.4
$2.9
42.7
43.8
hopper-medium
67.6
98.1
96.9
99.4
97.5
98.3
walker2d mediumm
74,0
83.8
88.8
90.8
76.8
81.0
86.8
93.9
96.1
97.2
87,1
90.7
Acknowledgements
hoper--meiim--epeett
107.6
111.8
113.4
113.5
111.0
112.3
walker2d wederndddddiimmppppee expert
108.1
112.7
112.6
114.2
109.6
111.0
This project is supported by the STI 2030- MajoorMajoroooo Projects (No. 2021ZD0201405), the National Natural Science Foun- dation of China (Grant No. U23A20318, 62276195 and 622251 (2225113), the Shenzhen Basic Research Project (Natu- ral Science Foundation) Basic Research Key Project (NO. JCYJ2024 211222212021220 220112222224400 124430041), the Science and Technology Major Project of Hubei Province (Grant No. 2024BAB046), the Tencent JR2025TEG002,.te Foundation for Inno- vative Research Groups of Hubei Province (Grant No. 2024AFA017) and the National Research Foundation, Singa- pore, under its NRF Professorship Award No. NRF-P2024- 001. The numerical calculations in this paper have been done on the supercomputing system in the Supercomputing Center of Wuhan University.
Average
74.6
84.7
91.4
94.4
82.7
84.7
inffeennnnieeeecee strategy comparison. We designed three schemes: (1) The hypernetwork uses single token input for training and inference without needing an auxiliary predictor. (2) The hypernetwork uses a causal sequence as input for training and inference without needing an auxiliary predictor. (3) The hypemetwork uses the entire sequence as input for training, providing data for the binary classifica- tion training of the auxiliary predictor, which then used for inference (adopted by DM). As shown in Table 6, (3)
Table 6. Ablation on training-inference strategies. We compare three schemes to investigate the impact of consistent or hicrarchical training-inferene designs on task performance.
Impact Statement
Datasets
(1)
(2)
Decision Mixer contributes to developing more efficient and scalable offline RL foundation models, enabling broader applications in autonomous systems, robotics, and decision- making processes. While the immediate societal implica- tions are primarily related to improving computational ef- ficiency in RL tasks, it is crucial to remain mindful of the ethical considerations surrounding deploying RL models in real-world scenarios. These include the potential risks of unintended biases in decision-making, the impact of au- tomated systems on jobs and industries, and the need for responsible AI practices to ensure fairness, transparency, and accountability in autonomous systems.
halfchhetthhmmedung
27.5
40.9
43.5
hopper-medium
82.7
94.7
98.1
malder2d-medium
55.2
83.4
83.8
maze2d-umaze
59.1
83.2
86.9
antmaze-umaze
80.3
75.0
100.0
Average Score
61.0
75.4
82.5
demonstrated the best performance and stability across all tasks. Although (1) and (2) maintained consistency between training and inference, they struggled with convergence due to insufficient utilization of global information from the training data. In contrast, (3) adopts a hierarchical task decomposition strategy for training and inference. The aux-
9

---

Decision Mixer:
Integrating Long-term and Local Dependencies via Dynamic Token Selection for Decision-Making
References
S., Radford, A., Sutskever, I., and Amodei, D. Lan- guage models are few-shot learners. In Proceedings of the 34th International Conference tion Processing Systems, NIPS '20, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 9781713829546.
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S.,
Neural Informa-
Anadkat, S., et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08772 2023.
Agarwal, R., Schuurmans, D., and Norouzi, M. An opti- mistic perspective on offline reinforcement learning. In International conference on machine learning, pp. 104- 114. PMLR, 2020.
Chebotar, Y., Vuong, Q., Hausman, K., Xia, F., Lu, Y., Irpan, A., Kumar, A., Yu, T., Herzog, A., Pertsch, K., al. Q-transformer: Scalable offline reinforcement learning via autoregressive q-functions. In Conference on Robot Learning, pp. 3909-3928, PMLR, 2023.
Ainslie, J., Lei, T., de Jong, M., Ontanon, S., Brahma, S., Zemlyanskiy, Y., Uthus, D., Guo, M., Lee-Thorp. J., Tay, Y., Sung, Y.-H., and Sanghai, S. CoLT5: Faster long- range transformers with conditional computation. In Pro- ceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 5085-5100, Singapore, 2023. Association for Computational Linguisties.
Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P.. Srinivas, A., and Mordatch, I. De- cision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 14:15084-1509,, 2021.
Ching, W.-K., Huang, X., Ng, M. K., Siu, T.-K.. Ching, W.- K., Huang, X., Ng, M. K., and Siu, T.-K. Higher-order markov chains. Markov Chains: Models, Algorithms and Applications, pp. 141-176, 2013.
Ajay, A., Du, Y. Gupta, A., Tenenbaum, J. B., Jaakkola, T. S., and Agrawal, P. Is conditional generative modeling all you need for decision making? In The Eleventh International Conference on Learning Representations, 2023.
Dai, D., Deng, C., Zhao, C., Xu, R., Gao, H., Chen, D., Li, J., Zeng, W., Yu, X., Wu, Y., Xie, Z., Li, Y., Huang, P., Luo, F. Ruan, C., Sui, Z., and Liang, W. DeepSeekMoF: To- wards ultimate expert specialization in mixture-of-experts language models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguis- ties (Volume I: Long Papers), pp. 1280-1297, Bangkok, Thailand, August 2024. Association for Computational Linguistics.
An, G., Moon, S., Kim, J.-H., and Song, H. O. Uncertainty- based offline reinforcement learning with diversified q- ensemble. Advances in neural information processing systems, 34:7436-7447, 2021.
Bengio, E., Bacon, P.-L.,, Pineau, J., and Precup, D. Condi- tional computation in neural networks for faster models. arXiv preprint arXi::1511.00233 2015.
Bommasani, R., Hudson, D. A., Adeli, E., Altman, R., Arora, S., von Arx, S., Bernstein, M. S., Bohg, J., Bosse- lut, A., Brunskill, E., et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.07258. 2021.
Fu, J., Kumar, A., Nachum, O., Tucker, G., and Levine, S. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
Fujimoto, S. and Gu, S. S. A minimalist approach to offline reinforcement learning. Advances in neural information processing systems, 34:20132-20141 2021.
Brandfonbrener, D., Whitney, W., Ranganath, R., and Bruna, J. Offline rl without off- off-policy evaluation. Advances in neural information processing systems, 34:4933-4946, 2021.
Fujimoto, S., Meger, D., and Precup, D. Off-policy deep reinforcement learning without exploration. In Interna- tional conference on machine learning, pp. 2052-2062. PMLR, 2019.
Brandfonbrener, D., Bietti, A., Buckman, J., Laroche, R., and Bruna, J. When does return- condin--ndiiionedd supervised learning work for offline reinforcement learning? Ad- vances in Neural Information Processing Systems, 35: 1542-1553, 2022.
Hou, L., Pang, R. Y., Zhou, T., Wu, Y., Song, X., Song, X., and Zhou, D. Token dropping for efficient BERT pretraining. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 3774-3784, Dublin, Ireland, 2022. Association for Computational Linguistics.
Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert- A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D. M., Wu,
J., Winter, C., Hesse, C., Chen, M.. Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish,
Hu, S., Shen, L., Zhang, Y., and Tao, D. Graph decision transformer. arXiv preprint arXiv:2303.03744 2023.
10

---

Decision Mixer: Integrating Long-term and Local Dependencies via Dynamic Token Selection for Decision-Making
Hu, S., Fan, Z., Huang, C., Shen, L., Zhang. Y., Wang, Y., and Tao, D. Q-value regularized transformer for offline reinforcement learning. In Proceedings of the 41st In- ternational Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp- 19165-19181, PMLR, 21-27 Jul 2024a.
Liu, T., Shi, L.,, Hong, R., Hu, Y., Yin, Q. and Zhang, L. Multi-stage vision token dropping: Towards effi- cient multimodal large language model. arXiv preprint 0tt::4411.1000. 2024.
Liu, 7.,, Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, 7., Lin, S., and Guo, B. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 10012-10022, 2021.
Hu, S., Shen, I., Zhang, Y., Chen, Y., and Tao, D. On trans- forming reinforcement learning with transformers: The development trajectory. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2024b.
Muennighoff, N., Rush, A., Barak, B., Le Scao, T., Tazi, N., Piktus, A., Pyysalo. S., Wolf, T., and Raffel, C. Scaling data-constrained language models. Advances in Neural Information Processing Systems, 36:50358-50311 2023.
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. arXiv preprint wrXiv:2011.0836/ 2020.
Pravitra, J., Ackerman, K. A.. Cao. C., Hovakimyan, N., and Thcodorou, E. A. Ll-adaptive mppi architecture for robust and agile control of multirotors, In 2020 IEEE/RSJ International Conference on Intelligent Robots and Sys- tems (IROS), pp. 7661-7666. IEEE, 2020.
Kidambi, R., Rajeswaran, A., Netrapalli, P., and Joachims, T. Morel: Model-basee offline reinforcement learning. Advances in neural information processing systems, 33: 21810-21823, 2020.
Kim, J., Lee, S., Kim, W., and Sung, Y. Decision con- vformer: Local filtering metaformer is sufficient for decision making. In The Twelfth International Confer- ence on Learning Representations, 2024. URL https: //openreview. wwtrriiwweetteeeeee forrm//ormmme wwt.nddddaaaaaaaaaammmmmma
Press, O., Smith, N., and Lewis, M. Train short, test long: Attention with linear biases enables input length extrapo- lation. In International Conference on Learning Represen- tations, 2022. URL https:/ opwnreview.eeeeeeeeeeeeeeeooee.a net// forum? (rrumidd==Bs@PPCCO
Kostrikov, I., Nair, A., and Levine, S. Offline reinforcement learning with implicit q-learning. In International Confer- on Learning Representationss 2022. URL https: //openreview net forum?iiioo 68n222d22dd
Qu, Y., Wang, B., Shao, J., Jiang, Y., Chen, C., Ye, Z., Line, L., Feng. Y., Lai, L., Qin, H., et al. Hokoff: Real game dataset from honor of kings and its offline reinforcement learning benchmarks. Advances in Neural Information Processing Systems, 36:221666 16:22166-22190, 2023.
Kumar, A., Fu. J., Soh, M.. Tucker, G., and Levine, S. Stabilizing off off-policy q-learning via bootstrapping error reduction. Advances in neural information processing systems, 32, 2019.
Raposo, D., Ritter. S.. Richards, B.. Lillicrap, T.. Humphreys, P. C., and Santoro, A. lixture-of-depthss Dynamically allocating compute transformer-based lan- guage models. arXiv preprint arXiv:2404.00 022558, 2024.
Kumar, A., Zhou, A., Tucker, G., and Levine, S. Con- servative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33: 1179-1191, 2020.
Schuster, T., Fisch, A., Gupta, J., Dehghani, M., Bahri, D., Tran, V., Tay. Y., and Metzler, D. Confident adaptive language modeling. Advances in Neural Information Processing Systems, 35:17456-17472 2022.
Lee, K.-H., Nachum, O., Yang, M. S., Lee, L., Free- man, D., Guadarrama, S., Fischer, I., Xu, W., Jang, E., Michalewski, H., et al. Multi-game decision transformers. Advances in Neural Information Processing Systems, 35; 27921-27936 2022.
Shang, J., Kahatapitiya, K., Li, X., and Ryoo, M. S. Star- former: Transformer with state-action-reaar represen- tations for visual reinforcement learning. In European conference on computer vision, pp. 462-479. Springer, 2022.
Lei, T., Bai, J., Brahma, S., Ainslin, J., Lee, K., Zhou, Y., Du, N., Zhao, V., Wu, Li, B., et al. Conditional adapters: Parameter-efficient transfer learning with fast inference. Advances in Neural Information Processing Systems, 36: 8152 8172, 2023.
Si, C., Yu, W., Zhou, P., Zhou, Y., Wang, X., and Yan, S. Inception transformer. Advances in Neural Information Prucessing Systems, 35:23495-23509 2022.
Li, W., Luo, H., Lin, Z., Zhang, C., Lu, Z., and Ye, D. A survey on transformers in reinforcement learning. Trans- actions on Machine Learning Research, 2023.
Vaswani, A. Attention all you need. Advances in Neural
Information Processing Systems, 2017.

---

Decision Mixer: Integrating Long-term and Local Dependencies via Dynamic Token Selection for Decision-Making
Wang, H., Zhang, Z., and Han, S. Spatten: Efficient sparse attention architecture with cascade token and head prun- ing. In 2021 IEEE International Symposium on High- Performance Computer Architecture (HPCA), pp. 97-110. IEFE, 2021.
Advances in Neural Information Processing Systems, 33: 621-632. 2020a.
Ye, D., Liu, Z., Sun, M., Shi, B., Zhao, P., Wu, H., Yu, H., Yang, S., Wu, X., Guo, Q., et al. Mastering complex control in moba games with deep reinforcement learn- ing. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 6672-6679, 2020b.
Wang, J., Karanasou, P., Wei, P., Gatti, E., Plasencia, D. M., and Kanoulas, D. Long-short decision transformer: Bridging global and local dependencies for generalized decision-making, In The Thirteenth International Confer- ence on Learning Representations, 2025.
Yu, T., Kumar, A., Rafailov, R., Rajeswaran, A., Levine, S., and Finn, C. Combo: Conservative offline model-based policy optimization. Advances in neural information processing systems, 34:28954-28888 28967, 2021.
Wang, Y., Yang, C.. Wen, Y., Liu, Y., and Qiao, Y. Critic- guided decision transformer for offline reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence. pp. 15706-15711 14, 2024.
Zheng, H., Shen, L.., Luo, Y., Liu, T., Shen, J., and Tao, D. Decomposed prompt decision transformer for efficient unseen task generalization. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024.
Wang, Z., Hunt, J. J., and Zhou, M. Diffusion policies as an expressive policy class for offiiee reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023, URL https: //tpe//eeeeeeeeeeeeeeeeeee.eeeer net/ met/forune forum? A-AFFDPP-
Zheng, H., Shen, L., Tang. A., Luo, Y., Hu, H., Du, B., Wen, Y., and Tao, D. Learning from models beyond fine-tuning, Nature Machine Intelligence, pp. 1 1-12, 2025.
Wei, H., Ye, D., Liu, Z., Wu, H., Yuan, B., Fu, Q. Yang, W., and Li, Z. Boosting offline reinforcement learning with residual generative modeling. In 30th International Joint Conference Artificial Intelligence, IJCAI 2021, pp- 3574-3580. International Joint Conferences on Artificial Intelligence, 2021.
Zheng, Q. Zhang, A., and Grover, A. Online decision trans- former. In international conference on machine learning. pp. 27042-27059. PMLR, 2022.
Zhong, Q., Ding, L.,, Liu, J., Liu, X., Zhang, M., Du, B., and Tao, D. Revisiting token dropping strategy in efficient BERT pretraining. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 10391-10405, Toronto, Canada, July 2023, Association for Computational Lin- guistics.
Wu, Y.H.., Wang, X., and Hamaya, M. Elastic decision transformer. Advances in Neural Information Processing Systems, 36, 2024.
Xie, Z., Lin, Z., Ye, D., Fu, Q., Wei, Y., and Li, S. Future- conditioned unsupervised pretraining for decision trans- former. In International Conference on Machine Learn- ing. pp. 1888033 PMLR, 2023.
Zhou, Y., Lei, T., Liu, H., Du, N., Huang, Y., Zhao, V., Dai, A. M., Le, Q. V., Laudon, J., al. Mixture-of- experts with expert choice routing. Advances in Neural Information Processing Systems, 35:7103-711 101411 2022.
Xu, K., Zhang, Y., Ye, D., Zhao, P., and Tan, M. Relation- aware transformer for portfolio policy learning. In Pro- ceedings of the twenty-ninth international conference on international joint conferences on artificial intelligence, pp. 4647-4653, 2021.
Xu, M., Shen, Y., Zhang, S., Lu, Y., Zhao, D., Tenenbaum, J., and Gan, C. Prompting decision transformer for few- shot policy generalization. In international conference on machine learning, pp. 14641-2 PMLR, 2022.
Li, Yao, Z. Wu, Li, C., Holmes, C., Zhang, M., Li, C., and He, Y. Random-Itd: Random and layerwise token drop- ping brings efficient training for large-scale transformers. arXiv preprint arXiv:2211. 5166111861111 2022.
Ye, D... Chen, G., Zhang, W., Chen, S., Yuan, B., Liu, B., Chen, J., Liu, Z., Qiu, F., Yu, H., et al. Towards play- ing full moba games with deep reinforcement learning.
12

---

Decision Mixer: Integrating Long-term and Local Dependencies via Dynamic Token Selection for Decision-Making
Appendix
The appendix organized into several sections, each providing additional insights and details related to different aspects of the main work.
Hyperparameters Configuration
13
Network Architecture Details
13
Algorithm Pseudocode
14
Environment Details
14
D.I Main Environment
14
D.2 Meta-RL Enviromment
15
E Supplementary Experiment
16
A. Hyperparameters Configuration
Table 7. Common hyperparameters configuration of Decision Mixer.
Hyperparameters
Value
K (length of context)
20
training batch size learning rate weight dacay
512
1e-3
le-4
pet.traj number of layers number of mixer layers
number of transformer layers number of attention heads
embedding dimension activation dropout
256
GeLU
0.1
num_workers
64
B. Network Architecture Details
Table 8. Network architecture details.
Network
Layer
Input
Output
Linear
embed dim
Linear
lengtxt_length embed dim
512
Lea
Linear
512
0
Linear
embed_dim
embed_dim/1
SiLU
Linear
embed_dimm/12
2
13

---

Decision Mixer: Integrating Long-term and Local Dependencies via Dynamic Token Selection for Decision- Makingg
C. Algorithm Pseudocode
Algorithm Training Process of Decision Mixer
Input:
Training data D
Initialize: Model иTa, Router Rt, Hypernetwork Hi, Auxiliary predictor 0! for each epoch do
ouz
for each batch (Xo. yb) in training data do
Initialize X X₂ (input sequence)
for each layer = 1 to L do
Obtain the router weights: w (x.)
Compute the dynamic threshold:k = Hx()
Perform token selection:
X < Select tokens based on wa and k
Apply attention mechanism and Feed- (eed-Foowward Network by Equation
Update token sequence by Equation 4
end for
Train the auxiliary predictor independently:
Compute the Laux by Equation 5
Update 6 using Laux (independent training, no gradient propagation to TT))
Compute the LDM by Equation 6
Backpropagation to update Wg
end for end for
Algorithm 2 Sampling Process of Decision Mixer
1: Initialize: Model :T.. Router Rj, Auxiliary predictor 0lup
2: for each test sample X do
3: Initialize X X (input sequence)
4: for each layer 1 to L do
5: Obtain the router weights: = R
6: Use the auxiliary predictor e to decide token selection:
7:
1, + Auxiliary predictor output (whether token xa is selected)
8:
Apply token selection based the predicted f;
9:
Apply attention mechanism and Feed-Forwaal Network by Equation 3
10:
Update token sequence by Equation 4
11: end for
12: Compute predicted action à using X
13: end for
D. Environment Details
D.I. Main Environment
Gym tasks:
The Gym-MuJoCo tasks (Hopper, HalfCheetah, Walker2d) are popular benchmarks used in offline deep RL. They are relatively straightforward and characterized by datasets with significant proportion of near-optimal trajectories and smooth reward functions.
Adroit tasks: The Adroit domain involves controlling a 24-DoF simulated Shadow Hand robot to perform tasks such as hammering nail, opening door, twirling a pen, or picking and moving ball. This domain is chosen to study the impact narrow expert data distributions and human demonstrations on sparse-reward high-dimensional robotic manipulation tasks. Since these tasks are primarily derived from human behavior, they exhibit a limited state-action space, requiring robust policy regularization to ensure consistent agent performance.
14

---

Decision Mixer: Integrating Long-term and Local Dependencies via Dynamic Token Selection for Decision-Making
The Kitchen domain involves controlling 9-DoF Franka robot in kitchen environment with everyday household items such microwave, kettle, overhead light, cabinets, and oven. The goal is interact with these items to achieve a desired state configuration. This domain benchmarks the impact of multitasking behaviour in realistic, naviigaatn environment, where the "stitching" challenge arises from complex paths through the state space. Consequently, algorithms must generalize to unseen states rather than rely solely on training trajectories. The environment requires the agent to complete multiple sequential sub-tasks, further emphasizing the need for robust generalization.
Kitchen tasks:
Maze2D tasks: The Maze2D domain is navigation task in which 2D agent must reach fixed goal location. It tests offline RL algorithms' ability to stitch together previously collected sub-trajectories to find the shortest path to the goal. Three maze layouts are provided: the "maze," "medium," and "large mazes. These tasks evaluate the algorithm's capability to effectively combine sub-trajectories and identify the shortest path to the set goal.
AntMaze tasks: The AntMaze domain extends the Maze2D task by replacing the 2D ball with "Ant" quadruped robot, presenting a more demanding navigation challenge. This domain is introduced to test the stitching challenge with morphologically complex robot, better representing real-world robotic navigation tasks. The task uses a sparse 0-1 reward, activated upon reaching the goal.
complex 8-DoF
D.2. Meta-RI. Environment
There are 40 tasks in Cheetah-vel with different goal velocities. The target velocities are uniformly sampled from the interval [0,3]. The agent is penalized with 12 errors to the target velocity. We hold out 5 tasks to construct the testing set and train with the remaining 35 tasks.
Cheetah-vel:
Ant-dir: There are 50 tasks in Ant-dir with different goal directions uniformly sampled in 2D space. The 8-joints ant rewarded with high velocity along the goal direction. We sample 5 tasks for testing and leave the rest for training.
Mera-World MI10: Meta- World ML10, the task is to control a Sawyer robot's end-effector to reach a target position 3D space. The agent directly controls the XYZ location the end-effector. Fach task has different goal position. We train in 10 tasks and test in unseen 3 tasks.
Meta-Wort ML45: In Meta- World ML45, each task has different goal position. We train in 45 tasks and test in unseen tasks.
Table 9. Training and testing task indexes when testing the generalization ability in unseen tasks.
Cheetah-vel
Training set of size 35
1,3 6,8 14, 16 22, 24 25, 27 39)
Testing set size 5
2, 15, 23, 26]
Ant-dir
Training set of size 45
10-5,7 16, 18 22, 1001-29,333
29, 31
40, 42 49
Testing set of size 5
[6, 17, 30, 30,41) 41]
Meta- World ML10
Training set of size 10
[0,9, 19, 29, 33, 36, 39, 40, 48, 49]
Testing set size 3
11, 11,24,4 41]
Meta- World ML45
Training set of size 45
00 -- 10100,12 12 16, 18 24, 26 -3. -40, 42 49
Testing set of size 5
17, 25, 36, 41]
15

---

Decision Mixer: Integrating Long-term and Local Dependencies via Dynamic Token Selection for Decision-Making
E. Supplementary Experiment
Another perspective on the token selection mechanism. worth noting that token selection can also be viewed from the stitching perspective. Existing methods (Wu al., 2024; Hu al., 2024a) demonstrate effectiveness in stitching capabilities but often require additional optimization, such as complex representation learning objectives and statistical measures, which complicate the training process and increase the computational burden. In contrast, DM employs intuitive and straightforward token selection mechanism to reorganize the trajectories, capturing specific features that lead to performance improvements. Significant inconsistencies in sequence lengths are observed in standard and non-standard Markov task datasets of the same quality collected under the same policy. This further demonstrates that DM focuses on the internal features of the data rather than merely the quality of the trajectories themselves.
Table 10. The average number of tokens selected on the three mixer layers by the token selection mechanism when the model size is (6, 4, 256). For simplicity, have removed the version numbers after each task, which does not affect understanding.
Dataset
Ith Mixer Layer
2nd Mixer Layer
3rd Mixer Layer
Average
15.07
31.99
37.29
28.12
hopper-medium-rrelay
40.59
33.60
15.76
29.98
walker2d-medium-te
37.26
16.10
28.25
27.20
halfcheetah-meedimm
25.15
34.65
3.47
21.09
Gym Tasks
hopper-medium
36.54
38.62
6.29
27.15
walker2d-medium
30.48
44.51
5.04
26.68
halfcheetaa-mmccuu--experrr
37.95
35.36
31.70
34.67
hopper-mediim--exee
18.41
10.60
41.66
23.56
walker22dmmeiium-cee
33.81
23.64
45.49
34.31
pen-human
17.17
34.42
2.83
18.14
hammer-human
40.97
27.43
10.93
26.44
door-human
13.32
44.12
45.03
34.16
Adroit Tasks
pen-cloned
22.04
10.69
9.07
13.93
hammer-cloned
21.18
13.27
26.81
20.42
door-cloned
9.38
53.63
3.88
22.30
kitchen-complete
33.89
42.88
7.15
27.97
Kitchen Tasks
kitchen-partial
29.03
28.55
32.34
29.97
maze2d-umaze
19.74
19.98
22.67
20.80
Maze2D Tasks
maze2d-medium
21.08
32.77
41.23
31.69
antmaze-umaze
21.36
12.87
30.37
21.53
AntMaze Tasks
autmaze-umazz-divvrrse
21.70
20.27
32.00
24.66
antmaze-meddiun-divesss
10.56
40.59
57.40
36.18
Impact of context length K As shown in Table 11, conducted experiments in the Gym environment with DM context lengths of 8, 20, 60, and 120, keeping the other parameters consistent with those in Table A. We observed that DM performance improved to varying degrees with increasing context length, indicating that DM, like DT, exhibits excellent extendability. However, when K was set to 120, the model's performance decreased compared to K=60. This highlights that an indiscriminate increase in context length can negatively impact model performance. Therefore, is important to carefully analyze the relationship between data scale and data quality and design appropriate hyperparameters and architecture sizes
16

---

Decision Mixer: Integrating Long-term and Local Dependencies via Dynamic Token Selection for Docision-Making
Tokents Selected
Laper
helfo heficheeeettl www.eeeeaa-eeeeiiteeeee meteeeeeiinnnn cnnerrrrrpey epply
happer modium replay
halccheethh mod moddum
hoggar mondmmddda
halfchreeaah eegeet
wall walker wwwaardddddddiiidd mes---maddu medium dappert
pan Nusuaa
per-olored
hammer- hammer-obreed coreed
door cloodd
compless
mazze Pezeett-ummoos
mazzede medemmedda
antmaaze umaze
antrmazo diverss
walkerdd enedium
hopper wweeemeedmm
hommor
Avarage Selected
helfihheele meteeeeeeeeeeeeee spley
hatrchestan medium
happer-madiun aopar-madium
waterr 20-madium
halicheetah eepeeeeer
hopper expert
welkke 220 modi exporrt
pan- human
hammar- human
per- pen-ddoned
hammmer hammer--onnee doned
doar claned
knchonn complemm
mazedd amrecco
artimaze uimazze
klchett-pert
mazx24-mmdin
hopper medium
wakkrree medium
ant
Figure Average tokens selected across three mixer layers token selection mechanism.
Table 1. Experimental results of DM with different context lengths (K) in Gym tasks.
Datasets
DT(20)
DM (8)
DM (20)
DM (60)
DM(120)
sulfihhetta-medium--eee
36.6
36.9
39.6
39.2
39.5
hopper-medium replly
82.7
95.6
95.4
94.5
96.6
79.4
82.4
85.5
87.4
82.9
Siiichectah-mddinne
42.6
44.2
43.5
43.8
44.0
hopper-medium
67.6
96.0
98.1
99.1
99.7
falker2d-medium
74.0
85.1
83.8
84.1
85.5
selichhcctah-meddiie expert
86.8
91.1
93.9
92.4
93.6
hopper-medium-expper
107.6
112.2
111.8
112,4
113.2
walker2dd-medium-exper
108.1
114.0
112.7
114.6
110.6
Average
76.2
84.2
84.7
85.3
85.1
ensure the token selection mechanism fully function.
17

---

Decision Mixer: Integrating Long-term and Local Dependencies via Dynamic Token Selection for Decision-Making
Selected
Unselected
(a) meal
(b)
haffcheetch-meddum--eppaay
(c) walker2d-medium-replee
(d) haffcheetah-medlum
(e) walker2d-medium
(f) hopper-medium-expert
(g) halfcheetaa-mmediu--xeppett
(h) pen-human
(i) hommer-hummn
(j) door-human
(k) pen-cloned
(I) hammer-cloned
(m) door-cloned
(n) kitchen-complete
(o) kitchen-partial
(p) mazed2d-umaze
(q) antmaze-une untmaze-umazz-dde
(r) mawuraa--mmdiimmiiieeeeeee
Figure 7. The visualization results of the token selection mechanism. The mixer and standard transformer layers are arranged alternately. from bottom to top, corresponding the first to the sixth layer. The input sequence length is 60, with tokens marked in blue indicating those selected and tokens marked in beige indicating those not selected.
18


Reviews:
[{"review_id": "wfcmfRCyPJ", "reviewer": "Reviewer_5tKr", "summary": "The main contribution is a novel dynamic token selection mechanism termed Decision Mixer (DM), inspired by MoE to enhance CSM for offline reinforcement learning. DM adaptively selects key tokens for attention computation while preserving information from unselected tokens via feature concatenation, improving efficiency and mitigating information loss. Additionally, an auxiliary predictor in the autoregressive sampling process enhances long-term decision-making. Experiments show that DM achieves SOTA performance with reduced computational cost.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the careful review of our work. Due to the strict word limit, we have tried to address the reviewers' comments carefully. These additional experiments and suggestions will be incorporated into the updated main text.\n\n**W1: The robustness of the dynamic mechanism in sparse scenarios needs to be enhanced**  \nWe fully agree that there is still potential for improvement. An intuitive approach is to adopt data augmentation. Considering that complex tasks exhibit significant noise variations mainly at the environmental state level, we introduce Gaussian noise perturbation to all state dimensions in the training data, formulated as:\n$$\\tilde{s}_j=s_j+\\alpha \\cdot \\sigma_j \\cdot \\epsilon_j.$$\nHere, $\\alpha$ is a controllable noise intensity, which we set to 0.05. $\\sigma_j$ represents the standard deviation of the $j$-th dimension, which has been computed during data preprocessing. $\\epsilon \\sim \\mathcal{N}(0, I)$ denotes a noise vector sampled from the standard normal distribution. We generate two random seeds to add noise to the initial data $\\tau$, obtaining two noisy datasets, $\\hat{\\tau_1}$ and $\\hat{\\tau_2}$. Each training iteration uses a total of $3 \\times$ batch data. $\\tau$, $\\hat{\\tau_1}$, and $\\hat{\\tau_2}$ share identical values except for the perturbed states. The processed data is then used for training and referred to as DM_enhanced. The experimental results on three tasks are as follows:\n||DM|DM_enhanced\n|-|-|-|\nhopper-medium|98.1 $\\pm$ 3.6|98.0 $\\pm$ 1.9 \nmaze2d-umaze| 86.9 $\\pm$ 1.9|88.9 $\\pm$ 1.3\nantmaze-umaze|100.0 $\\pm$ 0.5|100.0 $\\pm$ 0.1\n\nDM_enhanced performs better than DM on the maze2d-umaze task, with a minor standard deviation. Although DM_enhanced does not perform better on the reward-dense task hopper-medium, a similar stability improvement is observed. This phenomenon suggests that the data augmentation employed in DM_enhanced enhances its stability and robustness. Due to rebuttal time constraints, exploring more instructive designs for robustness suppression will be part of our future work.\n\n**Q1: Could the auxiliary predictor make incorrect selections due to distribution shift**\n\nThe only potential distribution shift is the difference between the training and inference inputs. We designed three schemes: (1) The hypernetwork uses a single token as input for training and inference without needing an auxiliary predictor. (2) The hypernetwork uses a causal sequence as input for training and inference without needing an auxiliary predictor. (3) The hypernetwork uses the entire sequence as input for training, providing data for the binary classification training of the auxiliary predictor, which is then used for inference (the approach adopted by DM).\n\n||(1)|(2)|(3)| \n-|-|-|-\n|halfcheetah-medium|27.5 $\\pm$ 5.1|40.9 $\\pm$ 0.6|43.5 $\\pm$ 0.7| \n|hopper-medium|82.7 $\\pm$ 18.0|94.7 $\\pm$ 5.0|98.1 $\\pm$ 3.6| \n|walker2d-medium|55.2 $\\pm$ 4.9|83.4 $\\pm$ 2.7|83.8 $\\pm$ 0.8| \n|maze2d-umaze|59.1 $\\pm$ 4.7|83.2 $\\pm$ 5.8|86.9 $\\pm$ 1.9|\n|antmaze-umaze|80.3 $\\pm$ 12.4|75.0 $\\pm$ 9.1|100.0 $\\pm$ 0.5| \n|average score|61.0|75.4|82.5|\n\nScheme (3) demonstrated the best performance and stability across all tasks, with significantly shorter training times. Although schemes (1) and (2) ensured consistency between training and inference, they struggled with convergence due to insufficient utilization of global information from the training data. Scheme (3) adopts a task decomposition approach to implement training and inference hierarchically, with the auxiliary predictor trained based on the prediction data from each round of the hypernetwork and router. Figure 1 of the anonymous pdf shows the training loss curves, which indirectly suggest that binary classification of data on a specific Mixer layer for a given task is relatively easy to learn. Experimental results in Table 1 confirm that potential distribution shifts were addressed through synchronized training.\n\n**Q2: How does the \"dynamic token selection mechanism\" reduce computational overhead**  \nIn a specific Mixer Layer $l$, the dynamic token selection mechanism reduces computational complexity by adaptively selecting the top-$k$ tokens (via router $R^l$ and hypernetwork $H^l$) for attention processing while skipping others. When the sequence length input to the attention layer is reduced from $S$ to $k$, the attention complexity decreases from $O(S^2d)$ to $O(k^2d)$, and the computational complexity of the FFN layer reduces from $O(Sd)$ to $O(kd)$, where $d$ is the hidden dimension. For example, in the Key/Value projection stage, if $k = S/2$, the computation is reduced to 25% of the original FLOPs. Experiments show that DM reduces FLOPs by 47.0% compared to DT and achieves better clock time performance, as shown in the table below.\n||DT|DC|DM\n|-|-|-|-|\nhalfcheetah-medium|$\\approx$ 12.0h|$\\approx$ 8.0h|$\\approx$ 7.5h\nhopper-medium|$\\approx$ 10.0h|$\\approx$ 6.5h|$\\approx$ 7.5h\nantmaze-umaze|$\\approx$ 10.0h|$\\approx$ 6.0h|$\\approx$ 6.5h", "claims_and_evidence": "Yes, the claims made in the submission are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "Yes, the proposed methods and evaluation criteria, including the benchmark datasets, are well-aligned with the problem and application at hand.", "theoretical_claims": "The factorization in Equation (2) correctly follows from the conditional probability decomposition, ensuring a proper reweighting mechanism based on future returns. The formulation aligns with standard importance sampling principles, making it a valid approach. The theoretical justification for this reweighting perspective is well-grounded.", "experimental_designs_or_analyses": "The study presents a well-structured evaluation of DM across multiple D4RL domains, supported by detailed ablation studies that effectively highlight the contributions of its core components. One notable strength is the computational complexity analysis (Table 3), demonstrating DM's efficiency in reducing memory usage and FLOPs compared to baseline methods like DT and DC. This analysis underscores DM's practical advantages and aligns with the broader goal of developing scalable and resource-efficient offline RL methods, and exploring how DM's efficiency scales with larger datasets or more complex environments would provide further insights into its applicability in real-world settings.", "relation_to_broader_scientific_literature": "Scaling law in offline reinforcement learning has not been explored well, especially regarding how CSM methods based on the Transformer architecture can maximize the performance advantages of Transformers, which is a fascinating question. While some previous works have scaled up the architecture and parameters of DT, they have done so at the cost of a proportional increase in computational resources. In contrast, DM proposes a feasible solution for computational efficiency while maintaining performance.\n\n[1] Lee K H, Nachum O, Yang M S, et al. Multi-game decision transformers. Advances in Neural Information Processing Systems, 2022.\n[2] Reed S, Zolna K, Parisotto E, et al. A generalist agent. arXiv preprint arXiv:2205.06175.", "essential_references_not_discussed": "All essential related works that provide the necessary context for understanding the key contributions of the paper have been cited and discussed.", "other_strengths_and_weaknesses": "Strengths:\n+ The paper tackles underexplored challenges in offline RL, such as handling non-standard Markov properties and improving generalization in suboptimal trajectories. These contributions are original and fill important gaps in the literature. DM's ability to handle standard and non-standard Markov tasks broadly applies to various RL benchmarks, including Gym, Adroit, Kitchen, AntMaze, and Maze2D.\n+ The paper is well-organized, with clear explanations of the motivation, methodology, and results. The use of visualizations enhances understanding of the model's behavior.\n\nWeaknesses:\n- While the dynamic token selection mechanism can adaptively adjust computational load, it may introduce selection bias in certain long-sequence or high-noise tasks (e.g., complex maze tasks with sparse rewards) due to fluctuations in router weights caused by local noise. Therefore, the robustness of the dynamic mechanism in extremely sparse scenarios still needs to be enhanced.", "other_comments_or_suggestions": "The authors could consider providing a more detailed explanation of this aspect, as the autoregressive nature of sampling may cause the input distribution to gradually deviate from the training data. Since the auxiliary predictor relies on token selection labels generated during training, could this distribution shift potentially affect the quality of the generated sequences?", "questions_for_authors": "1. The training of the auxiliary predictor relies on token selection labels generated during training.  However, the model generates sequences autoregressively during actual sampling, which may slightly cause the input distribution to deviate from the training data. In this case, could the auxiliary predictor make incorrect selections due to distribution shift, thereby affecting the quality of the generated sequences?\n\n2. How does the \"dynamic token selection mechanism\" mentioned in the paper reduce computational overhead? Specifically, how does it lower computational complexity by reducing the number of tokens entering the attention layer?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Sq9uvZgFqu", "reviewer": "Reviewer_ekS2", "summary": "This paper introduces Decision Mixer (DM), a select-concatenate-compute mechanism that improves efficiency in offline reinforcement learning. Inspired by MoE, DM dynamically filters key tokens for attention while retaining information from unselected ones. It also integrates an auxiliary predictor to mitigate short-sightedness. Experiments demonstrate that DM outperforms existing methods while significantly reducing computational overhead.\n\n## update after rebuttal\n\nI will keep my score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the careful review of our work. Due to the strict word limit, we have tried to address the reviewers' comments carefully. All additional experiments and suggestions will be incorporated into the updated text.\n\n**It is recommended to adjust \"MOE\" to \"MoE\"**  \nWe apologize for this mistake. To ensure consistency and accuracy of terminology, we will change all occurrences of \"MOE\" to \"MoE\" in the main text.\n\n**Q1: Should regularization terms be introduced**  \nWe appreciate the reviewer's thoughtful insights. To ensure our algorithm's simplicity and ease of deployment, we have solely used Equation (6) for the entire model training. Introducing additional regularization terms, such as a routing weight smoothing term, can improve training stability. To investigate this, we conducted experiments by adding a regularization term with a weight of 0.1 on top of Equation (6). The specific details of the regularization are as follows:\n$$L_{smooth}=\\lambda \\sum_{i=1}^{S-1}\\|w_i-w_{i+1}\\|^2$$\nThe experimental results are presented in the table below. We refer to DM with the added regularization term as DM_s.\n||DM|DM_s\n-|-|-\nhalfcheetah-medium|43.5 $\\pm$ 0.7|33.1 $\\pm$ 0.2| \nhopper-medium|98.1 $\\pm$ 3.6|90.9 $\\pm$ 1.0| \nwalker2d-medium|83.8 $\\pm$ 0.8|80.1 $\\pm$ 0.1| \nmaze2d-umaze|86.9 $\\pm$ 1.9|86.4 $\\pm$ 0.8| \nantmaze-umaze|100.0 $\\pm$ 0.5|48.0 $\\pm$ 0.8| \n\nWe found that while DM_s exhibited improved stability, it experienced varying degrees of performance degradation across all tasks compared to DM. We hypothesize that this is primarily due to the introduction of the weight smoothing term, which constrains the flexibility of token selection based on task-specific characteristics, making the selection process overly conservative. Abrupt changes in token selection across neighboring positions can be reasonable, as they allow the model to flexibly concatenate trajectories based on data characteristics while preserving the Markov properties required for specific tasks. In future work, we will explore more refined approaches to achieving efficient token selection more smoothly.\n\n**Q2: Is the dynamic range of k constrained**  \nThe range of values for $k$ is unrestricted to ensure the algorithm's simplicity. Depending on the nature of the task, $k$ can approach either 0 or the entire sequence length $S$. The medium-replay dataset (high noise), medium dataset (moderate noise), and medium-expert dataset (low noise) from specific tasks serve as references for evaluating the robustness of our model's selection, as shown in Table 8 and Figure 6. For the hopper-medium task, the average number of selected tokens fluctuates slightly across datasets of different qualities. A similar pattern is observed in the other two tasks, with the overall average number of selected tokens remaining between 20 and 30, approximately half of the entire sequence length, without significant degradation.\n\nTo further investigate the hypernetwork’s predictions in highly noisy environments, we introduced Gaussian noise sampled from a standard normal distribution to all tokens in the hopper-medium-replay and halfcheetah-medium-replay datasets while keeping the original action labels unchanged. We refer to these modified datasets as hopper-medium-noise and halfcheetah-medium-noise. The results show that the hypernetwork outputs a higher $k$ value for hopper-medium-noise, whereas the opposite trend is observed for halfcheetah-medium-noise. This suggests that the number of selected tokens adapts to task-specific characteristics in highly noisy environments. No significant performance degradation was observed during the experiments, indicating the robustness and stability of our approach.\n||1th Mixer Layer| 2nd Mixer Layer|3rd Mixer Layer|Average\n-|-|-|-|-\nhopper-medium-replay|40.59|33.60|15.76|29.98 \nhopper-medium-noise|45.94|37.66|29.83|37.81 \nhalfcheetah-medium-replay|15.07|31.99|37.29|28.12   \nhalfcheetah-medium-noise|10.04|25.37|7.99|14.47\n\n**W1/Q3: A more balanced discussion would strengthen the paper**  \nWe have summarized DM's advantages and limitations to help readers understand our approach better. By dynamically selecting and concatenating important tokens at each layer, DM reduces computational complexity while effectively balancing the trade-off between capturing long-term dependencies and extracting local Markov features. This approach enhances efficiency and offers valuable insights into scaling laws for offline RL. Notably, the dynamic token selection during inference relies on the auxiliary predictor for online decision-making, which may introduce latency in scenarios with strict real-time requirements. Additionally, DM performs slightly worse than value-based methods on low-quality or noisy data. Future work will explore data augmentation or more efficient and robust token selection strategies—such as adversarial training or noise-adaptive mechanisms—to improve adaptability to noisy or low-quality data.", "claims_and_evidence": "The claims made in the submission are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are suitable for the real-world problem.", "theoretical_claims": "Yes, the theoretical claims and proofs have been examined. The auxiliary predictor in Section 3.3, formulated with Equation (5), provides a valid approach to handle autoregressive sampling constraints.", "experimental_designs_or_analyses": "The experimental design is robust, with comprehensive benchmarking across diverse D4RL domains and thorough ablation studies to dissect the contributions of DM's components. A powerful aspect is the visualization of the token selection mechanism (Figure 4), which provides intuitive insights into how DM adapts to different tasks, such as the positional proximity of selected tokens in standard Markov tasks versus their discrete distribution in non-standard Markov tasks. This visualization aligns with the theoretical motivation behind DM and offers a clear, interpretable understanding of its behavior across varying task complexities.", "relation_to_broader_scientific_literature": "Previous work, whether DT or DC, had a certain degree of data bias. DM introduces a dynamic token selection mechanism that dynamically balances local and long-term dependencies. MoE inspires this mechanism and leverages a router and hypernetwork to select tokens for attention computation. This innovation allows DM to adaptively focus on relevant features, addressing DT's limitations and improving task performance with standard and non-standard Markov properties.", "essential_references_not_discussed": "No, all essential related works have been cited discussed.", "other_strengths_and_weaknesses": "Strengths:\n\nThe paper creatively combines ideas from CSM, MoE, and conditional computation to design DM. The dynamic token selection mechanism, inspired by MoE, is a novel and innovative approach to balancing local and long-term dependencies in offline RL. By significantly reducing FLOPs and memory usage compared to DT, DM addresses a key practical challenge in scaling RL models, making it more feasible for resource-constrained environments.\n\nWeaknesses:\n\nThe paper does not explicitly discuss the limitations of DM, such as potential failure cases or scenarios where it may underperform compared to baseline methods. A more balanced discussion would strengthen the paper.", "other_comments_or_suggestions": "It is recommended to adjust some instances of \"MOE\" to \"MoE\" in the main text to ensure consistency and accuracy in terminology.", "questions_for_authors": "1. The existing loss functions mainly focus on the MSE loss for action prediction. Considering the stability of the training process, should additional regularization terms, such as a routing weight smoothing term, be introduced to prevent abrupt changes in token selection between adjacent tokens?  \n\n2. Is the dynamic range of k generated by the hypernetwork constrained? When the input sequence contains significant noise, can the hypernetwork generate extreme values (e.g., k=0 or k = S), causing the model to degrade into a purely convolutional or attention-based architecture?  \n\n3. Could the paper provide a more balanced discussion by highlighting DM's limitations, such as scenarios where it may underperform, to guide future work?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kA8P5hBaEF", "reviewer": "Reviewer_xDKp", "summary": "This paper introduces Decision Mixer (DM), a Transformer-based architecture for offline reinforcement learning. DM features a dynamic token selection mechanism, where a routing module learns to selectively attend to relevant past tokens during training. To enable efficient inference, an auxiliary predictor is trained concurrently to approximate token importance without access to future information. Experimental results across a diverse set of offline RL benchmarks—including standard locomotion tasks and MetaWorld—demonstrate that DM consistently outperforms existing approaches in both performance and efficiency.\n\n## update after rebuttal\nConsidering the authors' response and other reviews, I have changed my score to weak accept.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the detailed review of our work. Given the strict word limit, we have carefully addressed the reviewer's comments and will incorporate the additional experiments and suggestions into the updated version. Anonymous pdf: https://anonymous.4open.science/r/Decision-Mixer-1068/rebuttal.pdf\n\n**The details of $R$ and $H$**  \nWe found that both $R$ and $H$ can perform well using simple MLP, as shown below. $R$ and $H$ are incorporated as part of the main model, with no additional constraints added, and the training is based solely on equation (6). Additional thoughts can be found in Q1 for Reviewer ekS2.\nNetwork|Layer|Input|Output\n-|-|-|-\n$R$|Linear|embed_dim|1\n$H$|Linear|context_length×embed_dim|512\n||LeakyReLU\n||Linear|512|1\n$θ_{aux}$|Linear|embed_dim|embed_dim//2\n||SiLU\n||Linear|embed_dim//2|2\n\nThe auxiliary predictor $\\theta^l_{aux}$ is trained with gradient isolation using dynamic selection results from $R$ and $H$. The main model selects the top-$k$ tokens from $X^l$ and generates binary labels $z_i \\in {0,1}$ as ground truth for $\\theta^l_{aux}$. The predictor outputs binary logits $\\hat{y}_i$, optimized via equation (5) to match $\\sigma(\\hat{y}_i)$ with $z_i$. We found the token selection distribution in binary classification easy to learn, with rapid loss convergence shown in Figure 1 of the anonymous pdf. Main experiments and ablation studies confirm prediction accuracy. Alternatives are discussed in our response to reviewer EJaz on \"Inconsistency of training and inference.\"\n\n**Lacks evaluation on specific environments**  \nWe have added experimental results for DM in the Atari, reporting the average performance across three random seeds. Table 1 of the anonymous pdf shows that DM performs well in discrete action environments, demonstrating the method's generalizability. This success was expected due to the precedents set by DT and DC.\n\n**More engage with the existing literature**  \nWe will provide a more detailed discussion in the main text. Token dropping, initially proposed to reduce BERT inference costs [1,2], was adapted by Hou et al. [3] to improve training efficiency. Random-LTD [4] advanced this with random-layer token dropping and learning rate scheduling. While prior work focuses on static efficiency strategies for vision and language tasks [5,6], they often lack dynamic adaptation, risking semantic disruption. In contrast, DM dynamically selects tokens using a router and hypernetwork, aligning with offline RL’s Markovian nature for better performance. Its plug-and-play design and synchronized training offer greater flexibility than existing methods.\n\n[1] Train short, test long...  \n[2] SpAtten: Efficient Sparse...  \n[3] Token dropping for...  \n[4] Random-ltd...  \n[5] Revisiting token dropping...  \n[6] Multi-Stage Vision Token...  \n\n**Further empirical evidence**  \nIn the paper, we have included baseline comparisons in RL scenarios, component ablation studies, token selection statistics and visualizations, computational efficiency, training curves, generalization, portability, and context length experiments. Following reviewer feedback, we added DM’s Atari results and adapted Random-LTD to DT in Gym. Table 2 of the anonymous pdf shows that Random-LTD underperforms DM in RL tasks, likely due to its reliance on random token dropping. We will attempt more transplant comparison experiments in future work.\n\n**Essential References Not Discussed**  \nWe will include LSDT in the related work and experimental comparison sections. LSDT integrates DT's self-attention and DC's dynamic convolution via branch design for decision-making. DM extends these environments and outperforms LSDT in most experiments (Table 3 of the anonymous pdf).\n\n**Other Comments**  \n- We apologize for the typo and will change \"MOD\" to \"model\".\n- We have placed the updated Figure 2 in the anonymous PDF.\n- We will replace the relevant terms with \\texit{k}.\n\n**Q1: Why does an auxiliary predictor improve performance**  \nThe auxiliary predictor addresses issues from inconsistent input data formats during training and inference. $H$ uses the complete sequence to predict $k$, but during inference, the autoregressive nature of the transformer makes the following sequence invisible, preventing sequence-level predictions. The token-level auxiliary predictor enables filtering without full sequence information. $R$ and $H$ jointly serve as teacher models, providing binary classification training data for the auxiliary predictor. This hierarchical training approach reduces the difficulty of training and fully utilizes all available data.\n\n**Q2: Include a clock-time measurement**  \nWe measured the clock time from training start to convergence across multiple tasks. The results in Table 4 (anonymous pdf) show that DM has a smaller time overhead than DT and is competitive with DC. Despite adding networks, DM's minimal parameters and dynamic token selection mechanism shorten the sequence length, avoiding significant clock time increases.", "claims_and_evidence": "Although the motivations and solutions seem to be novel, there are a number of unclear factors in the paper that required clarification: I find it unclear how the router R and the hypernetwork H were trained in the proposed approach. Specifically, the training process for these components is not well-explained in the manuscript. What is the architecture of these networks (not included in the manuscript) and whether the training process is only based on the downstream loss function in Eq (6)? Additionally, I am curious about how the auxiliary predictor is trained. More details are needed on what constitutes the ground truth for training this predictor and how its predictions are evaluated during the learning process.", "methods_and_evaluation_criteria": "The benchmarks that were for evaluation in this paper are diverse and suitable. These are benchmarks that have been used in previous DT related research.", "theoretical_claims": "There are no theoretical proofs in this paper.", "experimental_designs_or_analyses": "Experiments design have been checked, and they are reasonably formulated. The paper includes standard evaluation on Mujoco benchmarks, in addition, they also include evaluation on MetaWorld environments. However, the paper lacks evaluation on discrete-action-space environments (Atari games are often being used in previous DT research).", "relation_to_broader_scientific_literature": "- The paper introduces Decision Mixer (DM), a low-complexity architecture aimed at balancing local and long-range dependencies in Decision Transformer (DT) models through a layer-wise token selection mechanism. Inspired by the Mixture-of-Experts (MoE) architecture, the authors propose a router network to select the tokens.\n- While this method is novel in its application to offline reinforcement learning, token selection and dropping strategies have been explored in prior Transformer research. For example, Token Dropping for Efficient BERT Pretraining (arXiv:2203.13240) proposes dropping less important tokens mid-layer to improve training speed without sacrificing performance. Similarly, Random-LTD: Random and Layerwise Token Dropping (arXiv:2211.11586) presents a technique that skips computation for random subsets of tokens at intermediate layers to reduce cost.\n- The authors should more thoroughly engage with the existing literature on token dropping in Transformers, including a comparison and contrast to clearly articulate the novelty of their approach. Additionally, further empirical evidence is needed to convincingly demonstrate the superiority of the proposed method in reinforcement learning settings.", "essential_references_not_discussed": "Missing baseline: The paper “Long-Short Decision Transformer: Bridging Global and Local Dependencies for Generalized Decision-Making” acknowledges the same problem and proposes solution that should be compared here.", "other_strengths_and_weaknesses": "The paper presents a well-motivated problem, and the proposed solution seems to be novel in RL settting, particularly in its ability to dynamically select token weights. Empirical results across multiple tasks and environments demonstrate performance improvements when applied to various base models, such as DT, QDT, and ODT.", "other_comments_or_suggestions": "- There is a mention of \"MOD\" in line 343, but I was unable to find what it refers to.\n- For Figure 2, consider adding axis titles or specifying in the caption what the attention scores represent to improve clarity. For example, in DC, the relevant discussion section explicitly mentions that the attention scores reflect relationships between token level (state, action, and returns-to-go). Providing similar context here would make the figure more interpretable. \n - The term “threshold k” is somewhat misleading, especially when the output of hypernetwork is denoted as \\textit{k}. Changing top-k to top-\\texit{k} might help reader in understanding that this refers to the same value.", "questions_for_authors": "- Can you explain more about the results in Table 2. as to why incorporating such auxiliary predictor improves the overall performance compared to just dynamically token selection?\n- I believe it is necessary to include a clock-time measurement in addition to the complexity measurement, as the method involves the additional training of multiple networks. This would provide a more comprehensive assessment of the method's efficiency and resource requirements.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jW7QfJ1E9a", "reviewer": "Reviewer_EJaz", "summary": "This paper proposes adapting the attention block in transformers based on the Mixture of Experts (MoE) design to effectively balance capturing long-term dependencies and extracting local features. By adapting it to the Decision Transformer architecture, extensive experiments demonstrate its superior performance in the areas of offline reinforcement learning.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the careful review of our work. Due to the strict word limit, we have tried to address the reviewers' comments carefully. All additional experiments will be incorporated into the text.\n\n**Theoretical Claims**  \nWe provided a brief analysis of the CSM method from the perspective of re-weighting in equation (2). Given the word limit and the need for academic rigor, we have removed the statement \"theoretical analysis\" from the main text.\n\n**Concern 1: The inconsistency of training and inference**  \nWe carefully considered using a more consistent approach to achieve our objectives. We designed three schemes: (1) The hypernetwork uses a single token as input for training and inference without needing an auxiliary predictor. (2) The hypernetwork uses a causal sequence as input for training and inference without needing an auxiliary predictor. (3) The hypernetwork uses the entire sequence as input for training, providing data for the binary classification training of the auxiliary predictor, which is then used for inference (adopted by DM).\n||(1)|(2)|(3)\n-|-|-|-\nhalfcheetah-medium|27.5 $\\pm$ 5.1 (48.0h)|40.9 $\\pm$ 0.6 (20.5h)|43.5 $\\pm$ 0.7 (7.5h)\nhopper-medium|82.7 $\\pm$ 18.0 (48.0h)|94.7 $\\pm$ 5.0 (20.5h)|98.1 $\\pm$ 3.6 (7.5h)\nwalker2d-medium|55.2 $\\pm$ 4.9 (48.0h)| 83.4 $\\pm$ 2.7 (20.5h)|83.8 $\\pm$ 0.8 (7.5h)\nmaze2d-umaze|59.1 $\\pm$ 4.7 (24.0h)|83.2 $\\pm$ 5.8 (15.0h)|86.9 $\\pm$ 1.9 (6.9h)\nantmaze-umaze|80.3 $\\pm$ 12.4 (24.0h)|75.0 $\\pm$ 9.1 (16.0h)|100.0 $\\pm$ 0.5 (6.5h)\n\n(3) demonstrated the best performance and stability across all tasks, with significantly shorter training times. Although (1) and (2) ensured consistency between training and inference, they struggled with convergence due to insufficient utilization of global information from the training data. (3) adopts a task decomposition approach to implement training and inference hierarchically, with the auxiliary predictor trained based on the prediction data from each round of the hypernetwork and router. Figure 1 of the pdf (https://anonymous.4open.science/r/Decision-Mixer-1068/rebuttal.pdf) shows the training loss curves, which indirectly suggest that binary classification of data on a specific Mixer layer for a given task is relatively easy to learn. Experimental results in Table 1 and the ablation study in Table 2 confirm that potential distribution shifts were addressed through synchronized training.\n\n**Concern 2: The novelty and contribution seem to be incremental**  \nAll prior works have tended to introduce expert routing mechanisms in FFN or attention layers without involving token selection, where a fixed number of experts handle different tokens in the complete sequence. We emphasize that DM differs fundamentally from existing works in component design and training approach. DM innovatively designs (1) a dynamic token selection mechanism to address sequence modeling conflicts specific to offline RL, differing from conventional static MoE. DM handles incomplete sequences and uses generalized residual connections after each layer to ensure the consistency of output and input lengths. (2) During inference, we also designed a unique auxiliary predictor from the task decomposition perspective to address inconsistencies between training and inference. (3) We deploy the selection before the attention layer, making DM a more flexible plug-and-play architecture. It is more cost-efficient than conventional combinations of MoE and transformers, providing a reasonable direction for exploring scaling laws under the DT architecture.\n\n**Essential References Not Discussed**  \nWe will discuss the combination of MoE and transformers in the text. GShard[1] introduced MoE into transformers to address load imbalance via routing and expert capacity constraints. Switch Transformers[2] proposed a Top-1 gating mechanism to reduce computation and communication overhead. DeepSeekMoE[3] optimized expert utilization with fine-grained segmentation and expert isolation. MoEUT[4] combined MoE with the Universal Transformer, addressing the parameter-computation efficiency trade-off. Additionally, several works[5,6,7] explored MoE integration in visual tasks from the perspectives of data processing[5], multi-task learning[6], pre-training[7], and global training[8].\n\nWe discuss MoE because DM can be viewed as using a single expert to filter tokens before the standard transformer layer. However, this does not mean we have merely made a simple transplant. The collaboration between the router and the hypernetwork in DM, along with the auxiliary predictor and generalized residual connections, are seamlessly integrated, forming an efficient and practical framework that paves the way for future paradigms in offline RL.\n\n[1] Gshard...  \n[2] Switch transformers...  \n[3] Deepseekmoe...  \n[4] Moeut...  \n[5] Scaling vision with sparse mixture...  \n[6] Adamv-moe: Adaptive multi-task vision...  \n[7] Moe jetpack: From dense checkpoints...  \n[8] Mod-squad: Designing mixtures...", "claims_and_evidence": "Extensive experiments on standard benchmark datasets empirically demonstrate its superior performance compared to various offline RL methods, including both value-based and CSM-based approaches. Additionally, experiments highlight its computational efficiency over existing CSM methods, suggesting its potential to support scalable research in offline RL. However, the claim of theoretical consideration is unreasonable due to the absence of any formal theoretical analysis.", "methods_and_evaluation_criteria": "The module design is generally reasonable for its intended purpose. However, I have two concerns:\n\n1. The inconsistency of training and inference. During training, the selection of tokens passed through the attention block is determined by the hypernetwork, which considers the entire input sequence—an approach that is not feasible during inference. To address this, the authors introduce an auxiliary predictor to approximate the hypernetwork’s decision. However, this predictor operates on an incomplete sequence, missing crucial information available to the hypernetwork. This raises doubts about its ability to make accurate predictions. A more consistent approach might involve ensuring the hypernetwork also processes a causal sequence by masking future steps during training.\n\n2. The novelty and contribution seem to be somewhat incremental. Many studies have explored the combination of MoE and transformers to improve computational efficiency. While adapting this to the Decision Transformer is a reasonable extension, the paper's novelty and impact are limited due to insufficient discussion and comparison from a broader perspective beyond its specific application area.", "theoretical_claims": "Although the authors claimed theoretical consideration in Introduction, there is no theoretical analysis in this paper.", "experimental_designs_or_analyses": "The experiments are thorough and well-executed. The ablation study and computational complexity analyses effectively validate the strength of their design.", "relation_to_broader_scientific_literature": "Scaling offline RL is both a compelling and significant topic. However, there has been many works exploring combing MoE and transformer to enhance the computation efficiency. Although adapting it to the Decision Transformer is reasonable, the novelty and impact of this paper are limited due to the lack of broader discussion and comparative analysis beyond a specific scenario.", "essential_references_not_discussed": "Several previous works exploring combing MoE and transformer to enhance the computation efficiency:\n\nFedus, William, Barret Zoph, and Noam Shazeer. \"Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity.\" Journal of Machine Learning Research 23.120 (2022): 1-39.\n\nDai, Damai, et al. \"Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models.\" arXiv preprint arXiv:2401.06066 (2024).\n\nCsordás, Róbert, et al. \"Moeut: Mixture-of-experts universal transformers.\" arXiv preprint arXiv:2405.16039 (2024).\n\nLepikhin, Dmitry, et al. \"Gshard: Scaling giant models with conditional computation and automatic sharding.\" arXiv preprint arXiv:2006.16668 (2020).", "other_strengths_and_weaknesses": "The paper is well-written, providing clear explanations and precise notations.", "other_comments_or_suggestions": "No other comments.", "questions_for_authors": "No other questions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: Based on the strong consensus among reviewers regarding the paper's novelty within the RL context, significant empirical contributions (SOTA performance across diverse benchmarks), demonstrated computational efficiency, and the authors' thorough and successful rebuttal addressing initial concerns, I recommend accepting this paper. The proposed Decision Mixer offers a valuable and well-validated contribution to the field of offline reinforcement learning, particularly concerning efficient and effective sequence modeling. The authors have convincingly differentiated their approach from prior work and provided strong evidence for its utility.
Decision: Accept (poster)