Less is More: an Attention-free Sequence Prediction Modeling for Offline Embodied Learning
Authors: Wei Huang
Jianshu Zhang
Leiyu Wang
Heyue Li
Luoyi Fan
Yichen Zhu
Nanyang Ye
Qinying Gu
Venue: NeurIPS 2025 poster ()
Keywords: Offline Reinforcement Learning
Embodied Learning
Primary area: reinforcement_learning

Abstract: Offline reinforcement learning (offline RL) is increasingly approached as a sequence modeling task, with methods leveraging advanced architectures like Transformers to capture trajectory dependencies. Despite significant progress, the mechanisms underlying their effectiveness and limitations remain insufficiently understood. We conduct a thorough analysis on the representative Decision Transformer (DT) model using an entropy analysis and identify the inconsistencies in state-action-reward ($\langle s, a, R \rangle$) distributions causing attention ``dispersal". To address this, we propose a hierarchical framework that decomposes sequence modeling into intra-step relational modeling—handled by a Token Merger that fuses each $\langle s, a, R \rangle$ triplet—and inter-step modeling—handled by a Token Mixer across timesteps. We investigate several Token Merger designs and validate their effectiveness across various offline RL methods. 
Furthermore, our theoretical analysis and experimental results suggest that while Token Mixers are important, lightweight architecture can also achieve even better performance to more complex ones. We therefore propose a parameter-free Average Pooling Token Mixer, which, combined with a convolutional Token Merger, forms our final model, Decision HiFormer (DHi). DHi achieves a \textbf{73.6\%} improvement in inference speed and an \textbf{9.3\%} gain in policy performance on the D4RL benchmark compared to DT. DHi also generalizes well to real-world robotic manipulation tasks, offering both practical benefits and insights into sequence-based policy design for offline RL. Code and models are public at \href{https://wei-nijuan.github.io/DecisionHiFormer/}{project page}.

Less is More: an Attention-free Sequence Prediction Modeling for Offline Embodied Learning
Wei Huang!.
Jianshu Zhang!
Leiyu Wang'
Heyue Li*
Luoyi Fan
Yiehen Zhu
Nanyang Ye
Qinying Gu
Shanghai Al Laboratory
Tsinghua University SShanghai Jiao Tong University
WWuhan University *Midea Group
fhuanguei, gugiaying)apjiabb.rgg org..m
{j fjianshuzhang, leiyuwang33, luoyi_fan L ynylincoln) edu.
liheyue@whu. edu. cn yichen_ www@foxmail.ccmm
Abstract
Offlie reinforcement learning (offline RL) is increasingly approached as a se- quence modeling task, with methods leveraging advanced architectures like Trans- formers to capture trajectory dependencies. Despite significant progress, the mech- anisms underlying their effectiveness and limitations remain insufficiently under- stood. We conduct a thorough analysis on the representative Decision Transformer (DT) model using an entropy analysis and identify the inconsistencies in state- action-reward ((s, o, I?)) distributions causing attention "dispersal". To address this, we propose hierarchical framework that decomposes sequence modeling into intra-step relational modeling handled by Token Merger that fuses each (s, a, R) triplet-and and inter-step modeling- handlee by Token Mixer across timesteps. We investigate several Token Merger designs and validate their effectiveness across various offline RI. methods. Furthermore, our theoretical analysis and experimental results suggest that while Token Mixers are important, lightweight architecture can also achieve even better performance to more complex ones. We therefore propose a parameter-free Average Pooling Token Mixer, which, combined with a convolutional Token Merger, forms our final model, Decision HiFormer (DHi). DHi achieves a 73.6% improvement in inference speed and an 9.3% % gain in pol- icy performance on the D4RL benchmark compared to DT. DHi also generalizes well real-worke robotic manipulation tasks, offering both practical benefits and insights into sequence- based policy design for offline RL. Code and models are public at project page
Introduction
Offline Reinforcement Learning (offline RL) focuses on learning policies from fixed, static datasets without further environment interaction. Traditional methods often rely policy regularization or value function approximation. Recently, shift has occurred toward framing offline RI. sequence prediction problem, where policies are learned by predicting actions based sequences of past state-action-reward triplets. This reformulation enables the use of advanced sequence modeling techniques, particularly various token mixers, to capture dependencies within trajectories.
representative work is Decision Transformer (DT) which leverages Transformers 29 model entire trajectories through global self-attention and has inspired numerous follow-up methods 4
'This work was done when Wei Huang interned at Shanghai Artificial Intelligence Laboratory. Qiiyiing Gu and Nanyang Ye co-corresponding authors.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

10 12 14 IS 18
22 25 27 33], Despite its great success in this domain, some studies highlight its limitations in local dependency modeling. To address the Decision Converrmer (DC) 15 replaces self-attention with convolutional layers to enhance local feature extraction. StARformer 25 adheres to the self-attention mechanism by explicitly separating short-term and long-term dependencies through interleaved Step and Sequence Transformers, enhancing both local and global trajectory modeling. Similarly, Graph Decision Transformer (GDT) [12] uses a graph transformer to capture potential dependencies between adjacent states, actions, and rewards and differentiate the impact these different tokens.
The aforementioned previous studies demonstrate that the standard DT architecture is not well- suited for modeling local dependencies in a Markov process [T5] 22 (7). This is because the DT tends to consider all past context preventing it from effectively "attending" to useful states. To mitigate this, alternative operators such as convolutions, linear attention, and state-space models proposed to replace the self-attention layers for improved performance. However, several key questions remain unresolved: (1) Why is the attention mechanism in DT unable to "attend" useful states in the context? (2) Beyond performance scores, can we measure the model's ability to capture local dependencies using internal metrics? (3) Does offline RL exhibit significant differences across various operators, and are there simpler, more effective operators available?
To address the above questions, we make the following contributions: (I) We introduce atten- tion entropy as a novel metric theoretically explain the limitations of DT in modeling lo- cal dependencies, attributing attention disper- sion to inconsistent state-action-rewall distribu- tions; (2) We propose Token Merger to har- monize intra-step representations through ex- plicit modality alignment, demonstrating univer- sal effectiveness across different architectures for enhanced trajectory modeling; (3) We de- velop parameter-fiee free Token Mixer using win- dowed average pooling [34] to capture inter-step dependencies and preserve the Markov prop- erty, validated through theoretical analysis and experiments for efficiency and generalization; (4) We present generalized hierarchical frame- work that decomposes sequence modeling into intra-step modeling (Convolution-batt based Token Merger) and inter-step modeling (Pooling Token Mixer) operations, achieving 73.6% increase in speed and an 9.3% improvement in performance on the D4RL benchmark compared DT, as shown in Figure T (5) We demonstrate the generalization of our method to real-world robotic arm operations, achieving high success rate, smooth trajectory execution, and fast inference speed.
Performance Speed
95
90-
Performince
60-
75
Better
Model
Contect Langll (K)
Del
K-10
K-110
4500
5000
6000
6500
7000 x7
7500
Speed (step/s)
Figure 1: Performance and speed of the proposed Decision HiFormer (DHi) are compared with DT and the recently proposed DC on the MuJoCo Hopper-medium task. The results show that our method achieves significant improvements in both performance and speed.
Related Work
2.1
Offline RL
Offline reinforcement learning (offline RL) learns policies from fixed datasets without environment interaction, but suffers from distributional shift between the learned and behavior policies [19]. Q-learning method is one of the most prominent categories to address this by learning conservative
or constrained value functions [IO 16
17 28). For example, Batch-Constrained Q-learning (BCQ) [1] constrains the learned policy to remain near the dataset support by using generative model and action perturbation. Conservative Q-learning (CQL) [77 penalizes Q-values for actions not in the datasets, explicitly enforcing conservatism. Implicit Q-learning (IQL) [16] avoids behavior constraints by decoupling policy and value learning, using Dovaantage-wiighte regression. While effective, these methods rely value estimation and are sensitive to bootstrapping error.

---

2.2 Sequence Modeling for Offline RL
Recent advances in offline RL have shifted focus toward sequence modeling 6 [14, which reframe policy learning as predicting future actions from historical trajectories using architectures that capture temporal dependencies. pioneering work, DT [6], uses Transformers to model long-term dependencies by conditioning actions past states and future returns, eliminating the need for bootstrapping. DT's success inspires various follow-ups leveraging Transformers for temporal feature modeling. For example, EDT [T] introduces dynamic history length mechanism to address DT's trajectory stitching limitation, while Waypoint Transformer (2) focuses on implicit goal generation for trajectory planning. GDT 22 model complex state-action relationships. StARformer [25] uses interleaved Step and Sequence Transformers to explicitly separate short- and long-term dependencies.
combines Graph Neural Networks (GNN) with Transformers to
Despite their success, vanilla self-attention's global focus can overlook local patterns and incurs significant computational overhead. To alleviate these issues, alternative lightweight structures have been proposed. For instance, RvS [8] uses a two-layer MLP to condition on current state and return; DC [15) replaces self-attention with local convolution to enhance local modeling; and FCNet [27] adopts Fourier transforms to capture iow-frequency trajectory patterns for real-time decision maciii..-making.... State-space Models (SSMs), such as S4 [3] and Mamba 77 22) have also been adopted as lightweight alternatives to self-attention, offering enhanced local dependency modeling.
3 Method
3.1 Preliminaries: Trajectory Modeling and Attention Structure
In offline reinforcement learning (RL), the goal is to learn a policy that maximizes the expected cumulative return from a fixed dataset of trajectories, without any further environment interaction. Each trajectory r = (sg: ao: ro....; 81, or, rr.) collected under an unknown behavior policy, where 8t, dz and denote the state, action, and reward at timestep I, respectively. To enable autoregressive modeling, a common
caannncch-eemplfieee by DT represeeettt trajectories as sequences of return- to-go (RTG), state, and action tokens.
For each timestep f, subtrajectory of length is constructed in the form:
T- = (f 81- GI-K+ Re,
(1)
where the return-to-ga is defined as R, Σt - Tr. Each token in the subtrajectory is embedded into a d-dimensional vector and passed through a multi-layer Transformer encoder.
Given an input of length 3K, multi-head self-attention is applied. For each attention head h {1..... H}, attention weights computed using linearly projected queries and keys:
(ar W T Tax VTx
a softmax
where E Rd is the embedded token at position i, and W W E Rdxds are learnable projection matrices. These attention weights form the basis for downstream value aggregation and temporal dependency modeling.
3.2 Quantifying Attention Focus with Attention Entropy
Mathematical Formulation Although Transformers possess strong sequence modeling capabili- ties, their attention mechanism inherently position-agnostic and does favor temporally local structure. In offline RL, however, effective decision-making often relies on short-range temporal dependencies- e.g., transitions from or 2 due the Markovian nature of the environment. These short-range inter-step relationships. referred to here as local dependencies, play crucial role credit assignment and behavioral consistency across trajectories. To evaluate whether the model captures such local dependencies, we analyze the sharpness of its attention distributions. Specifically, we adopt attention entropy as a proxy measure. Let a denote the attention weight from token to token j in head h of the Transformer. These weights can be interpreted as probability distribution
3

---

(a) DT
Ag
86
Attention Entropy
3.5
Rewards
82
Reference Entrupy
3.0
84
2.5
90
2.13
2.0
1.5
(b) DT with Token Merger
Attention Entropy
Rewards
Reference Entropy
41
2.5
2,13
2.0
www.mm
1.5
Step
200
400
600
800
111
Figure 2: Entropy-reward correlation (left) and attention maps (right) for DT (top) and DT with Token Merger (bottom)) where lower entropy signifies more focused attention and thus stronger local dependency modeling. DT exhibits stateddominate attention, with sparse weights on RTG and action tokens, leading to higher attention entropy and a strong correlation between local entropy maxima and reward minima. In contrast, DT with Token Merger achieves significantly lower entropy, more balanced inter-step dependencies, and notably higher rewards.
over keys j for each query i, conditioned on the attention head h. The average attention entropy for head h is defined as:
h 0g
log atj.
(3)
N
1=1
where N = 3K is the total number of input tokens in the subtrajectory (with each timestep contribut- ing RTG, state, and action tokens). The overall attention entropy is obtained by averaging across all H heads: и Σi=1 Sn. lower value of S indicates more focused attention distributions, often associated with stronger local dependency modeling, while high entropy implies dispersed, possibly uninformative focus.
Empirical results We empirically evaluate this relationship in the MuJoCo Hopper- medium en- vironment by measuring attention entropy and its correlation with rewards from the first layer of DT after a warm-up of 50 steps. As shown Figure left), clear negative correlation emerges: spikes in entropy align with sharp reward drops, while periods of low entropy coincide with improved returns. Notably, low entropy does not always guarantee high reward reward-likely due to delayed credit assignment-butt high entropy consistently correlates with degraded performance. These patterns indicate that attention dispersion undermines stable decision behavior.
We further disentangle intra-step and inter-step behaviors by analyzing attention maps of DT from the perspective of attention entropy (Figure a, right)). Within a timestep, attention is highly concentrated on state tokens, with RTG and action tokens receiving minimal weight. We refer to the attention entropy within a single time step as intra-step modeling entropy. This modality imbalance indicates that DT's intra-step attention is focused but biased emphasizing state over other decision- relevant inputs. When extending this analysis across multiple timesteps, however, the attention scores for the dominant state become more dispersed across time steps, forming a full lower triangular matrix, aligning with the observations in [15) We treat the state, action, and RTG single time step as whole and refer to the attention score entropy across different time steps inter-step modeling entropy. The inter-step modeling entropy consistently higher, indicating that while DT exhibits structured intra-step attention behavior, it struggles capture inter-step dependencies. This

---

observation supports a key insight: DT's local modeling limitations stem not from weak intra-step fusion, but from its inability to propagate relevant information across adjacent timesteps. Note that in this context, local dependency modeling does not refer to intra-step modeling but the modeling between very adjacent timesteps.
Go-1
de
Forward
Model
Model
Feed Forward
Inter Inter-Step Modeling
Re-1
8--1 Re
8t
Re-1
Rg
Straightforward
Concatenation
Token Mixer
de-1
du
4--1
Pooling
Norm
Model
Model
Norm
Self-Attention
Self-Attention
Cony
Conv
Intra Intra-Start Stop Modeling
Token Merger:
R
5t-1
R
85
O--2
R
St
O4
R
8t
(c) Self-Attention
(d) Convolution
Figure 3: Hicrarchical modeling framework and micro structure of Token Merger. Token Merger merges < R, s, a > triplets into unified representations, while Token Mixer models dependencies
across timesteps.
3.3 Structured Entropy Reduction: Token Merger
To address the limitations in inter-step modeling observed in DT, natural idea would be to train the model via backpropagation capture dependencies across timesteps. However, since Transformers model pairwise interactions between individual tokens, lcarning relationships between coherent (s, a, H) triplets becomes challenging. The standard DT struggles with this due to the absence of timestep-level structure in its token sequence. Our proposed solution is to first merge the components of each (s, a, R) triplet into single composite token, and then model dependencies across these merged tokens. To this end, introduce two modules: the Token Merger for intra-step modeling and the Token Mixer for inter-step modeling.
The Token Merger merges the state, action, and RTG within each timestep into single unified representation, dynamically balaneing the contributions of each component to prevent Transformer from over-relying on state while neglecting action and return-to-go This design ensures that decision- making accounts for the dynamic relationship between state transitions and reward optimization. Additionally, in its original form, DT was responsible for both intra-step and inter-step modeling, but it struggled to balance these tasks. The Token Merger relieves this burden by providing timestep-level representations, allowing the Transformer focus exclusively that the total attention entropy remains unchanged after adding the Token Merger, we can apply Theorem B.I. to conclude that DT's overall attention entropy should decrease.
inter-step modeling. Assuming
Theorem 3.1. Let P (pi...., .P)) represent the original attention probability distribution. By = =Eje,
grouping tokens into triplets: G, = (pai-2, P.i--1,PP.... p' "(91;900 (991922....)
The entropy can be decomposed using the chain rule: H(P) = H(P") + Σ g. H(G)), H(G)) =
Inter-group
Intra-group
- EjEc,
Pi log Es Since H(G.) 2 0, we can conclude: H(P) 2 H(P)). 9:
Remark: This inequality guarantees that triplet grouping reduces attention entropy, enabling more focused attention across timesteps.
We further analyze the attention entropy and attention map of DT with a Token Merger in the same Hopper-medium environment discussed Sec.3.2. As shown in Figure b left), introducing Token Merger significantly reduces average attention entropy and yields a more stable and higher reward trajectory, supporting the conclusion of Theorem 3.1.. In terms of attention distribution (Figure b. right)), compared to the original DT model--wichh exhibits modality imbalance within
5

---

Table I: Performance in the Gym domain with different Token Mixer and Merger configurations. All models are trained for 20k steps, and scores are averaged over nine task-dataset pairs (combinations of Halfcheetah, Hopper, Walker2d tasks and medium- medium, medium-replay datasets) and five random seeds.
Merger
Straightforward
Concatenation
Attention
Convolution
Pooling
Average
Mixer
None (MLP only)
60.7
75.6
79.8
79.3
75.5
74.2
Linear Attention
74.6
78.1 1
78.8
80.4
80.0
78.4
Attention
73.5
76.9
76.4
78.9
75.4
76.2
Convolution
76.7
78.6
76.4
80.1
79.8
78.3
Pooling
75.0
78.5
78.5
80.7
78.1
78.2
Average
72.1
77.5
78.0
79.9
77.8
77.1
each timestep-the the merged attention pattern is more structured and focused across timestep-level tokens. This suggests that the Token Merger only compresses intra-step information but also promotes a more focused and temporally coherent inter-step attention profile.
3.4 Hierarchical Modeling Framework
This subsection outlines our hierarchical trajectory modeling framework. We first define the Token Merger for intra-step fusion and evaluate several merger operations, ultimately selecting convolution for its effectiveness. We then analyze the role of the Token Mixer in inter-step modeling, examining its necessity, architectural choices, and interaction with the Token Merger. Motivated by the short-range nature of dependencies in offline RL, demonstrate through both empirical results and theoretical analysis that a lightweight, non-parann con-aarmettric average pooling operation is an even preferable Token Mixer more complex alternatives. Together, these components form final model, shown in Figure 3
Token Merger. As mentioned, Token Merger maps each timestep's ' triplet (R, s, a) into a unified representation f R23dd + Rd. To ensure the alignment of Rt and a in the sequence, we obtain the input by performing an action shift to the trajectory +ii, restructuring as: 7, f
(Ot-K, K, Rr-K+1 K+1, 8t- K+1, a K+l...., Re, St) This adjustment ensures that the model predicts a, using (a,-1, R2) rather than (st, R2, az), reinforcing structured input format.
In standard sequence modeling architectures for offline RL- comprisins an embedding layer, encoder blocks, and prediction head the Token Merger is applied to trajectory sequences before they enter the encoder blocks (Figure J. We evaluate four types of single-step Token Mergers: Concatenation directly merges the state, action, and return vectors by combining them before applying an embedding layer; Belf-Attention introduces local attention mechanism over adjacent tokens to capture short- range dependencies [25]; Convolution applies kernel to extract local patterns from neighboring tokens; Pooling performs an average operation over adjacent tokens to aggregate their representations into single merged token. Experiments in the Gym domain show that performance is lowest when no merger is used, and highest with convolution (Table DT. Our experimental results are consistent with intuition that convolution performs the best because it effectively captures local dependencies between the tokens and smooths the information fusion.
Token Mixer. As mentioned, recent studies have explored various token mixers, such self- attention (6), linear attention
[27] and convolution [15), all which have demonstrated strong performance in offline RL tasks. This raises three key questions: (1) Is the Token Mixer essential for effective offline RL? (2) Do different Token Mixers significantly impact performance? and (3) Can the Token Merger mechanism generalize across different types of Token Mixers?
We evaluate various token mixers on nine Gym task-dataset pairs, as shown in Table I The experimental results indicate that: (1) The Token Mixer is essential, particularly in the absence of a Token Merger. Models without a Token Mixer (i.e., MLP only) perform significantly worse than those incorporating Token Mixers. However, with the Token Merger in place, the Residual MI.P (without Token Mixer) already performs reasonably well, similar to behavior cloning (BC) methods that cannot utilize historical trajectory information cumulative returns. (2) There is limited

---

performance variation between different token mixers. (3) Most Token Mixers achieve substantially better performance when combined with Token Merger, demonstrating that the effectiveness of the Token Merger extends across different Token Mixers.
Hierarchical Framework with Token Merger and Token Mixer. Based on these observations, we propose a hierarchical framework that integrates both a Token Merger and a Token Mixer. While the Token Merger alone, when combined with residual MI.,, can implicitly model temporal dependencies, it struggles with stability across timesteps due to the lack of context. Meanwhile, Transformer-style Token Mixers incur high computational costs and do not necessarily improve across- across-timestep dependency modeling, often overemphasizing long-range interactions that are less relevant in offline RL. Given that decision-making in offline RL is primarily driven by short-range dependencies, and that excessive modeling capacity can lead to instability or inefficiency, we propose non non-parametric average pooling operation the Token Mixer. This aggregates adjacent timestep representations, providing stable cross cooss-timestep smoothing without introducing additional parameters.
The convolution-bast Token Merger encodes each token as an (R, s, a) tuple, reducing the input sequence length to one-third of its original size and improving the computational efficiency of the multi-layer encoder. Given this reduction, pooling size of 2 is sufficient to capture temporal information from the current and preceding timesteps. We provide a theoretical proof of the pooling layer's expressive power in modeling Markov problems in Appendix A]and Appendix B We refer to our model as Decision HiFormer (DHi), a lightweight, hierarchical, and high-performance alternative to Transformer-basee architectures that demonstrates strong empirical and theoretical performance in modeling local dependencies.
4 Experiment
In this section, we present comprehensive evaluation of the model using the widely recognized D4RL benchmark [9]. A broad range of offline RL baselines are included. To bridge the gap between simulated benchmarks and real-world applications, Panda robotic arm for physical validation. Ablation studies are also conducted to examine the impact of different window sizes in the pooling layer and different Token Mergers model performance. We further demonstrate the outstanding performance of DHi in sparse reward environments as well as its generalization across context lengths.
further deploy DHi 7-DOF Franka Emika
4.1 Experiment Setting
Baselines We compare our approach with existing SOTA offline RL approaches, each excelling in
specific domain tasks. For value-based approaches, we include IQL [16]. BCQ [1], and CQL [17]
For supervised learning approaches, include DT (6) StAR 25 GDT 12 FCNet (27 and
DC (15). Further experimental details and more baseline comparisons are provided Appendix C
and Appendix D respectively.
D4RL Datasets We explore four domains in the D4RL benchmark, each presenting unique chal- lenges for evaluating offline RL. algorithms. The Gym-MuJoCo environments provide a straightfor- ward assessment platform with near-optimal trajectories and clear reward structures. contrast, the Adroit domain, based on human demonstrations, features a constrained state-action space where the agent must maintain reasonable behavior within strict limits. The Kitchen domain involves multi-task complexity, requiring agents to complete four sequential subtasks in a specific order and generalize beyond training data to reach the target configuration. Finally, the AntMaze domain uses an 8-DOF robotic agent with sparse rewards and increased locomotor complexity to make navigation even more difficult.
Worl-World Environment Datasets. We further evaluate the model's inference speed and gen- eralization ability in real-world settings using a 7-DOF Franka Emika Panda robotic arm equipped with a single-joint torque sensor. The experiments involve four tasks of varying complexity, as shown in Figure 4 In the Lift task, the robotic arm lifts green cube, while in the Place task, it picks up a teddy bear and places it into a green bowl. The Selective Lift and Selective Place tasks introduce additional challenges, with Selective Place incorporating obstacles to assess the model's
7

---

Table 2: Performance of DHi and baselines Gym, Adroit, Kitchen and AntMaze domains. All models are trained for 100k steps, and scores are averaged over five random seeds. Best and second- best scores denoted in bold and underlined, respectively.
Gym Tasks
CQL
IQL
BCQ
StAR
GDT
FCNet
Ours
malfchectta-madiim--eoo
91.6
86.7
69.6
55.2
86.8
93.7
93.2
91.2
93.0
94.2
hopper-medium-oxppe
105,4
91.5
109.1
52.5
107.6
111.1
111.1
110.5
110.4
111.7
108.8
109.6
67.3
107.5
108.1
109.0
107.7
108.8
109.6
109.6
halfcheetah--mae medimm
47.4
41.5
42.6
42.6
42.9
42.9
42.9
hopper-medium
69.4
66.3
65.1
52.9
67.6
59.5
77.6
57.8
92.5
90.1
walker2d-medium medium
83.0
78.3
52.0
75.3
74.0
73.8
76.5
75.2
79.2
79.9
45.5
44.2
34.8
36.6
36.6
36.8
40.5
39.8
41.3
41.5
hopper-medium--eplay
94.7
31.1
29.2
85.3
85.8
97.7
malker2d-medium--rphhe
77.2
73.9
13.7
32.3
79.4
39.8
77.5
63.5
76.6
81,2
Average
80.6
77.0
53.8
52.6
76.2
66,2
79.1
75.1
82.2
83.3
Adroit Tasks
CQL
IQL
BCQ
BC
StAR
GDT
FCNet
DC
Ours
pen-human
71.5
66.9
77.9
57.7
hammer-human
4.4
1.4
0.9
1,2
3.7
3.7
5.5
1,2
43.0
31.2
dnor-human
9.9
4.3
-0.1
2.0
14.8
1.5
20.6
0.4
22.6
25.2
pen-eloned
39.2
37.3
50.9
37.0
75.8
33.1
86.2
50.4
98.2
89.1
hammer-cloned
0.4
0.3
33.8
44.6
door-cloned
0.4
1.6
0.01
0.0
16.3
0.0
19.8
-0.2
23.6
23.6
Average
15.6
19.7
19.8
17.5
32.2
19.4
38.9
18.3
52.5
50.1
Kitchen Tasks
CQL
IQL
BCQ
DT
StAR
GDT
FCNet
Ours
kitchen-complete
43.8
62.5
8.1
65.0
50.8
40.8
43.8
28.0
50.0
55.0
kitchen-partial
49.8
46.3
18.9
33.8
57.9
12.3
73.3
32.5
75.0
75.0
Average
46.8
54.4
13.5
49.4
54.4
26.6
58.6
30.3
62.5
65.0
AntMaze Tasks
CQL
IQL
BCQ
BC
StAR
GDT
FCNet
DC
Ours
antmaze-umaze
87.5
78.9
54.6
51.3
85.0
86.9
antmaze-umaze-divrsee
84.0
62.2
55.0
45.6
53.0
45.6
69.0
82.0
78.5
84.0
Average
79.0
74.9
67.0
50.1
56.1
48.5
72.5
83.0
81.8
85.5
generalization capabilities. In Selective Lift, the model must identify red cube among various objects and place onto green plate.
To train the model, 50 image- trajectories were collected for the Lift and Place tasks (with the same set used for both Place and Selective Place tasks), and 60 trajectories were gathered for the Selective Lift task. The latter includes an equal split of trajectories with the red cube the left and the yellow cube on the right, adding complexity to the model's task execution.
Seletite
Lih
Plate
Sclective Place
Figure 4: Real-world applications.
4.2 Main Results
Results on D4RL Datasets We evaluate DHi
across four task domains from the D4RL benchmark: Gym, Adroit, Kitchen, and AntMaze, with results summarized Table To ensure fair comparison, we normalize the scores following the protocol established by [9. where a score of 100 represents expert-level performance. We report average scores from five random seeds. DHi demonstrates strong and consistent performance across all domains, surpassing baseline methods in many cases. In the Gym domain, DHi achieves the highest average score of 83.3, indicating its effectiveness in handling wide variety of tasks. In the Adroit domain, where offline RL struggles with extrapolation errors due to limited human demonstration data, DHi significantly outperforms most baselines and achieves highly competitive average score. For Kitchen tasks, which require generalization unseen states and long- long-term value optimization,

---

DHi delivers strong performance. indicating its high expressiveness, In the AntMaze domain, which presents challenges due to sparse rewards and sub-optimal trajectories, DHi consistently matches or surpasses existing methods. Overall, these results highlight the robustness of our approach across different domains and task complexities, reinforcing effectiveness in addressing the challenges posed by the D4RL benchmark.
Real- World Applications We further evaluate the performance of DHi and DT the real- world robotic dataset, measuring their success rates 10 trials and the inference time for generating 400 actions. As shown in Table DHi achieves 100% success rate across all tasks, whereas DT struggles with the Lift task and performs significantly worse in the more challenging Selective Lift task.
Table 3: Performance of real-world experiments. S-Lift and S-Place represent the Selective Lift and Selective Place tasks, respectively.
Policy
1int
5-Lint
Place
S-Place
Params/M
Time/s
DVT
0.90
0.40
1.00
1.00
9.7
2.3
DHI
1.00
1.00
1.00
1.7
Moreover, visualization results (Appendix G and videos in supplementary material) indicate that DHi produces significantly smoother robotic arm trajectories than DT and completes tasks more quickly. Under the same hyperparameter settings, DHi also demonstrates greater efficiency, requiring approximately 32% fewer parameters than DT (6.6M compared to 9.7M) and achieving a 35% improvement in inference speed (1.7s compared 2.3s). This efficiency gain is primarily due to replacing the self-attention layer with a non-parametric, constant-comllextt average pooling operation.
4.3 Ablation Study
The Design of Token Merger As discussed in the Method section, we evaluate the performance of different Token Mergers across various Token Mixers in the Gym domain (Table []. Among them, the convolution-based Token Merger consistently achieves the highest average performance. To further analyze effectiveness, we perform I-SNE visualizations on the same task using convolution- based token merging. Figure that before applying the Token Merger, embeddings exhibit random distribution of states, actions, and RTGs. After applying the Token Merger, the embeddings demonstrate improved local clustering, reinforcing its role in structuring input representations. Combined with the conclusion of Theorem B.T., these results suggest that the introduction of a Token Merger significantly improves the inter-step modeling consistency. This finding is further supported by evaluations under a more extensive training regime, as detailed Appendix E.
(a) without Token Merger
(b) with Token Merger
10
Stutes
Actions
RTGs
Merged Tokens
10
Figure 5: t-SNE visualization of embeddings with/without Token Merger. Post-merger embeddings
exhibit structured clustering, indicating enhanced local associations.
The Size Pooling Layer To analyze how pooling sizes influence the model's 's ability to capture token dependencies, experiment with pooling sizes of 1, 5, 10, 15 and 20 in the pooling layer. As shown in Appendix [2.2 the experimental results show that when the pooling size is set to (equivalent to MI.P only), the model's performance significantly decreases due the lack of contextual history. Smaller pooling sizes (2-10) result in better performance, while excessively large window sizes hinder the modeling of dependencies. We the pooling size 2 to align with the Markov property and ensure efficiency.
Generalization Across Context Lengths and Sparse Reward Settings We provide an in-depth discussion of our method's robustness to context length and its performance under sparse reward

---

settings in Appendix F These analyses further validate the design choices behind DHi, particularly its insensitivity to context size and its resilience in sparse reward conditions.
Conclusion
In this study, we address the challenges of Offline RL by proposing novel hierarchical decision modeling framework. Our approach effectively captures state-action-retumi relationships through intra-step and inter-step modeling. To quantify the model's ability to focus on relevant states, we introduce attention entropy as metric and reveal its negative correlation with performance. To tackle this issue, we propose the Token Merger to unify representations within each time step and design parameter-free Token Mixer based on Average Pooling for efficient inter-step dependency modeling. The resulting Decision HiFormer (DHi) outperforms existing baselines on the D4RL benchmark while enhancing computational efficiency. Furthermore, DHi has demonstrated strong generalization capabilities and practical benefits in real-world robotic manipulation tasks, underscoring its appli- cability beyond simulated environments. A limitation of this study is the lack of evaluation on the transferability of our approach to other RL paradigms, which we leave for future investigation.
6
Acknowledgments
This work was supported by the Shanghai Artificial Intelligence Laboratory and the National Natural Science Foundation of China (Grant Nos. 62572313 and 62106139).
References
Anurag Ajay. Yilun Du, Abhi Gupta, Joshua B. Tenenhaum, Tommi S. Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? Proc. of ICLR. OpenReview.net. 2023.
[2] Anirudhan Badrinath, Yannis Flet- Berliac, Allen Nic, Emma Brunskill. Waypoint transformer: Reinforcement learning via supervised learning with intermediate targets. Proc. of NeurlPS, 2023.
[3] Shmuull Bar-David, Itamar Zimerman, Eliya Nachmani, and Lior Wolf. Decision S4: efficient sequence- based RL state spaces layers. In Proc. of ICLR. OpenReview.net, 2023.
[4] Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Jasmine Hsu, et al. Rt-1: Robotics transformer for real-world control scale. ArXiv preprint, abs/2212.06811,, 2022.
[5] Yevgen Chebotar, Quan Vuong, Karol Hausman, Fei Xia, Yao Lu, Alex Irpan, Aviral Kumar, Tianhe Yu, Alexander Herzog, Karl Pertsch, al. Q-transformer: Scalable offline reinforcement learning via autoregressive fuunctiooss. Proc. of CoRI. pages 3909-3928. PMLR, 2023.
[6] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In Proc. of NeurlPS, pages 55884-1009) 2021.
[7] Yang Dai, Oubo Ma, Longfei Zhang, Xingxing Liang, Shengchao Hu, Mengzhu Wang, Shouling Ji, Jincai Huang, and Li Shen. Is mamba compatible with trajectory optimization in offline reinforcement lcarning? In Proc. of NeurlPS, 2024.
[8] Scott Emmons. Benjamin Eysenbach. Ilya Kostrikov. and Sergey Levine. Rys: What is essential for offline RL via supervised learning? In Proc. of ICLR. OpenReview.net, 2022.
[9] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RI:: datasets for deep data-driven reinforcement learning. ArXiv preprint, abx/2004.07219, 2020.
[10] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning, In Proc. of NewrlPS, pages 20132 20145, 2021.
[11] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without explo- ration. In Proc. ICML, pages 2052-2062. PMLR. 2019.
[12] Shengehao Hu, Li Shen, Ya Zhang, and Dacheng Tao. Graph decision transformer. ArXiv preprint. abs/2303.037477 2023.
10

---

[13] Shengchao Hu, Ziqing Fan, Chaoqin Huang, Li Shen, Ya Zhang, Yanfeng Wang, and Dacheng Tao, Q-value regularized transformer for offline reinforcement learning. In Proc. of ICML. OpenReview.nee, 2024.
[14] Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. In Proc. of NeurlPS, pages 1273-1286 2021.
[15] Jeonghye Kim, Suyoung Loe, Woojun Kim, and Youngchul Sung. Decision convformer: Local filtering in metaformer is sufficient for decision making. In Proc. of ICLR. OpenReview.net, 2024.
[16] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforeement learning with implicit q-leamiig. In Proc. of ICLR. OpenReview.ned 2022.
[17] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In Proc. of NeurlPS, 2020.
[18] Kuang-Huer Lee, Ofir Nachum. Mengjiao Yang. Lisa Lee. Daniel Freeman, Sergio Guadarrama, Ian Fischer, Winnie Xu, Eric Jang, Henryk Michalewski, and Igor Mordatch. Multi-game decision transformers. In Proc. of NewrIPS, 2022.
[19] Sergey Levine, Aviral Kumar, G. Tucker, and Justin Fu. Offfine reinforcement learning: Tutorial, review and perspectives on open problems. ArXiv preprint, abs/2005.01 643,, 2020.
[20] Shuang Li, Xavier Puig, Chris Paxton, Yilun Du, Clinton Wang, Linxi Fan, Tao Chen, De-An Huang, Ekin Akyürek, Anima Anandkumar, Jacob Andreas, Igor Mordatch, Antonio Torralba, and Yuke Zhu. Pre-trained language models for interactive decision-making, In Proc. of NewrlPS, 2022.
[21] Zuxin Liu, Zijian Guo, Yihang Yao. Zhepeng Cen, Wenhao Yu. Tingnan Zhang, and Ding Zhao, Constrained decision transformer for offiine safe reinforcement learning. In Proc. of ICMI, pages 2161 1-21630. PMLR, 2023.
[22] Qi Lvv Xiang Deng, Gongwei Chen, Michael Yu Wang, and I Liqiang Nie. Decision mamba: A multi-grained state space model with self-evolution regularization for offiin RL. In Proc. of NewriPS, 2024.
[23] Ashok Vardhan Makkuva, Marco Bondaschi, Alliot Nagle, Adway Girish, Hyeji Kim, Martin Jaggi, and Michael Castpar. Attention with markov: A curious case of single-layer transformers. In Proc. of ICLR. 2025.
[24] Ethan Perez, Florian Strub, Harm de Vries, Vincent Dumoulin, and Aaron C. Courville. Film: Visual reasoning with general conditioning layer. In Proc. of AAAI, pages 3942-3951. AAAI Press, 2018.
[25] Jinghuan Shang, Kumara Kahatapitiya, Xiang Li, and Michael S Ryoo. Starformer: Transformer with state-action-reward representations for visual reinforcement learning. In Proc. of ECCV, pages 462-479.. Springer, 2022.
[26] Mohit Shridhar, Lucas Manuelli, and Dieter Fox. Perceiver-actor: A multi-task transformer for rohotic manipulation. In Proc. of CoRL, pages 785-799. PMLR, 2023.
[27] Hengkai Tan, Songming Liu, Kai Ma. Chengyang Ying, Xingxing Zhang, Hang Su, and Jun Zhu. Fourier controller networks for real- time decision-making in embodied learning. In Proc. ICML OpenReview.net. 2024.
[28] Faraz Torabi, Garrett Warnell, and Peter Stone. Behavioral cloning from observation. In Proc. of IJCAI, pages 4950-4957. ijeai.org, 2018.
[29] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakoh Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Prov. of NeurlPS, pages 5998-6008, 2017.
[30] Muning Wen, Jakub Grudzien Kuba, Runji Lin, Weinan Zhang, Ying Wen, Jun Wang, and Yaodong Yang. Multi-agent reinforcement learning is a sequence modeling problem. In Proc. of NeurlPS, 2022.
[31] Yueh- Wu, Xiaolong Wang, and Masashi Hamaya. Elastic decision transformer. In Proc. of NeurlPS. 2023.
[32] Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, Joshua B. Tenenhaum, and Chuang Gan. Prompting decision transformer for few-shot policy generalization. In Proc. of ICMI,, pages 24631 1-44455. PMI.R, 2022.
[33] Taku Yamagata, Ahmed Khalil, and Raúl Santos-Rodriguez. Q-learning decision transformer: Leveraging dynamic programming for conditional sequence modelling offline RL. Proc. of ICML, pages 98989-39007. PMLR, 2023.
[34] Weihao Yu, Mi Luo, Pan Zhou, Chenyang Si, Yichen Zhou, Xinchao Wang. Jiashi Feng, and Shuicheng Yan. Metaformer actually what you need for vision. In Proc. of CVPR, pages 10819 10829, 2022.
11

---

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

---

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

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [No]
Justification: Once the blind review period is finished, we'll open-source all codes and instructions.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
ollicy for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips.ccppblle pwwppppguuuiieeppepppppppppupppp Sute:////uudde//eeSuuuubiii
(CodeSubiisss wwtplubmmsssioiio ssiiiioicccy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Yes, please Appendix C
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Yes, please see Appendix Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi-
dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for
example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
14

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: please see Table 3 and Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer:
[Yes]
Justification:
We followed the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- cration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: The paper focuses on technical contributions without direct societal implica- tions, and no explicit positive or negative societal impacts are discussed.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
15

---

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

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
17

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: This work does not involve the use of LLMs in the core methods.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.LM policy battap///eerripoc.cc///oooereeep for what should or should not be described.
18

---

A Pooling with Markov
Given the inherent complexity of deep neural networks, our work builds upon the mathematical framework proposed by [23], focusing on simplified scenarios involving first-order binary Markov data and single-layer pooling token mixer. We conduct principled analysis of pooling token mixer using Markov chains and further extend the framework to accommodate multi-state Markov processes. Our notation is summarized in Table ST
Table S1: Mathematical Notations and Definitions
Notation
Definition
Notation
Definition
Iz V
Scalars
Vectors
M.P.W. W
Matrices
II
f2-norm Frobenius norm
[K]
Set [l...., K}
Sequence from I. to.2,
a(2)
Sigmoid function
ReLU(2)
ReL.U function: max(0,
((AA)
Probability of event A
P(A B)
Probability given
X
Vocabulary, e.g., {0,
N
Length of input sequence
{In) )
Input sequence
S
Vocabulary size (multi-state)
Pij
Transition probability P(
Transition kernel matrix (Piy)
I 1)
π (!))
Probability distribution of state
n
(1)
Initial state distribution
P.g
Switching probabilities Pot
g(p.g)
Binary Markov kernel
P
p, P10 =
H(Y [X)
Conditional entropy
h/py)
Binary entropy function
H (( )..)
Markov chain entropy
Da. Q)
Kullback-Leibllr divergence
Data. We assume binary vocabulary X' = {0, 1} and model the input sequence {I, J first-order Markov chain (i.e., with memory order m 1). In such processes, the probability of the next state depends solely on the current state, independent of any preceding states:
P A P(rnt1 i x, 1) P( = j Ti, i)
For any ú,j E X, we consider Markov process governed by the transition kernel P = (Pt)), which determines the system's 's dynamics. Specifically, if n [0, 1||x represents the probability distribution of the state at time then the distribution evolves according to π (n+1) = t (n) P.
In this paper, we focus specific transition kernel defined over the binary state space:
P(p A
p
p
1 1-
where p, 9 E (0, 1) denote the switching probabilities (i.e., Poi p and Pio ②). We define
first-order binary Markov chain (#n)n21 characterized by the transition kernel P(p, q) and an initial distribution (1), Formally, we denote this process as (In)n21 ~ (I(1), (1), p)).))) When the initial distribution x (1) is clear from the context, we adopt the simplified notation (Ia+1 Fn)n21 ~ P(p.)) describe the conditional transitions.
The entropy rate of a first-order Markov chain, defined as H (#n+1 x,), quantifies the average uncertainty of the next state given the current state. For the binary Markov process considered in this work, the entropy rate given by
gh(p) ph(())
H(r++1 x.,))
p ++
where ((-) denotes the binary entropy function.
19

---

Model. To study concrete instance of a Pooling Token Mixer, we modify the model architecture from 233. The Celf-Attention layer in this base model is substituted with module comprising an average pooling operation of size k followed by a single-layer LayerNorm. Let ( J E {0, 1 denote an input sequence of length N. For each n. E [M], the operations within this model are defined as follows:
X = e + Pa E R
(Embedding)
9.
Σ X; E ERR
(Pooling)
knn
q. = O 9n
fp (p))
+ B1 E g
(LayerNorm)
VIo(qn) +
yn Xn + +Qn
(Shortcut)
In W W2ReLU(WiNya) ER"
(Feed-Forward)
2n = yn + tn E RR
(Shortcut)
Sn (a, zn) + b
(Linear)
fo(x)) = a (s,)) ).)c) E [0.
(Prediction)
In the above equations, the embedding dimension. The term e E Rd is the embedding vector, where specific vectors (for = 1) and en (for I,, 0) are utilized. P.. TR represents trainable positional encoding. For the LayerNorm operation applied after pooling, 7,. E RR are learnable parameters used for scaling and shifting, respectively, > 0 is a small constant for numerical stability. The functions f.),,, (-): : Rd + IR compute the mean and standard deviation, respectively, of their input features (here, the pooled features qn). The feedforward network (FFN) consists matrices W1 E IRrxd and W₂ E Rddr where the hidden dimension the FFN. The final linear layer is characterized by the weight vector E Rd and scalar bias b R, which produce the logits 8n- The probability for the symbol 1 is then computed from these logits using the sigmoid function o())
fa(xT) Po(xn+1 = 1 1 = (sx))
Note that single symbol probability suffices, the vocabulary is binary. The set of all trainable parameters in our model is given by
0 e (e, (p. 31 W1, a2,
Loss. We train the parameters 0 using a next-token prediction loss, defined as the cross-entropy loss between the predicted probability (=f") and the ground truth token xn+1 at each position [N]. Specifically, the loss function is given by
L
Σ NE[N]
Exaei In+1 log fe + (1 xn+1) log (1 - (o-fo(x))))
(I)
N
where the expectation taken over the data distribution of the sequence {an \ practice, this expectation is approximated by the empirical average computed over sequences sampled from the corpus, and stochastic optimizers such as SGD or Adam are used update the model parameters.
Objective. Building on the theoretical framework described above, we investigate the modeling capacity of the aforemeentooed single-layer pooling-based models (i.c., those utilizing a pooling token mixer). We address the following research question:
Can single-layer models utilizing a pooling token mixer effectively capture the dynamics of data sequences that exhibit Markovian properties?
By analyzing the global minimum of the loss function, we derive the theoretical conditions under which these single-layer pooling-base based models, when equipped with appropriately designed param- eters, can provably model the Markovian dynamics present in input streams with state transition dependencies.
20

---

Lemma 1. (Loss as KL Divergence) Let the input sequence {In ( be generated by Markov chain (T. q). P)) with fixed (p. q) (0,1)². Let 0 (en, c1; (paIn=::.... denote the model parameters. The cross-entropy loss I, as defined in Eg. ( can then be expressed as the sum of the KI. divergence between the Markov kernel and the predicted distribution, and the entropy rate of the Markov chain:
E RD
I(Ø) N
Σ EuT REIN]
Euy
DKL ( aa) Po(- =?))] H (xn+1 In);
where DRL (P.) I1 Q) denotes the Kullback-Leibler divergence between two distributions P and Q, and H (In+1 Fn) is the entropy rate of the Markov chain. The detailed proof is available in (23).
Lemma 2. (Gradient Computation) Retaining the setup from Lemma 1. let L(0) be the cross- entropy loss defined in Eq. (T). Then, for any parameter w & 0, the gradient V VL is expressed
as:
V L
1
Σ
[d E Patt1
-o v
N
DE[N]
Σ
Eay
(P(Tn+1=
N
ME[M]
The second equality is derived using the law of total expectation. The detailed proof is available in 123.
Theorem 1. (Global Minimum in the Weight-tied Case) Let the input sequence { J.V be gen- erated by Markov chain (T q): P(p, q)) with fixed (p, q) E (0, 1)2. For the weight-tied case, model parameters. Then, for any (p, q), there exists an
let 0 E D-d denote the complete explicitly constructed parameter 0* RD-d d such that 0 is global minimum the population loss I.()) defined in Eq. ① and its predictions perfectly match the Markov kernel. More precisely, we have:
(i) L(0)= 2 ((0")) all RD-4
(ii) Fe. (In+1 | P( = 1 xn))
(iii) L(0") H (2n+1 xn).
(iv) ( = 0, ie., 0* is stationary point.
The detailed proof is available in Appendices B. and B.2
Corollary 1. (Global Minimum in the Untied Case) The weight-fied solution 0- D-D from Theorem L. can be extended to untied parameter vector 8" & (0' a²) E R by introducing independent copy a'. This construction ensures that the prediction probability under O" matches that under 0* and thus the true Markov kernel:
Pø (In+1 = x²) = Po- (In+1 | = P(Int1 = xxn).
Since this prediction matching (which is property (ûi of Theorem I. adapted for 0*) is a sufficient condition for global optimality (as established in the proof of Theorem I.. e.g., Appendices B. and B.2. it follows that 0" is global minimum of the loss L(-) in R.P Consequently, O also satisfies properties analogous (i), (iii), and (iv) of Theorem I. with 0* replacing 0* and the parameter space being R
Multi-state Markov Chains. Our framework, originally formulated for the binary setting X = {0, 1). naturally generalizes to the multi-state setting X {0, 5,...,S- 1}. where S denotes the state (vocabulary) size. Let the input sequence (xn)n21 IP(p) be first-order Markov chain such that, for some E (0, 1), either 1,+1 = In with probability - p or #n+1 switches uniformly at random to different state with probability - The symmetric kernel P(p) described earlier (with y y in the binary case) generalizes naturally to this multi-state setting. For the model architecture, we use the same design as before, except that the embedding and linear layers now incorporate token embeddings.
21

---

B Proof of Theorem 1
B.1 Proof of Theorem 1 for q S 1, Weight-tied Case
Proof. Assume that + q S 1 and that we employ weight tying, i... the parameter set is given by
0
a, {p=) .... RRD-d.
According to Lemmas| T. and 2. if a set of parameters 0 configures our model such that its output fo(x))) perfectly matches the true conditional probability P( 1 vn), then 0 achieves global minimum for the loss function L(.). At this minimum, the loss equals the entropy rate of the process, and 0 is also stationary point.
Consequently, sufficient construct such a parameter set, which we denote 0", To inform the design of 0", we note that the explicit form of the target Markov kernel is:
P( xn.) Fn (1 p 9) + p.
Therefore, to enforce
it is sufficient for the model fo(=1) = In (1-p-91 pp
only the information from the current symbol 26, and ignore the preceding tokens z²-1_ According to the architecture described in Appendix [ natural way to
achieve this to set W1 = and 7 B the positional embeddings p.. 0 for all n, have
the LayerNorm module. Moreover, by choosing
X,, = on e, Yn = Xns
Zn
Xn.
The logits are then computed as
Sn (e, kn) + +b xn /le/ b. Since fo(=?) = a (8n) must equal the Markov kernel, we require that o (IIAlee]
In (1 p q) p. Taking the inverse of the sigmoid (i.e., the logit function) yields
+b) [
Zin (1 p 9) P
In eml|e|| + log
(1 p) xnn) + y Fin
By substituting and Vn 1, we obtain:
p
Fy 0: b= log
p
Fn. |||| + b log
- 9
q
Subtracting the first equation from the second, we obtain
|||| ==l log (1 p1 9)
pq
Thus, a global minimum must satisfy
||||2 = log
(1 p1 - 9)
and b = log
P
pq
1- p
Note that this choice is well defined since
(1pi((1- (1)))))
1 when p q implying
Pq
log (1-p)(1- O.
Although there are infinitely many solutions (c, Pa; b) satisfying the above conditions, canonical choice for the global minimum 0 = 0" is
(1 - p1 9)
e = a=
log
{p. = 0
0"
pq
P
71 = 3 = 0, = log
W = 0,
W 0
1--p)
where DRd denotes the all-one vector. This completes the explicit construction of 0* and hence the proof.
22

---

B.2 Proof of Theorem 1 for 1, Weight-tied Case
Proof. We follow similar strategy as in the p p+g g I case by constructing parameter RRnd such that
fo P(zn+1
zn) = In (1 p p-
q) p.
However, in this case we must utilize the ReLU component of the feed-forwal mechanism, unlike earlier case where we set W₂ = O. We now describe the construction of 0".
First, we set the embedding vectors as e = a 1, and the positional encodings as
Pn - 1, Vn 1,
where 1 E IRd denotes the all-one vector. Consequently, the embedded input becomes
Xn = In e + Pa Oh 1,
with
+
if Fn 1,
On
1 if Fn 0.
We also set the LayerNorm parameters to zero:
Y1 = 31 = 0.
For the feed-forward layer, we choose W1 and W2 so that
W2 ReI.U(W. B,
where 8n is scalar that may depend on Zn. Consequently, the output of the feed-forward block becomes
Z, x,, + W2 ReLU(W.
1 + Bn 1 (an + BBn) 1.
The logits are then computed as
Sn (a, 2n) + d (an +B,) b,
since a = 1 and (1.1) = d. As fo(z)) = (sn)) must equal the Markov kernel, we require
o (d(an, + 8n) +) In (1 p- 9) +p, In E {0,1). 13.
Taking the logit (i.e., the inverse of σ) yields
In - (1-p-
(a, + Bx) + log
In E 10.
(1 (1 xn.) + g
By substituting Fn 0 and nn separately and denote the corresponding values of B by 8o and
31. respectively, we obtain:
In 0: d(- ( (6) + log
p
p
1
In 1: d(} + 31) + = log
Subtracting the first equation from the second yields
a(1 + 31 -8) log
(1 9)
(2)
pq
It remains to choose 3n and 31 so that above equation holds. To this end, consider the design of the feed-forward layer. Let W, 11 and W2 -2 WT for some R2o Since ya 0₂ we have
WIyn 11' (a, 1) aa d1.
Thus,
0,
if at = - (i.e. In = 0),
ReLU(W: (.).)
Ød1, 2
if an 1 (i.e. In 1).
23

---

Then.
0,
if 2n 0,
We ReLU = ? ReLU BeLU(W1y.) --w² w2 df 1,
if 1.
Hence, we obtain 80 = and B1 -w² d. Substituting these into Eq. 2 gives
d w² dP) log ((1 p) 9)
pq
That is,
1 d
(1 p) ②)
1 1 w² de
log
pg
Let w = w be the solution:
1
(1 p1 9)
log
ds
d
pq
Next, substitute Bo = O in the equation for 0 to obtain the bias:
d( (-2 + b log
1
p
p
+ d
= log
I-p p-p/
1- p
Piecing everything together, we define the canonical global minimizer 0* as
e = a a
{Pm = - 1 71 31 = 0,
0"
d W1= W2 = - W.
b = log {- + 2'
This completes the explicit construction of 0* and the proof.
C Experimental Setup
C.1 D4RL
In this section, we detail the experimental setup for our Decision HiFormer (DHi) model on the D4RL benchmark (9]. We implement DHi by building upon the official Decision ConvFormer (DC) [5] codebase and incorporate our proposed pooling module. This module performs a lightweight temporal averaging operation: for an input sequence (e.g., of merged tokens) [x1, Ø2.-.. ... NNN, it
=N-1+>x] J. This specific form of pooling, with an effective window of two for the averaged elements, aligns with our goal of efficiently capturing
produces an output sequence [x], Z1 22izza
2
local temporal dependencies from current and preceding timesteps, as discussed for DHi.
Our experiments are conducted on several offline RL tasks from the D4RL benchmark (9) which span diverse challenges including continuous control, dexterous manipulation, and nanghhorizon navigation with sparse rewards. Specifically, we select representative tasks from the following domains:
This domain includes classic continuous control tasks such as HalfCheetah, Hopper, and Walker2d. These tasks generally feature dense reward signals and are used to evaluate the model's ability to learn fundamental motor skills from offline data. The datasets used are medium (generated by a mecdium-performann from the replay buffer during the training of medium policy, containing large volume of suboptimal states), and medium- medium-expert expert (a mixture of trajectories generated by medium and expert policies, providing both generalization challenges and high-quality demonstrations).
Gym-MuJoCo:
policy), medium-replay (collected
This domain consists of high- -inensninnsinn,,, fine- fine-graine graiind dexterous manipulation tasks such as Door, Hammer, Pen, and Relocate. These tasks originate from (potentially noisy) human demonstration data and present significant challenges due to action complexity and system dynamics. The datasets used include human (trajectories from real human operations, which are natural but may be unstable), cloned (data generated by a low- quality policy obtained through behavioral cloning from the 'human' data), and expert (data generated from expert policies).
Adroit:
24

---

AntMaze: This domain encompasses navigation tasks of varying scales, such as 'umazee. 'medium', and 'large" mazes. The core challenges lie in sparse rewards and long-term temporal dependencies, requiring the agent to perform long-term credit assignment and effective exploration. The datasets used are play noo--gol--oriented human teleoperation data, offering broad state coverage but low policy efficiency) and diverse (data collected from multiple policies, enhancing state coverage while increasing the learning difficulty).
Kitchen: This domain involves multi-task scenarios where the agent must complete sequence of four subtasks in a specific order to reach a target configuration, testing general- ization and sequential task completion. We use the complete dataset (trajectories showing successful, ordered completion of all subtasks) and the partial dataset (trajectories with incomplete or out-of-order subtask attempts, representing suboptimal data).
The hyperparameters for Decision HiFormer (DHi) in our D4RL experiments. detailed in Table S2 guide the model's architecture and training. DHi is trained for 10 steps with a batch size of 64. Architecturally. it consists of layers with GELU activation. The context length K, defining the length of historical dimension is optimized from {128, 256} based on preliminary experiments. For training, we use the Adam W optimizer with Cosine learning rate scheduler, selecting an initial learning rate from {10-A, 10 10-33. Regularization includes a weight decay of 10-4 and gradient clipping at 0.25 to ensure stable training.
sequences fed into the model, is set to 20. The hidden
Upon completion of training, the model is evaluated online in the respective D4RL task environments. We adopt the standard evaluation protocol from D4RL [9] to quantify performance in these test environments. Specifically, we report the normalized score, calculated as:
score random score
normalized score 100 x expert score random score
Here, score refers to the agent's raw score, random score is the average score achieved by a random policy, and expert score is the average score achieved by an expert policy for that task. This normal- ized score typically ranges from 0 to 100, where O corresponds to random policy performance and 100 corresponds to expert policy performance. For each task, we report the average of normalized scores over five random seeds, where each seed is evaluated over 100 episodes.
Table S2: Hyperparameter Settings of DHi in D4RL Dataset.
Hyperparameter
Value
Total number of updates
105
Number of layers
3
Context length K
20
Dropout
0.1
Batch size
64
Optimizer
AdamW
Learning rate scheduler
Cosine Scheduler
Learning rate
110-4,10-3)
Weight decay
10-4
Gradient norm clip
0.25
Nonlinearity function
GELU
We compare our approach with existing SOTA offline RL approaches, each excelling in specific domain tasks. For value-based approaches. we include Implicit Q-learning (IQL) [6). Batch- Constrained deep Q-learning (BCQ) [. and Conservative Q-Learning (CQL) [17] For supervised learning approaches, our comparisons encompass Behavioral Cloning (BC), Decision Transformer (DT) (6). StARformer (StAR) 255 Graph Decision Transformer (GDT) 112, Fourier Controller
Networks (FCNet) [27] and Decision ConvFormer (DC) [5]. For FCNet and DC, we report the results from the original papers, while results for the other baselines are cited from [3]
Table S3 presents the computation time for one training epoch, GPU memory usage, and the number of parameters. These metrics offer comparative analysis of the computational efficiency between
25

---

Table S3: Comparison of Model Complexity and Training Resources on D4RL.
Complexity
MLP
DT
DC
DHi (Ours)
Training time (s)
30.4
34.8
34.2
31.9
GPU memory usage
1.82 GiB
2.10
2.00 GiB
1.84 GiB
Total parameters
1.84 M
M
1.86 M
1.84
Token Mixer parameters
0.0
789.5K
13.8
0.0
MLP vs DT vs. DC DHi (Ours), all of which are trained on single RTX 4090 GPU. It is observed that DHi is more efficient than DT and DC terms of training time, GPU memory usage, and the number of parameters.
C.2 Real- World Applications
We further evaluate DHi's inference speed and generalization ability real-world settings using DOF Franka Emika Panda robotic arm equipped with a single-joint torque sensor. The model's input consists of RGB images captured by cameras at three positions (right, left, and wrist), as well as dimensional joint position data. The input images are encoded using ResNer50 for visual information, and the joint position features are fused with the visual features through FiL. (2).. These fused features are then passed through a linear layer to transform them into 51 512-dimensional state vector, which input into DHi. The final output is a 10-dimensional action. The hyperparameters of Decision HiFormer and the environmen used for real-world evaluation are summarized in Table [S4] and Table S5
respectively.
Table S4: Hyperparameter Settings.
Table S5: Latency Measurement Environment.
Hyperparameter
Value
Configuration
Value
Total number of updates
2 x 105
Warm Up
10 dummy state inputs
Number of layers
3
CPU
Intel Xeon Platinum 8383C
Hidden dimension
512
CPU Clock
2.70 GHz
Context length
16
Memory Clock
3200 MHz
Batch size
8
OS
Ubuntu 20.04 (Kernel 5.15)
Optimizer
AdamW
Python
3.10.16
Learning rate scheduler
MultiStepLR
GCC
9,4.0
Learning rate
10-1
Torch Version
2.1.1
Nonlinearity function
GELU
D
Supplementary Experiments
D.1 More Comparison
For extensive comparison, compare DHi with more baselines: TD3+BC 100, model-free actor-critic method, utilizing behavior cloning for regularization; RvS 8 a neward-conditional
behavior cloning approach framed supervised learning; Decision Diffuser (DD) [I]. diffusion- based generative model for iterative trajectory refinement; Waypoint Transformer (WT) [2].
transformer-based sequence-modeling approach using automatically generated subgoals; Decision Mamba (DM) [22], an SSM-based decision-making model with multi-grained architecture. The results shown in Table S6
D.2 Further Validation
To assess whether the experimental patterns observed earlier persist under an extended training regime, we conduct this experiment exclusively on the Hopper-medium dataset, increasing the total number parameter updates from 10 to x 105, performing evaluations every 10 steps. All other experimental settings remain unchanged.
26

---

Table S6: Performance of DHi and more baselines Gym domains, where dataset abbreviations 'medium' as 'm', 'medium-replay' as 'm-r', and 'medium-expert' as 'm-e'. The results are all cited from their original papers.
Task Name
TD3+BC
RvS
DD
WT
DM
Ours
halfcheetah-m-e
97.9
92.2
90.6
93.2
93.5
94.2110
hopper-m-e
112.2
101.7
111.8
110.9
111.9
111.742..6
walker2d-m-e
101.1
106.0
108.8
109.6
111.6
109.611..0
halfcheetah-m
42.8
41.6
49.1
43.0
43.8
13.4++100
hopper-m
99.5
60.2
79.3
63.1
98.5
90.1 2.99
walker2d-m
79.7
71.7
82.5
74.8
80.3
79.912..7
halfchcetah-m--
43.3
38.0
39.3
39.7
40.8
41.5+
hopper-m-r
31.4
73.5
100.0
88.9
89.1
97.7+
walker2d-m-r
25.2
60.6
75.0
67.9
79.3
81.2++8.3
Average
70.3
71.7
81.8
76.8
83.2
83.3
As shown in Table S7, the results obtained by extending the training budget largely corroborate the conclusions drawn from our experiments with a smaller training budget. Specifically, we reconfirm that: (1) the Token Mixer is essential for achieving higher performance (average scores of 81.3-86.2 with mixers vs. 76.6 for MLP-only), though an MI.P-only model can perform reasonably well when combined with effective Token Merger (e.g., 81.2 with the Convolution merger); (2) the performance variation among different Token Mixers remains limited (average scores ranging from 81.3 to 86.2); and (3) Token Mergers consistently improve performance across different Token Mixers and the MLP-only baseline, highlighting their general utility. These findings under constrained training regime further validate the importance of the Token Mixer component and the synergistic benefits derived from employing Token Mergers.
Table S7: Performance in Hopper-medium task. report the mean normalized scores obtained
across five random seeds.
Merger
Straightforward
Concatenation
Attention
Convolution
Pooling
Average
Mixer
None (MLP only)
64.0
76.5
79.4
81.2
82.1
76.6
Linear Attention
76.9
89.5
79.1
94.8
90.8
86.2
Attention
67.6
78.9
85.6
85.6
81.3
Convolution
92.5
89.2
64.4
88.2
83.9
Pooling
75.6
84.3
84.7
90.1
87.1
84.4
Average
75.3
83.7
78.6
88.0
86.8
82.5
D.3 Standard Deviation
Due to limited space in Ta the main text, the standard deviations for the Gym domain reported in Table S6 while those for the remaining tasks are shown in Table
Table S8: Standard deviations of DHi across different domains.
Adroit Domain
Kitchen Domain
Task
DHi
Task
DHi
pen-human
86.6 = 11.8
kitchen-complete
55.0 4.5 75.0 5.3
hammer-human
31.2 5.6
kitchen-partial
door-human
25.2 = 1.9
pen-cloned
89.1 11.1
AntMaze Domain
hammer-cloned
44.6
antmaze-umaze antmaze-umaze-diveree
86.9 11.9
door-cloned
23.6 2.1
84.0 + 13.6
27

---

E
Ablation Study
E.1
The Design of Token Merger
The Token Merger fuses the per-timestep return-to-go (RTG), state, and action into a unified rep- resentation, which enables the subsequent encoder layers to focus modeling inter-timestep de- pendencies. To investigate the impact of different Token Merger, we compare several Token Merger variants-AAttentionn Pooling, Concatenation, and Convolution-against a Straightforward baseline that directly feeds separate RTG, state, and action tokens into the model without explicit merging. The comparative results in the Gym domain are summarized in Table
Experimental results show that all Token Merger variants outperform the Straightforward baseline, indicating the importance of explicitly integrating RTG, state, and action. Among the different approaches, the Convolution based Token Merger consistently achieves the best average performance, it effectively captures local dependencies between the input tokens and smooths the information fusion. In contrast, the Attention-based variant, known for its strength in global modeling, shows poor performance this task.
Table S9; Performance of DHi with various Token Merger in Gym domain, where dataset abbrevia- tions are: 'medium' as 'm', 'medium-replay" as 'm-r', and 'medium-expert' as 'm-e'.
Task Name
Straightforward
Attention
Pooling
Concatenation
Convolution
halfcheetah-m--e
89.3
92.7
88.9
90.8
94.2
hopper-m-e
108.9
111.7
110.8
111.0
111.7
walker2d-m-e
108.4
96.2
108.0
107.3
109.6
halfcheetah-m
43.1
43.0
43.1
43.1
43,4
hopper-m
75.6
84.7
87.1
84.3
90.1
walker2d-m
80.1
78.8
79.5
79.0
79.9
halfchectah-m-rr
38.8
38.8
36.0
38.0
41.5
hopper-m-r
83.4
92.6
95.1
97.1
97.7
walker2d-m-r
71.3
74.4
75.4
76.4
81.2
Average
77.8
79.2
80.4
80.8
83.3
E.2 Analysis of Pooling Size
In this section, we detail our investigation into the impact of varying pooling window sizes on the model's capacity to learn inter-step dependencies. The of this investigation revolves around the Pooling Module, which executes lightweight temporal averaging operation. Given an input sequence tokens [z1, Z2, .2,,.., and a pooling size 1:
For the initial P positions in the output sequence (i.e., .h...., Vr-1), simpler averaging or direct mapping is typically employed. For instance, one might set Ei Ty for
P.
For all subsequent positions i P, the output token is computed the average of the P P-1
input tokens ending r;: 34 Σ kk=0 Ti-k-
the specific instance where P 2, as is used in some of our primary configurations, this module's
output sequence starts with 8/1 = z1, followed by y/4 z; t for é 1.
2
We experiment with pooling sizes (P) of 1, 2, 5, 10, 15, and 20 for this module. The results, shown in Figure [ST a), indicate clear trend. Notably, pooling size of P (which effectively reduces the model to an MLP architecture by negating any temporal averaging beyond the current token if we consider y; .z;) results in a significant performance drop due the lack of contextual history integration. In contrast, better results are achieved with smaller pooling windows, particularly for P values between and However, further increasing the window size beyond this range proves detrimental, hindering effective dependency modeling. Based on the design in [15 whhchh sets filter length to current and previous timesteps per Markov assumption choose pooling size of P = 2 to maintain efficiency and temporal locality.
28

---

Expert Level:
Expar
94.8
Performunce
90.4
90.1
89.0
87.6
86.1
80
10
20
10
20
160
(a) Pooling Size
(b) Context Length
Figure S1: Performance comparison Hopper-medium task: (a) Impact of varied pooling size; (b) Impact of varied context length.
F
Discussion
F.1
Context Length
In our experiments on the Hopper-medine task, DHi demonstrates robust performance across both small and large context lengths by replacing the attention length-with wwtth size-2 average pooling layer, as shown in Figure [ST.b. For example, the model achieves a performance score of 86.1 with a context length of 10 and 89.0 with context length of 160. This consistent performance, regardless of the context window length, underscores the model's resilience to variations in context length. Furthermore, by substituting the attention mechanism with a pooling layer, our approach mitigates the potential issue of out-of-distribution generalization, which can occur when there is mismatch between context lengths during training and testing. In summary, our method effectively reduces sensitivity to context length, ensuring that the model maintains high quality across range of settings.
dependent context
F.2 Sparse Reward Processing
We assess DHi's performance in a delayed (sparse) reward setting, where rewards are not provided throughout the trajectory but are instead accumulated and given at the final timestep [3] TIb displays the results for both delayed and dense reward conditions the D4RL-Hopper dataset. The results show that delayed returns minimally affect DHi and DT while CQL. collapses, indicating the resilience of sequence modeling approaches in such conditions. Additionally, compared the straightforward method, DHi with convolution- Token Merger achieves notable performance improvement, suggesting that sparse rewards do not affect the local modeling capability of CNNs.
Table $10: Results on the Hopper datasets with sparse rewards, where dataset abbreviations are: 'mediumm as 'm', 'medium-replay" as 'm-r', and 'medium-exper'' as 'm-e'.
Dense Setting
Sparse Setting
Task Name
CQL
DT
DHi
CQL
DT
DHi
halfcheetah-m
49.2
42.6
43.4
1.0
42.2
43.5
hopper-m
69.4
67.6
90.1
23.3
57.3
walker2d-m
83.0
74.0
79.9
0.0
69.9
79.6
halfcheetah-m-rr
45.5
36.6
41.5
7.8
33.0
40.4
hopper-m-r
95.0
82.7
97.7
7.7
50.8
94.3
walker2d-ma m-r
77.2
66.6
81.2
3.2
51.6
74.6
Average
69.9
61.7
72.3
7.2
50.8
69.0
29

---

G Visualization
G.1 D4RL Visualization
We observe that the "attention dispersal" problem appears to be core architectural issue, rather than direct consequence of trajectory quality. As shown in Figure SSZ our visualizations across the medium-expert and medium-replay datasets show similar dispersed attention patterns in the standard DT. Importantly, these visualizations also indicate that this issue effectively mitigated in each scenario after applying our Token Merger, resulting in a more structured and focused attention maps. These observations support our conclusion that the benefit of our hierarchical structure tied to specific data type, but rather provides more robust foundation for policy learning across the various datasets of differing qualities found in the D4RL benchmark.
(a) DT
(b) DT with Token Merger
Figure S2: Attention maps for DT (left) and DT with Token Merger (right) MuJoCo Hopper- medium-expert and Hopper-medie pooper-medium-replyy environments.
G.2 Real-World Visualization
Visualization DHi in real-world applications, as shown in Figure
S3a
S3b
S3c
(a) Lift cube by DHI..
(b) Pick up teddy bear and put in the bowl by DHi
(c) Lift red cuhe and put it the plate by DHi.
(d) Lift red cuhe and put it the plate by DT. Figure S3: Visualization of DHi and DT real-world applications
30


Reviews:
[{"review_id": "XMyci0N86D", "reviewer": "Reviewer_ZvZX", "summary": "This paper investigates the limitations of Decision Transformer (DT) in offline RL, attributing its performance bottleneck to attention dispersal, i.e., high attention entropy that prevents effective local dependency modeling. To address this, the authors propose a hierarchical architecture, Decision HiFormer (DHi), composed of a Token Merger (fusing ⟨state, action, reward⟩ triplets at each timestep) and a Token Mixer (capturing inter-step dependencies via lightweight average pooling). The authors show that attention entropy correlates negatively with rewards, and their Token Merger reduces this entropy, improving performance.", "questions": "- Why does avg pooling inplace of attention even work? It's only capturing information from adjacent previous tokens and that seems to be sufficient. Is this because the hopper task which is locomotion does not really require much past history context and hence pooling works fine? Consequently some other tasks that might require history (in terms of credit assignment) would give different results?\n- Instead of self-attention in tokenmerger if the authors could use cross attention with a learned query I wonder if it might give better performance?\n- Does dataset for real robot experiment has any suboptimal trajectories? \n- I'm open to raising my score if authors provide confidence intervals.", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- The entropy–reward correlation is well-motivated and supported by attention map visualizations. Authors re-identify the issue of local dependency modeling and mitigates it with proposed Merger and Mixer methods.\n- Strong empirical results across D4RL tasks, with thorough ablations on important design decisions (e.g., Token Merger and Mixer variants).\n\nWeakness:\n- No standard errors or confidence intervals are reported, which is crucial for offline RL benchmarks with high variance. Table 2 has some values underscored and it's not clear what that signifies.\n- The proposed approach is essentially a variant of DT and thus inherits its core limitation of being a single-step method, lacking the ability to perform multi-step dynamic programming. As a result, its performance is expected to degrade significantly on datasets with very few or no near-optimal  trajectories, such as antmaze-medium and antmaze-large.", "was_revised": true, "final_justification": "My concerns were fairly addressed and I'm happy with additional experiments that authors did to address those. Hence I've raised my score.", "rebuttal": "Dear Reviewer ZvZX:\n\nMany thanks for your valuable comments and questions, which help us a lot to improve our work. We address your questions as follows.\n\n> W1, Q4: No standard errors or confidence intervals are reported, which is crucial for offline RL benchmarks with high variance. Table 2 has some values underscored and it's not clear what that signifies.\n\n[A1]  We sincerely apologize for the unclear statement. The underlined values indicate the second-best performing method for each task. We will clarify this in the table caption in the revised version.\n\nWe reported the standard deviations on the Gym domain in Appendix Table 5 due to limited space in Table 2 in the maintext. The scores and standard deviations for all tasks are shown in the table below:\n\n**Table R1. Performance of DHi in Gym Domain.**\n\n| **Gym Tasks**       | **DHi** |\n| --------------------------- | --------------- |\n| halfcheetah-medium-expert | 94.2 ± 1.0   |\n| hopper-medium-expert      | 111.7 ± 2.6  |\n| walker2d-medium-expert    | 109.6 ± 1.0  |\n| halfcheetah-medium        | 43.4 ± 1.0   |\n| hopper-medium             | 90.1 ± 2.9   |\n| walker2d-medium           | 79.9 ± 2.7   |\n| halfcheetah-medium-replay | 41.5 ± 1.0   |\n| hopper-medium-replay      | 97.7 ± 1.9   |\n| walker2d-medium-replay    | 81.2 ± 8.3   |\n\n**Table R2. Performance of DHi ​**​**in**​**​ Adroit Domain.**\n\n| **Adroit Tasks** | **DHi** |\n| ------------------------ | --------------- |\n| pen-human              | 86.6 ± 11.8  |\n| hammer-human           | 31.2 ± 5.6   |\n| door-human             | 25.2 ± 1.9   |\n| pen-cloned             | 89.1 ± 11.1  |\n| hammer-cloned          | 44.6 ± 9.5   |\n| door-cloned            | 23.6 ± 2.1   |\n\n**Table R3. Performance of DHi in Kitchen Domain.**\n\n| **Kitchen Tasks** | **DHi** |\n| ------------------------- | --------------- |\n| kitchen-complete        | 55.0 ± 4.5   |\n| kitchen-partial         | 75.0 ± 5.3   |\n\n**Table R4. Performance of DHi in AntMaze Domain.**\n\n| **AntMaze Tasks** | **DHi** |\n| ------------------------- | --------------- |\n| antmaze-umaze           | 86.9 ± 11.9  |\n| antmaze-umaze-diverse   | 84.0 ± 13.6  |\n\n> W2: The proposed approach is a variant of DT and inherits its limitation as a single-step method. As a result, its performance is expected to degrade significantly on datasets with few near-optimal trajectories, such as antmaze-medium and antmaze-large.\n\n[A2]  We thank the reviewer for this insightful comment. We agree that as a thoughtful improvement on Decision Transformer, our DHi inherits its single-step, autoregressive nature. Consequently, on challenging tasks like `antmaze-medium` and `antmaze-large`, its performance may be constrained by the quality of the dataset's trajectories, a recognized challenge for this family of models.\n\nTo further investigate this limitation, we experimented with integrating a learned Q-function into DHi, enabling explicit bootstrapping across timesteps. This significantly improved performance on `antmaze-medium` (e.g. from 14.0 to over 40.0), demonstrating that our architecture is fully compatible with long-horizon bootstrapping strategies and benefits greatly from them. We will include detailed results and analysis in the revised Appendix. We believe these are complementary strengths, and combining structured sequence modeling with value estimation is a promising future direction.\n\n> Q1: Why avg pooling inplace of attention work? Would other tasks that require history give different results?\n\n[A3]  In our framework, long-term credit assignment is handled via the return-to-go (RTG) conditioning, which encodes the global task objective. The role of the pooling layer is to ensure that, at each step, the local policy execution remains coherent and responsive to this conditioning. Our analysis (Sec. 3.2, Appendix A/B) shows that this Markovian bias helps reduce representational noise and attention dispersion, especially in offline RL where historical context may be suboptimal or misleading. By enforcing a local Markov structure, it prevents the policy from being distracted by a noisy, distant past, because one cannot achieve a long-term goal if the short-term steps are erratic.\n\nFurthermore, this principle demonstrates strong generalization on tasks that require significant credit assignment. Our state-of-the-art results on complex domains like AntMaze (requiring trajectory stitching) and Kitchen (mastering sub-routines) already show that a robust local policy is a prerequisite for long-horizon success. To further validate this on a classic benchmark, we also evaluated our model in the ​**Atari domain**​. Atari is renowned for requiring credit assignment across long horizons, as rewards are often significantly delayed from the actions that cause them, and the agent must learn from high-dimensional image inputs. **In the Atari domain, ​**​**as shown in Table R5, ​**​**our DHi model achieved an average score of 36.2, outperforming both the original DT (34.8) and DC (35.9).**\n\n**Table R5. Performance in the Atari Domain.**\n\n| **Atari Tasks** | **DT** | **DC** | **DHi (Ours)** |\n| ----------------------- | -------------- | -------------- | ---------------------- |\n| Asterix               | 5.2          | 6.5          | 7.0                  |\n| Frostbite             | 25.6         | 27.8         | 29.7                 |\n| Pong                  | 105.6        | 106.5        | 105.6                |\n| Seaquest              | 2.7          | 2.6          | 2.3                  |\n| **Average**     | 34.8         | 35.9         | 36.2                 |\n\n> Q2: Instead of self-attention in tokenmerger if the authors could use cross attention with a learned query I wonder if it might give better performance?\n\n[A4]  We thank the reviewer for this insightful suggestion. We replaced our convolution-based Token Merger with a cross-attention module using a variable number of learned query vectors (e.g., 1 or 3) and conducted experiments on the D4RL MuJoCo Gym benchmark. The results are summarized in Table R6. In comparison, our proposed DHi model (with a convolution merger) achieves an average score of 83.3 on these tasks, outperforming the cross attention variants by a large margin.\n\n**Table R6. Performance of DHi with Cross-Attention Token Merger in Gym Domain.**\n\n| Gym Tasks         | Number of Query Vectors = 1 | Number of Query Vectors = 3 |\n| ------------------- | ----------------------------- | ----------------------------- |\n| halfcheetah-m-e   | 85.4 ± 7.8                 | 77.2 ± 11.2                |\n| hopper-m-e        | 111.0 ± 0.5                | 111.1 ± 0.2                |\n| walker2d-m-e      | 105.8 ± 3.3                | 104.1 ± 4.1                |\n| halfcheetah-m     | 42.9 ± 0.6                 | 43.3 ± 0.5                 |\n| hopper-m          | 72.7 ± 2.6                 | 62.2 ± 6.0                 |\n| walker2d-m        | 78.2 ± 6.2                 | 77.1 ± 5.3                 |\n| halfcheetah-m-r   | 38.0 ± 0.8                 | 37.1 ± 1.0                 |\n| hopper-m-r        | 36.7 ± 0.0                 | 31.1 ± 0.6                 |\n| walker2d-m-r      | 13.4 ± 5.0                 | 14.3 ± 4.1                 |\n| **Average** | 64.9                        | 61.9                        |\n\nWhile the cross-attention approach performs reasonably well on the `medium` and `medium-expert` tasks, its overall average score is significantly brought down by its poor performance on the `medium-replay` datasets. The `medium-replay` datasets are characterized by a high degree of noise and suboptimality, as they contain a mixture of trajectories from different policies of varying quality.\n\nWe hypothesize that this is where the flexibility of cross-attention becomes a liability. On cleaner data, it can learn an effective fusion rule. However, on the noisy `medium-replay` data, this highly flexible mechanism is more prone to overfitting to spurious correlations present in the suboptimal trajectories, failing to learn a generalizable representation. In contrast, the strong inductive bias of a simple 1D convolution (leveraging locality and weight sharing) acts as a robust regularizer, allowing it to distill a more consistent representation even from noisy data.\n\nThis outcome provides strong empirical support for our paper's central \"Less is More\" argument. It suggests that for this specific task, a simpler architecture with a stronger inductive bias (like a 1D convolution that leverages locality and weight sharing)  is more effective and robust. We are grateful for this suggestion, as this ablation study further solidifies our architectural choices and provides a valuable data point that strengthens the core message of our paper.\n\n> Q3: Does dataset for real robot experiment has any suboptimal trajectories?\n\n[A5]  We thank the reviewer for this thoughtful question. To clarify, the real-world dataset presented in our main paper consists of successful human demonstrations. While these trajectories achieve the task goal, they inherently contain natural sub-optimality—slight hesitations, inefficient paths, and variability in style—which is characteristic of realistic expert data. Our model's 100% success rate on this data reflects its ability to learn from such realistic imperfections.\n\nMotivated by this question, we conducted an additional test to evaluate DHi's robustness under explicitly corrupted datasets. For the “Place task” —where the robot must pick and insert a small cube into a cup—we introduced 20% **failure demonstrations** into the training set (e.g., dropping the cube, failing to grasp). This creates a challenging dataset with noisy and contradictory supervision. Despite this,  DHi still achieved a ​**success rate of 7 out of 11 trials (\\~64%)**, ​demonstrating its ability to selectively extract useful structure from partial failures and learn an effective policy. We will add this new result with video demonstrations and discussion to the final manuscript."}, {"review_id": "0sBIu8C9H5", "reviewer": "Reviewer_sz54", "summary": "The paper \"Less is More: an Attention-free Sequence Prediction Modeling for Offline Embodied Learning\" introduces Decision HiFormer (DHi), an alternative to Transformer-based models in offline reinforcement learning.\n\nThe authors analyze the Decision Transformer's (DT) attention mechanism, finding \"attention dispersal\" due to inconsistent state-action-reward distributions, which hinders local dependency modeling.\n\nDHi is a hierarchical framework that addresses this by:\n1.  **Intra-step modeling:** A \"Token Merger\" (convolutional) fuses each (state, action, reward) triplet into a single representation.\n2.  **Inter-step modeling:** A lightweight, parameter-free \"Token Mixer\" (Average Pooling) captures dependencies across timesteps, preserving the Markov property.\n\nDHi significantly outperforms DT on the D4RL benchmark, achieving a 73.6% improvement in inference speed and a 9.3% gain in policy performance. It also demonstrates strong generalization in real-world robotic manipulation, yielding smoother trajectories and faster, more efficient inference.\n\nIn summary, the paper offers a new understanding of attention in offline RL, proposes an efficient and effective hierarchical model (DHi), and validates its benefits in both simulated and real robotic tasks.", "questions": "*   **Relation to other \"Attention-Free\" Transformers:** While the paper compares to DC (Decision ConvFormer) which uses convolutions, a more in-depth discussion of how DHi's hierarchical decomposition (Merger then Mixer) differentiates it from other non-Transformer architectures that also aim to reduce attention (e.g., S4, Mamba) would strengthen the originality claims. The specific combination of a convolutional merger and pooling mixer is novel, but the underlying ideas of local processing and simple mixers have precursors. \n\n\n*   **Computational Cost of Training vs. Inference:** While inference speed improves, the paper doesn't detail training time or resources needed for DHi versus DT or other baselines. This is vital for practical use and could affect cost-effectiveness.", "limitations": "Please see the weakness section above.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths:\n\n*   **Significance:** Addresses crucial offline RL limitations (Transformer inefficiencies, local dependency issues), boosting performance (9.3% gain) and efficiency (73.6% faster inference), and demonstrating real-world robotic applicability.\n*   **Originality:** Introduces \"attention entropy\" for analyzing Transformer limitations. Proposes a novel hierarchical architecture (DHi) with a unique convolutional \"Token Merger\" for intra-step and a parameter-free \"Average Pooling Token Mixer\" for inter-step modeling.\n*   **Clarity:** Well-structured, clearly written, and effectively uses visualizations and tables to explain complex concepts and present results.\n*   **Quality:**\n    *   **Thorough Analysis of DT's Limitations:** The paper's empirical analysis of DT's attention mechanism using attention entropy provides a solid foundation for its proposed solution. The correlation between high attention entropy and performance drops is clearly demonstrated.\n    *   **Systematic Design and Evaluation of DHi Components:** The authors systematically investigate various Token Merger designs (straightforward, concatenation, self-attention, convolution, pooling, average) and Token Mixers, providing ablation studies (Table 1) to justify their architectural choices (convolutional Token Merger + Average Pooling Token Mixer).\n    *   **Strong Empirical Results:** DHi consistently outperforms various state-of-the-art baselines on the challenging D4RL benchmark across different domains (Gym, Adroit, Kitchen, AntMaze). The real-world robotic experiments further strengthen the empirical validation.\n    *   **Theoretical Justification:** The paper offers theoretical analysis, including Theorem 3.1 on entropy reduction with the Token Merger and proofs in the appendix for the pooling layer's expressiveness in Markov problems, bolstering the architectural decisions.\n\n\n\n### Weaknesses:\n\n*   **Quality:**\n    *   **Scope of \"Markov Property\" Preservation:** The claim of the pooling layer preserving the Markov property (Line 60) and its theoretical proof (Appendix A, B) are significant, but the implications of this for long-term credit assignment or complex, non-Markovian real-world tasks could be more thoroughly discussed. \n    *   **Robustness to Diverse Trajectory Quality:** The D4RL benchmark has varying trajectory qualities (e.g., medium-expert vs. medium-replay). While DHi performs well overall, a more nuanced discussion on how the attention dispersal issue and DHi's solution specifically manifest or perform across different levels of trajectory quality could provide deeper insights.\n\n*   **Significance:**\n    *   **Limited Generalization to Other RL Paradigms:** The paper explicitly states \"A limitation of this study is the lack of evaluation on the transferability of our approach to other RL paradigms, which we leave for future investigation\" (Lines 349-350). This suggests that while the hierarchical modeling is effective for offline RL (specifically sequence prediction), its broader applicability to online RL, model-based RL, or different policy representation forms remains an open question. It would be helpful to specify the type of problem where this method performs best.", "was_revised": true, "final_justification": "I sincerely appreciate the authors' effort to address the questions. The additional experimental results are helpful, and their detailed responses largely clarify my questions. I will maintain my positive score.", "rebuttal": "Many thanks for your valuable comments and questions, which help us a lot to improve our work. We address your questions as follows:\n\n> [W1, about the markov property] The claim of the pooling layer preserving the Markov property  and its theoretical proof are significant, but the implications of this for long-term credit assignment or complex, non-Markovian real-world tasks could be more thoroughly discussed.\n\n[A1]  We sincerely thank the reviewer for this insightful point, which helps us clarify the implications of our theoretical claims.\n\nIn our framework, long-term credit assignment is handled via the return-to-go (RTG) conditioning, which encodes the global task objective. The role of the pooling layer is to ensure that, at each step, the local policy execution remains coherent and responsive to this conditioning. Our analysis (Sec. 3.2, Appendix A/B) shows that this Markovian bias helps reduce representational noise and attention dispersion, especially in offline RL where historical context may be suboptimal or misleading. By enforcing a local Markov structure, it prevents the policy from being distracted by a noisy, distant past, because one cannot achieve a long-term goal if the short-term steps are erratic.\n\nSimilarly, for complex, non-Markovian real-world tasks, our claim is not that the environment is Markovian, but that introducing a Markovian architectural bias serves as a powerful regularizer. It embodies a \"Less is More\" trade-off, where a short, robustly processed history is often more effective than a long, noisy one that is susceptible to spurious correlations. Our strong performance on inherently non-Markovian tasks like the Kitchen benchmark and real-world robotics provides empirical validation for this design choice. These results support our claim that local Markovian stability complements—rather than restricts—global temporal reasoning. We are grateful for this comment and will add a more thorough discussion of these nuances to the revised version of our paper.\n\n> [W2, about the varying trajectory qualities] Robustness to Diverse Trajectory Quality: The D4RL benchmark has varying trajectory qualities (e.g., medium-expert vs. medium-replay). While DHi performs well overall, a more nuanced discussion on how the attention dispersal issue and DHi's solution specifically manifest or perform across different levels of trajectory quality could provide deeper insights.\n\n[A2]  We thank the reviewer for this insightful question regarding the model's robustness to varying trajectory qualities. Our analysis suggests that the \"attention dispersal\" problem appears to be a core architectural issue, rather than a direct consequence of trajectory quality. Our visualizations across the `medium-expert`​ ​​、​`medium`​ ​and `medium-replay` datasets show similar state-dominated, dispersed attention patterns in the standard DT. Importantly, these visualizations also indicate this issue is effectively mitigated in each scenario after applying our Token Merger, resulting in a more structured and focused attention map. These observations support our conclusion that the benefit of our hierarchical structure is not tied to a specific data type, but rather provides a more robust foundation for policy learning across the various datasets of differing qualities found in the D4RL benchmark. We will add these visualizations and a more detailed discussion to the appendix to make this point explicit.\n\n> [W3, about the application] This suggests that while the hierarchical modeling is effective for offline RL (specifically sequence prediction), its broader applicability to online RL, model-based RL, or different policy representation forms remains an open question. It would be helpful to specify the type of problem where this method performs best.\n\n[A3]  We sincerely thank the reviewer for raising this important question regarding the broader applicability of our work. Our response is twofold: first, we clarify the problem domain where DHi is best suited; second, we discuss its potential in other reinforcement learning paradigms.\n\nOur method is specifically designed for the offline reinforcement learning setting and demonstrates its strongest performance in this context, especially when learning from large, static datasets that may contain significant noise and suboptimal trajectories. The core strength of DHi lies in its ability to identify reliable local patterns while filtering out spurious long-range correlations from a noisy history, enabling the learning of stable and robust policies.\n\nRegarding its broader applicability to other RL paradigms:\n\n* **Online RL.** Directly applying DHi is non-trivial, as it is designed for imitation from a fixed dataset, not active exploration. However, its architecture could be highly valuable as a stable and efficient policy network (actor) within a larger actor-critic framework. The stability it provides could potentially lead to smoother and more reliable online policy updates.\n* **Model-Based RL. ​**While DHi is a policy model, its core architectural principles may be adaptable to model-based reinforcement learning. In this context, the goal is to learn a world model that predicts the next state s′ given the current state s and action a—a task that could potentially exhibit the attention dispersal issue we identified. A similar hierarchical approach, in which a Token Merger is used to fuse the state-action pair and the predicted next state into a compact representation before sequence modeling, may lead to more stable and accurate world models.\n\nIn summary, while the current work is empirically grounded in and validated for offline RL, we believe its underlying principles for stable sequence modeling have exciting potential for broader applications. We see this as a promising avenue for future research and have clarified this focus in the revised manuscript.\n\n> [Q1, about the relation to other \"Attention-Free\" Transformers] While the paper compares to DC (Decision ConvFormer) which uses convolutions, a more in-depth discussion of how DHi's hierarchical decomposition (Merger then Mixer) differentiates it from other non-Transformer architectures that also aim to reduce attention (e.g., S4, Mamba) would strengthen the originality claims.\n\n---\n\n[A4]  We thank the reviewer for this insightful question. Our key contribution lies not merely in the removal of attention, but in the problem-driven diagnosis and principled hierarchical design that this enables. This conceptually differentiates DHi from other \"attention-free\" architectures like S4, Mamba, and DC.\n\nOur work originates from a specific problem in offline RL: the \"attention dispersal\" that occurs when feeding three semantically distinct token types—`⟨R, s, a⟩`—into a flat sequence model. To solve this, we propose a hierarchical decomposition where a Token Merger first addresses the intra-step fusion problem before a Token Mixer models the inter-step dependencies between these unified tokens.\n\nThis two-stage approach clearly distinguishes DHi from other architectures. S4 and Mamba aim to improve general-purpose sequence modeling by replacing attention with efficient alternatives that  excel at long-range dependency modeling. Our contribution is orthogonal; one could theoretically use S4 or Mamba as the Token Mixer within our framework. Similarly, while DC also uses convolutions, it does so in a single-stage manner, co-mingling the intra- and inter-step modeling that our approach explicitly separates.\n\nIn summary, while S4, Mamba, and DC offer powerful alternatives to the standard attention mechanism, DHi offers a new framework designed for the unique structure of the offline RL problem. Our originality lies in identifying the `⟨R, s, a⟩` fusion as a key representational bottleneck and proposing a targeted, hierarchical solution. This is a conceptual contribution distinct from simply proposing a more efficient backbone, a distinction we will further clarify in the related work section.\n\n> [Q2, about the training  resources] While inference speed improves, the paper doesn't detail training time or resources needed for DHi versus DT or other baselines. This is vital for practical use and could affect cost-effectiveness.\n\n[A5]  We thank the reviewer for this critical comment. Table R7 presents the computation time for one training epoch, GPU memory usage, and the number of parameters. These metrics offer a comparative analysis of the computational efficiency between MLP vs DT vs. DC vs DHi (Ours), all of which are trained on a single RTX 4090 GPU.  It is observed that DHi is more efficient than DT and DC in terms of training time, GPU memory usage, and the number of parameters.\n\n**Table R7. Comparison of Model Complexity and Training Resources on D4RL.**\n\n| Complexity             | MLP  | DT      | DC     | DHi (Ours) |\n| ------------------------ | ---------- | --------- | -------- | ------------ |\n| Training time (s)      | 30.4     | 34.8    | 34.2   | 31.9       |\n| GPU memory usage       | 1.82GiB  | 2.10GiB | 2.0GiB | 1.84GiB    |\n| Total parameters       | 1.84M    | 2.63M  | 1.86M  | 1.84M      |\n| Token Mixer parameters | 0.0K     | 789.5K  | 13.8K  | 0.0 K      |"}, {"review_id": "qXC6MxkgtQ", "reviewer": "Reviewer_TaqZ", "summary": "This paper proposes to decompose sequence modeling for RL into a “token merger” that combines states, actions, and rewards for a specific timestep and a “token mixer” that combines across timesteps. They study various token mixers and mergers and propose the best one as Decision HiFormer, which uses pooling for the mixer and convolution for the merger. They show that empirically this outperforms DT on D4RL.", "questions": "See Strengths and Weaknesses", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths\n- The proposed hierarchical architecture is novel to the best of my knowledge, and it is interesting that you can have a version of DT that does not use attention\n- The authors perform extensive analysis of attention entropy and systematically consider multiple different ways of combining the tokens\n- The authors show that their model outperforms DT and other baselines on D4RL\n\n\nWeaknesses\n- This work seems incremental in the sense that how to combine the states, actions, and rewards seems more like a design choice/hyperparameter. It is not surprising that by adding a bit of complexity you can improve performance, but this is a tradeoff. The main appeal of DT is its simplicity.\n- It’s confusing that the choice of pooling for the mixer does not seem supported by Table 1, which shows that linear attention would actually be better. The text says that “There is limited performance variation between different token mixers” but the table seems to show significant variations in performance. Could the authors explain these discrepancies? \n- The main advantage of mechanisms that handle long-range dependencies is in tasks that require temporally extended reasoning. How do you ensure you do not lose this kind of reasoning ability with an architecture focused on local dependencies?\n- Empirically the proposed model performs very similarly to DC, which suggests that the benefit of the hierarchical structure might not be large", "was_revised": true, "final_justification": "I appreciate the authors’ detailed response clarifying the empirical benefits of their hierarchical formulation. I have increased my score to 4.", "rebuttal": "Dear Reviewer TaqZ:\n\nMany thanks for your valuable comments and questions, which help us a lot to improve our work. We address your questions as follows.\n\n> [W1] This work seems incremental in the sense that how to combine the states, actions, and rewards seems more like a design choice/hyperparameter. It is not surprising that by adding a bit of complexity you can improve performance, but this is a tradeoff. The main appeal of DT is its simplicity.\n\n[A1]  We sincerely thank the reviewer for this critical comment. We agree that simplicity is a key appeal of the original Decision Transformer (DT). However, we respectfully argue that our work offers a fundamental insight rather than being an incremental design choice.\n\nPrior works have noted the limitations of DT in modeling local dependencies (see [A.1, A.2]), but lacked a precise explanation. In contrast, our work identifies that the unstructured input of state, action, and return tokens is the key cause of attention dispersion, which we quantify via the proposed attention entropy metric (Sec. 3.2, Theorem 3.1). Building on this diagnosis, our proposed Token Merger directly aims at correcting this identified representational flaw, rather than a heuristic tweak. By fusing `⟨s, a, R⟩` triplets before sequence modeling, we reduce entropy and enable more coherent temporal reasoning. The observed gains stem from this representational improvement—not from added complexity. In fact, our model significantly reduces architectural complexity, as shown in Table R8.\n\nImportantly, this motivates a hierarchical modeling framework that explicitly decouples intra-step and inter-step reasoning—rather than treating everything as a flat token sequence. We find that even lightweight, parameter-free average pooling can outperform standard attention while improving inference speed by 73.6% (Figure 1). We also prove that such a pooling operator is sufficient to retain the expressiveness needed for Markovian decision-making, given properly fused intra-step representations (Appendix A and B). Therefore, our final model, DHi, is not only compact and easy to implement, but follows the “less is more” philosophy more closely than the original DT, by replacing costly attention layers with provably sufficient lightweight components.\n\nWe hope this clarifies that our contribution lies in restructuring the modeling paradigm, rather than incrementally tuning architecture design.\n\n> [W2] It’s confusing that the choice of pooling for the mixer does not seem supported by Table 1, which shows that linear attention would actually be better. The text says that “There is limited performance variation between different token mixers” but the table seems to show significant variations in performance. Could the authors explain these discrepancies\n\n[A2]  We thank the reviewer for pointing out this discrepancy and we apologize for the lack of clarity that led to this confusion. While Table 1 shows linear attention achieving the highest performance on the `hopper-medium` task, our statement regarding “limited performance variation” refers to a broader evaluation across the full Gym benchmark. As shown in Appendix Table 6, when averaged over nine Gym tasks, linear attention (78.4) and pooling (78.2) achieve nearly identical scores, supporting our claim that performance differences among token mixers are minor at scale. We will replace Table 1 for Appendix Table 6 in the revised version.\n\nFurthermore, the practical computational cost of linear attention is high in offline RL tasks. While linear attention reduces quadratic complexity to linear in certain implementations, its actual runtime cost is often dominated by high embedding dimensions (d) rather than sequence length (N), especially in offline RL where N ≪ d.\n\nThe table below presents the computation time for one training epoch, GPU memory usage, and the number of parameters. These metrics offer a comparative analysis of the computational efficiency among MLP, Linear Attention (denoted as the model where DHi's pooling is replaced with Linear Attention), DT, DC, and DHi (Ours), all of which are trained on a single RTX 4090 GPU. We observe that DHi outperforms Linear Attention, DT, and DC in terms of training time, GPU memory usage, and the number of parameters.\n\n**Table R8. Comparison of Model Complexity and Training Resources on D4RL.**\n\n| Complexity             | MLP  | Linear Attention | DT      | DC     | DHi (Ours) |\n| ------------------------ | ---------- | ------------------ | --------- | -------- | ------------ |\n| Training time (s)      | 30.4     | 37.4             | 34.8    | 34.2   | 31.9       |\n| GPU memory usage       | 1.82GiB  | 1.90GiB          | 2.10GiB | 2.0GiB | 1.84GiB    |\n| Total parameters       | 1.84M    | 2.42M            | 2.63M  | 1.86M  | 1.84M      |\n| Token Mixer parameters | 0.0K     | 983.8K           | 789.5K  | 13.8K  | 0.0 K      |\n\nGiven that pooling achieves comparable performance to linear attention across the full Gym domain while being significantly more efficient both in theory and in practice, we concluded that it offers a superior trade-off. This choice aligns with our goal of developing a lightweight, effective, and cost-efficient model.\n\nWe are grateful for this question, as it has allowed us to clarify this crucial design decision. We will ensure this reasoning is more clearly articulated in the revised version.\n\n> [W3] The main advantage of mechanisms that handle long-range dependencies is in tasks that require temporally extended reasoning. How do you ensure you do not lose this kind of reasoning ability with an architecture focused on local dependencies?\n\n[A3]  We appreciate the reviewer’s question. Our architecture is carefully designed to enable, rather than sacrifice, long-range reasoning by first ensuring local stability.\n\nAs shown in our entropy analysis (Sec. 3.2), DT’s global self-attention often leads to dispersed focus across unrelated tokens, hindering long-horizon credit assignment. To address this, we first introduce a Token Merger to unify intra-step `⟨R, s, a⟩` triplets, stabilizing local representations. This enables the Token Mixer to focus on inter-step reasoning, where temporal dependencies emerge.\n\nIn our framework, long-term credit assignment is handled via the return-to-go (RTG) conditioning, which encodes the global task objective. The role of the pooling layer is to ensure that, at each step, the local policy execution remains coherent and responsive to this conditioning. Empirically, DHi achieves strong performance on long-horizon tasks such as AntMaze and Kitchen, where effective temporal credit assignment is essential. These findings suggest that local consistency forms a stable foundation for long-range reasoning, particularly in offline RL where distant context may be noisy or misleading.  We further validate this in Appendix F, showing that DHi remains robust across varying context lengths and performs well in sparse-reward Gym domain, where long-horizon credit assignment is critical. Additionally, our framework is modular and compatible with alternative token mixers (e.g., Attention Mixer) where appropriate (Appendix E).  We chose pooling for our final DHi model as it provided a reasonable balance between performance and efficiency in our experiments. Finally, our model is naturally suited for integration into hierarchical systems, where a higher-level planner manages long-term strategy and DHi executes low-level actions reliably. This decoupling offers both generalization and control modularity, consistent with paradigms in embodied decision-making.\n\n> [W4] Empirically the proposed model performs very similarly to DC, which suggests that the benefit of the hierarchical structure might not be large.\n\n[A4]  We thank the reviewer for this observation. While DHi and DC achieve similar scores on several benchmarks, we argue that the benefits of our hierarchical design extend beyond final performance metrics. First, DHi introduced a principled decomposition of intra- and inter-step modeling, grounded in attention entropy analysis (Sec. 3.2). This separation provides clearer inductive biases, better interpretability, and modular extensibility, which are not present in DC's design. As shown in our visualizations (Figure 5), Token Merger creates more structured and meaningful internal representations.\n\nSecond, DHi is substantially more efficient than DC, requiring less training time (31.9s vs. 34.2s), fewer parameters (1.84M vs. 1.86M), and lower GPU memory, while also achieving  significantly faster inference. Furthermore, this efficiency does not sacrifice capability on the most challenging tasks; DHi outperforms DC on the hard-exploration AntMaze domain (85.5 vs. 81.8 average score).\n\nAlso,  the hierarchical formulation makes DHi better suited for integration into complex decision-making pipelines, including hierarchical or multi-agent systems where modular policy design is crucial.\n\nIn summary, DHi offers clear advantages in ​**efficiency, modularity, and interpretability**​, making it a more robust and extensible solution.\n\n### **References**\n\n[A.1] Jinghuan Shang, Kumara Kahatapitiya, Xiang Li, and Michael S. Ryoo. \"StARformer: Transformer with State-Action-Reward Representations for Visual Reinforcement Learning.\" ECCV 2022.\n\n[A.2] Jeonghye Kim, Suyoung Lee, Woojun Kim, and Youngchul Sung. \"Decision ConvFormer: Local Filtering in MetaFormer Is Sufficient for Decision Making.\" ICLR 2024."}, {"review_id": "tTgX3flhDM", "reviewer": "Reviewer_dUjn", "summary": "The paper employs an entropy-based analysis to show that the Decision Transformer (DT) suffers from attention dispersal, which stems from inconsistent attention allocation across the state–action–reward triplet, and that DT also struggles to capture dependencies across timesteps, ultimately limiting its performance. Motivated by these findings and supported by extensive comparisons, the authors propose DHi. The model first uses a convolutional Token Merger to handle intra-step dependencies by fusing each triplet into a single token, and then applies a parameter-free pooling Token Mixer to model inter-step dependencies. DHi achieves strong results on multiple D4RL tasks and real-world robotic manipulation benchmarks.", "questions": "- **Q1:** The heat map on the right-hand side of Figure 2 is a bit confusing. If I understand correctly, in the upper plot each \"$s_i$\" actually corresponds to three columns (or rows), yet DT attends only to the middle column representing the \"state\". Moreover, the lower plot shows temporal attention scores, but DHi is attention-free; does this mean the figure merely replaces DT's triplet input at each timestep with the single token produced by the convolutional token merger? Finally, how is DT's \"inter-step modeling entropy\" across different timesteps computed? Figure 2 does not appear to include a corresponding illustration.\n\n\n- **Q2:** When introducing Theorem 3.1, the authors posit the assumption: “Assuming that the total attention entropy remains unchanged after adding the Token Merger.” What is the intuition behind this assumption? Furthermore, because the merged token changes in form after the Token Merger, the distribution of attention scores might also change. In that case, does the decomposed form stated in Theorem 3.1 still hold?\n\n- **Q3:** Building on Weakness 3, how would DHi perform on maze environments more complex than umaze? For long-horizon planning tasks, might an attention-based mixer be more appropriate?\n\n- **Q4:** How would DHi perform on Gym’s random and expert datasets?\n  \n\n**Other suggestion:** It might be better to replace Table 1 in the main text with Table 6 from the appendix. In the current Table 1, using a convolutional mixer seems to nullify, or even harm, the effectiveness of every merger. What could be causing this?", "limitations": "Please see weakness 3.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "## Stengths: \n- This work is well-motivated. The authors conduct a thorough analysis that pinpoints the shortcomings of DT and, based on these insights, devise the new DHi architecture, which improves DT’s performance.\n- A large suite of experiments on both D4RL benchmarks and real-world tasks convincingly demonstrates DHi's effectiveness.\n- The design of DHi is further supported by theoretical analysis and comprehensive ablation studies.\n- The paper is well written and easy to follow.\n\n## Weaknesses:\n- While Figure 2 conveys the main insights about attention-entropy, several finer details remain confusing.\n- The Theorem 3.1 would benefit from additional discussion or intuition explaining its assumptions.\n- The design of DHi seems to assume that decision making in offline RL is mainly driven by short-range dependencies, which could hinder DHi's performance on tasks requiring long-range planning (e.g., navigation in more complex maze environments).", "was_revised": true, "final_justification": "I very much appreciate the authors' effort to respond to the concerns. The additional experimental results are helpful and the detailed response largely address my concerns/confusion. I will keep my score.", "rebuttal": "Many thanks for your valuable comments and questions, which help us a lot to improve our work. We address your questions as follows.\n\n> [Q1, W1] The heat map on the right-hand side of Figure 2 is a bit confusing. If I understand correctly, in the upper plot each \"si\" actually corresponds to three columns (or rows), yet DT attends only to the middle column representing the \"state\". Moreover, the lower plot shows temporal attention scores, but DHi is attention-free; does this mean the figure merely replaces DT's triplet input at each timestep with the single token produced by the convolutional token merger? Finally, how is DT's \"inter-step modeling entropy\" across different timesteps computed? Figure 2 does not appear to include a corresponding illustration.\n\n[A1]  Thank you for these observations; your understanding of Figure 2 is largely correct and this gives us a valuable opportunity to clarify.\n\n1. Regarding the upper plot, your interpretation is correct that it shows attention overwhelmingly focused on state tokens. To be precise, each label `si` on the axis only marks the column representing the state token. The column to its left represents the return-to-go (`Ri`) and the one to its right represents the action (`ai`); we omitted these extra labels due to space constraints. This visualization demonstrates the \"modality imbalance\" we aimed to highlight.\n2. For the lower plot, your guess is exactly right. This figure is a crucial ablation study showing a *standard DT's* attention map after its input is pre-processed by our Token Merger. Its purpose is to isolate and demonstrate that our merger successfully resolves the attention dispersal problem within the Transformer architecture itself.\n3. The \"inter-step modeling entropy\" is computed by first aggregating all attention weights between pairs of timesteps (i.e., summing the weights between all tokens in timestep `i` and all tokens in timestep `j`) to form a `KxK` matrix, and then calculating the average Shannon entropy from the resulting probability distributions.\n\nWe apologize for the lack of clarity and will revise the figure captions and add a more formal definition to the appendix to make these points explicit. Thank you for helping us improve the paper.\n\n> [Q2, W2] When introducing Theorem 3.1, the authors posit the assumption: “Assuming that the total attention entropy remains unchanged after adding the Token Merger.” What is the intuition behind this assumption? Furthermore, because the merged token changes in form after the Token Merger, the distribution of attention scores might also change. In that case, does the decomposed form stated in Theorem 3.1 still hold?\n\n[A2]   We thank the reviewer for this thoughtful question. The assumption—\"Assuming that the total attention entropy remains unchanged\"——is not meant as a literal description of the model’s behavior, but rather as a conceptual tool to isolate and clarify the mathematical effect of token grouping.\n\nOur goal with this assumption is to construct a controlled thought experiment: holding the overall attention entropy constant allows us to examine how redistributing attention mass over fewer, semantically richer tokens affects entropy concentration. This is grounded in the **chain rule of entropy** and the broader principle that grouping variables typically reduces entropy, assuming no added uncertainty. In other words, if attention were spread across ⟨s, a, R⟩ separately, merging them into one token—without adding noise—should encourage a more focused distribution.\n\nIn practice, the Token Merger alters the token representations and thus the attention distribution. Therefore, Theorem 3.1 is not intended as a formal proof of the final system's state but rather as a strong theoretical motivation for our design choice. It establishes that grouping is a principled way to encourage a more focused attention structure. The ultimate validation of this theoretically-motivated design is provided by our empirical results. As clearly shown in Figure 2, the introduction of the Token Merger leads to a significant and directly measurable decrease in the model's actual attention entropy, which in turn correlates with improved performance.\n\n> [Q3, W3] Building on Weakness 3, how would DHi perform on maze environments more complex than umaze? For long-horizon planning tasks, might an attention-based mixer be more appropriate?\n\n[A3] We thank the reviewer for this thoughtful question.\n\n1. The question rightly points to the performance limitations in environments like `antmaze-medium-diverse` or `antmaze-large-diverse`. In these tasks, DT is likely to fail (with scores of zero) and DHi also achieves only partial success (i.e., a score of 14.0 on `antmaze-medium-diverse`). This is due to a core limitation of DT-based models: being single-step methods, they lack the ability to perform multi-step dynamic programming. This becomes particularly challenging in AntMaze-style environments where success depends on trajectory stitching, which is fundamentally a severe credit assignment problem that requires identifying useful behaviors from datasets composed almost entirely of poor-quality data.\n2. Furthermore, we further conducted experiments with an attention-based mixer. We found that an attention-based mixer also fails on these extremely difficult maze tasks. This suggests the primary bottleneck is not simply a lack of long-range dependency modeling—which attention excels at—but the more fundamental algorithmic challenge of credit assignment from poor data, which likely requires different methods like integrating Q-learning for guidance. Tackling this is an important direction for future work but is beyond the scope of our current paper.\n3. On the other hand, we agree that for long-horizon tasks where complex temporal reasoning—rather than trajectory stitching from poor data—is the main challenge, an attention-based mixer could be a more appropriate choice. This, however, highlights a key strength of our overall hierarchical framework: its flexibility. The framework itself is not rigidly tied to a single mixer type. As demonstrated in our ablation studies (Table 1), the pooling mixer can be seamlessly swapped for a self-attention or linear attention mixer to cater to the specific demands of a task. Our final DHi model uses pooling because it offered the best performance-efficiency trade-off across the broad set of D4RL benchmarks, reinforcing our \"Less is More\" thesis for general offline RL problems.\n\n> [Q4] How would DHi perform on Gym’s random and expert datasets?\n\n[A4]  We sincerely thank the reviewer for this critical question.  We have conducted new experiments on the `expert` and `random` datasets to address this. The results are summarized in Table R9 below. \n\nAs the `expert` and `random` datasets represent the best and worst cases for data quality, the models (DT, DC, and DHi) exhibit consistent and predictable behavior. On the `expert` data, all models achieve near-optimal performance, with DHi consistently matching or slightly outperforming the baselines. Conversely, on the `random` ​dataset, all models perform poorly, as expected. These datasets lack any coherent behavioral structure to learn from, and thus no method can extract a meaningful policy. DHi's performance in this setting is comparable to DT and DC, showing low return and low variance.\n\nThe results indicate that DHi does not overfit or hallucinate policy structure when none exists, highlighting the stability of its architectural bias.  DHi behaves robustly and as expected at the boundaries of data quality, reinforcing its reliability as an offline RL algorithm. We will add this ablation to the appendix.\n\n**Table R9: Performance on Expert and Random Datasets.**\n\n| Task Name             | DT           | DC           | DHi          |\n| ----------------------- | -------------- | -------------- | -------------- |\n| halfcheetah-expert    | 90.1 ± 1.2  | 90.9 ± 1.1  | 94.3 ± 1.0  |\n| hopper-expert         | 110.7 ± 1.0 | 111.1 ± 1.1 | 111.9 ± 1.0 |\n| walker2d-expert       | 108.8 ± 0.6 | 109.3 ± 1.6 | 109.4 ± 0.4 |\n| halfcheetah-random    | 2.3 ± 0.0   | 2.3 ± 0.0   | 2.2 ± 0.0   |\n| hopper-random         | 8.7 ± 0.0   | 9.5 ± 0.1   | 8.9 ± 0.0   |\n| walker2d-random       | 6.9 ± 0.4   | 7.0 ± 0.9   | 7.7 ± 0.6   |\n| **mujuco-mean** | 54.6         | 55.0         | 55.7         |\n\n> **Suggestion: It might be better to replace Table 1 in the main text with Table 6 from the appendix. In the current Table 1, using a convolutional mixer seems to nullify, or even harm, the effectiveness of every merger. What could be causing this?**\n\n[A5]  We appreciate the reviewer’s insightful observation and constructive suggestion. We agree that focusing on a single task in Table 1 may not provide a complete picture, and we will feature the more comprehensive results from Appendix Table 6 more prominently in our revised manuscript. It is a sharp observation that the combination of a Convolutional Merger with a Convolutional Mixer shows a performance drop specifically on the `hopper-medium` task. However, this appears to be a task-specific interaction. When we evaluate the average performance across all nine Gym tasks (as shown in Appendix Table 6), this combination scores well, suggesting our hierarchical design is effective in a broader context.\n\nRegarding the performance on `hopper-medium`, we hypothesize that an \"over-smoothing\" effect from the dual convolutional layers may adversely affect that environment’s dynamics. We appreciate this question as it highlights the importance of comprehensive benchmarking, which we will clarify in the revision."}]

Author rebuttal:

Meta-review: The paper proposes Decision HiFormer (DHi), an alternative to Transformer-based models in offline reinforcement learning, and also a hierarchical framework that decomposes sequence modeling into intra-step relational modeling and inter-step modeling. DHi significantly outperforms DT on the D4RL benchmark, achieving a 73.6% improvement in inference speed and a 9.3% gain in policy performance. It also demonstrates strong generalization in real-world robotic manipulation, yielding smoother trajectories and faster, more efficient inference. 

The reviewers are all positively inclined toward this submission, with two recommending Accept and two Weak Accept. Given the methodological novelty, and strong empirical results from the reviewers, I am pleased to recommend acceptance of this paper.
Decision: Accept (poster)