Title: Decision Mamba: A Multi-Grained State Space Model with Self-Evolution Regularization for Offline RL
Venue: NeurIPS 2024
Accepted: yes
Keywords: Decision Transformer

Mamba

Offline Reinforcement Learning

Autoregressive Sequence Modeling

State Space Model
Abstract: While the conditional sequence modeling with the transformer architecture has demonstrated its effectiveness in dealing with offline reinforcement learning (RL) tasks, it is struggle to handle out-of-distribution states and actions.
Existing work attempts to address this issue by data augmentation with the learned policy or adding extra constraints with the value-based RL algorithm. However, these studies still fail to overcome the following challenges: (1) insufficiently utilizing the historical temporal information among inter-steps, (2) overlooking the local intra-step relationships among return-to-gos (RTGs), states, and actions, (3) overfitting suboptimal trajectories with noisy labels. To address these challenges, we propose $\textbf{D}$ecision $\textbf{M}$amba ($\textbf{DM}$), a novel multi-grained state space model (SSM) with a self-evolving policy learning strategy.
DM explicitly models the historical hidden state to extract the temporal information by using the mamba architecture. To capture the relationship among RTG-state-action triplets, a fine-grained SSM module is designed and integrated into the original coarse-grained SSM in mamba, resulting in a novel mamba architecture tailored for offline RL. Finally, to mitigate the overfitting issue on noisy trajectories, a self-evolving policy is proposed by using progressive regularization. The policy evolves by using its own past knowledge to refine the suboptimal actions, thus enhancing its robustness on noisy demonstrations. Extensive experiments on various tasks show that DM outperforms other baselines substantially.

Decision Mamba: A Multi-Grained State Space Model with Self-Evolution Regularization for Offline RL

Abstract

While the conditional sequence modeling with the transformer architecture has demonstrated its effectiveness in dealing with offline reinforcement learning (RL) tasks, it is struggle to handle out-of-distribution states and actions. Existing work attempts to address this issue by data augmentation with the learned policy or adding extra constraints with the value-based RL algorithm. However, these studies still fail to overcome the following challenges: (1) insufficiently utilizing the historical temporal information among inter-steps, (2) overlooking the local intrastep relationships among return-to-gos (RTGs), states and actions, (3) overfitting suboptimal trajectories with noisy labels. To address these challenges, we propose D ecision M amba ( DM ), a novel multi-grained state space model (SSM) with a selfevolving policy learning strategy. DM explicitly models the historical hidden state to extract the temporal information by using the mamba architecture. To capture the relationship among RTG-state-action triplets, a fine-grained SSM module is designed and integrated into the original coarse-grained SSM in mamba, resulting in a novel mamba architecture tailored for offline RL. Finally, to mitigate the overfitting issue on noisy trajectories, a self-evolving policy is proposed by using progressive regularization. The policy evolves by using its own past knowledge to refine the suboptimal actions, thus enhancing its robustness on noisy demonstrations. Extensive experiments on various tasks show that DM outperforms other baselines substantially.

1 Introduction

Offline Reinforcement Learning (RL) [13, 27, 29, 38] has attracted great attention due to its remarkable successes in the fields of robotic control [5, 36] and games [3, 32, 50]. As transformer [49] has exhibited powerful sequential modeling abilities in natural language processing [4, 43] and computer vision [10, 42], many efforts [6, 8, 25, 61] have been made on applying this architecture to offline RL tasks. Transformer-based methods view the reward/return-to-go (RTG), state, and action as a sequence, and then predict actions by using the transformer encoder. However, it often fails to make correct decisions when encountering out-of-distribution states or actions, showing limited robustness. Previous work attempts to address this issue from the perspective of data augmentation [51, 64] and objective constraints [6, 53, 61]. However, they introduce a significant number of noises or the overestimation bias. Thus, how to enhance model robustness remains a highly challenging and insufficiently explored issue.

In this study, we offer two novel perspectives on improving model robustness through both the model architecture and learning strategy. In terms of the model architecture, (1) although previous

studies have made some modifications to the transformer architecture [23, 44, 52], they have not fully utilized inter-step information, particularly historical information which is critical for decision-making processes. For example, the robot can adjust its subsequent routes based on the historical information of failed paths for completing the navigation task; (2) furthermore, most existing approaches adopt transformer to model the flattened trajectory as a sequence, while ignoring the structural trajectory patterns of the causal intra-step relationship among R TGs, s tates, and a ctions (RSAs). A RL policy typically predicts the next action given the current state based on the RTG. Thus, this kind of finegrained intrinsic connection among RSAs is intuitively beneficial for policy learning. As regards to the learning strategy, (3) there exists a large number of noisy labels in the suboptimal trajectories which hurt the performance of the policy significantly. Although the existing work that generates pseudo trajectories or actions alleviates this problem to some extent [57, 64], it also introduces other biases or errors.

To address the above issues, we propose D ecision M amba ( DM ), a multi-grained state space model with a self-evolving policy learning strategy for offline RL. In order to adequately leverage the historical information, we adopt mamba to explicitly model the temporal state among inter-steps, since mamba architecture [15, 18, 40] shows a more effective capability of extracting the historical information. Meanwhile, the causal intra-step relationship is beneficial for the model to understand the common patterns within the local dynamics. Thus, we introduce a fine-grained SSM module to extract the local features of structural patterns among the RSA triplet within each intra-step. Apart from modifying the model architecture and aligning it to the trajectory pattern, we also propose a learning strategy to prevent the policy from overfitting noisy labels. This is achieved by a progressive self-evolution regularization which leverages the past knowledge of the policy itself to refine and adjust the target label adaptively.

We conduct comprehensive experiments on Gym-Mujoco and Antmaze benchmark, containing 5 tasks with varying levels of noise and difficulties. The performance of DM surpasses other baselines by approximately 8% with respect to the average normalized score on the three classic Mujoco tasks, showing its effectiveness. In summary, the contributions of this paper are summarized as follows:

Different from the existing conditional sequence modeling work for offline RL with the transformer architecture, we propose Decision Mamba (DM), a generic offline RL backbone built on State Space Models, which leverages the historical temporal information sufficiently for robust decision making.

To extract the casual intra-step relationships, we introduce a fine-grained SSM module and integrate it to the original coarse-grained SSM in mamba, which combines the local trajectory patterns with the global sequential features, achieving the multi-grained modeling capability.

To prevent the policy from overfitting the noise trajectories, we adopt a self-evolving policy learning strategy to progressively refine the target, which uses the past knowledge of the learned policy itself as an additional regularizer to constrain the training objective.

2 Related Work

2.1 Offline Reinforcement Learning with Transformer-based Models

Offline Reinforcement Learning (RL) [7, 13, 22, 27, 29, 38, 56, 57, 67] is widely used for robotic control and decision-making. In particular, transformer-based methods [8, 25, 44] reformulate the trajectories as a state/action/RTG sequence, and predict the next action based on the historical trajectories. However, although the sequence modeling methods formulate offline RL in a simplified form, they can hardly deal with the overfitting problem caused by the suboptimial trajectories in offline data [11, 21, 59]. One line of approaches [34, 51, 64, 66] focused on exploiting data augmentation methods, such as generating additional data via the bootstrap method, or training an inverse dynamics model to predict actions for the large amount of unlabelled trajectories. Another line of work [6, 23, 25, 35, 44, 52] attempted to modify the transformer architecture to explicitly make use of the structural patterns within the training data. Furthermore, substantial efforts [37, 54, 58, 62] have also been made on applying regularization terms to learning policies, such as RvS [11] and QDT [61]. Nevertheless, previous work simply applies transformer to offline RL tasks while seldom considering about adapting the architecture to trajectory learning. Thus, these methods fail to extract

the historical information sufficiently and are unable to capture local patterns thoroughly from the trajectories. In this work, we address these issues by proposing DM, a tailored mamba architecture for offline RL tasks. A fine-grained SSM module is designed in DM to supply fine-grained intra-step information to the coarse-grained inter-steps features. Together with the architecture, we also present a self-evolving policy learning strategy to prevent the model from overfitting noise labels.

2.2 State Space Models for Linear-time Sequence Modeling

Recently, State Space Models (SSMs) show high potentials in various domains, including natural language processing [15-17, 19, 20, 40, 46], computer vision [30, 31, 33, 41, 63, 65] and time-series forecasting [55]. Stemming from signal processing, SSMs capture global dependencies from a sequence more effective in a lightweight structure and shows advantages in compressing the historical information, compared with the transformer architecture. Although SSMs have considerable benefits, it still struggles to perform contextual reasoning. Mamba [15] is thus proposed to alleviate this problem. It introduced a time-varying selective mechanism and a hardware-friendly design, making it as a competitive architecture against with transformer. The Mamba architecture is then adapted to different downstream tasks by considering the characteristics of these tasks. VIM [65] and VMamba [33] introduced 2D SSMs for image understanding. VideoMamba [30] introduced spatiotemporal scan for video understanding. In this work, we take the fine-grained trajectory patterns into consideration, and introduce a multi-grained mamba architecture tailored for RL tasks.

3 Method

3.1 Preliminaries

Decision Transformer for Offline RL. The fundamental Markov Decision Process [12] can be represented as M = ( S , A , T , r, γ ) , where S is the state space, A is the action space, T : S×A → S is the transition function, r : S × A → R is the reward function, and γ ∈ (0 , 1] is the discount factor. Given an offline dataset D µ collected by the behavior policy µ ( a | s ) , offline RL algorithms aim to maximize the rewards. Formally, the iteration process of learning a policy is as below ( k denotes the index of the learning iteration):

When updating the Q function, ( s, a, r, s ′ ) are sampled from D µ but the target action a ′ is sampled from the current policy π k -1 .

Inspired by the great success of sequence generation models in NLP [9, 39, 48], Decision Transformer [8] is proposed to model the trajectory optimization problem as an action prediction procedure. Specifically, it first obtains the return-to-go (RTG) with the reward, i.e., R t = ∑ T i = t r i . Then, the learned policy, which is based on the decoder-only transformer architecture [48], predicts the action sequence a i autoregressively, with the offline trajectory τ = ( s 0 , R 0 , a 0 , . . . , s T , R T , a T ) . The training objective is as follows:

where τ t -l : t = ( s j , R j , a j , . . . , s t , R t ) ( j = min( t -l, 0)) is the input trajectory and l is the length of context window.

SSMs for Linear-Time Sequence Modeling. The State Space Model (SSM) describes the probabilistic dependence between the continuous input signal x ( t ) and the observed output y ( t ) via the latent hidden state h ( t ) as Eq. (5):

Decision Mamba Encoder

0

0

0

1

1

1

2

2

2

Decision Mamba

Group

Norm

CG

Conv1d

CG

SSM

Activation

FG

Conv1d

FG

SSM

L x

Decision Mamba Encoder

Coarse-Grained Scan

Fine-Grained Scan

In order to apply the SSM model to the discrete input sequence instead of the original continuous signal, Structured SSM (S4) [18] discretizes it by a step size ∆ as Eq. (7).

where A = exp(∆ A ) , B = (∆ A ) -1 (exp(∆ A ) -I ) · ∆ B . To this end, the model can forwardpropagate in an efficient parallelizable mode with a global convolution. Due to the linear time invariance brought by the SSM model, it lacks the content-aware reasoning ability which is important in sequence modeling. Therefore, mamba [15] proposes the selective SSM which adds the length dimension to the original parameters (∆ , B , C ) , changing it from time-invariant to time-varying. It uses a parameterized projection to project the size of parameter B,C from (D , N) to (B , L , N) , and ∆ from (D) to (B , L , D) , where D , B , L and N denotes the channel size, batch size, sequence length and hidden size, respectively.

3.2 Decision Mamba

The transformer architecture has been well used in offline RL tasks. Despite its strong ability to understand complete trajectory sequence, it shows limited capabilities in capturing historical information. Thus, we propose a multi-grained space state model to extract the fine-grained local information to supply the coarse-grained global information, namely Decision Mamba (DM), for comprehensively learning the trajectory representation. Figure 1 presents the overall framework of DM.

3.2.1 Multi-Grained Mamba

Trajectory Embeddings Following the sequence modeling [8], we first use multilayer perceptrons (MLPs) to embed the RSAs from the given trajectory τ = ( R 0 , s 0 , a 0 , . . . , R T , s T , a T ) . Then, the trajectory embeddings are added the absolute step position embeddings to attach the position information, similar to the classic usage in the NLP field. Mathematically, it can be formulated as follows:

where e f i ∈ R B × L × N , and [; ] denotes the concatenate operation.

Coarse-Grained SSM Different from the transformer-based methods [8, 25], DM models the historical information before the current i -th step via the latent hidden state h i as shown in Eq. (6). It explicitly represents the feature of history information, rather than only learns such information implicitly. As the number of encoder layers increases, historical information is selectively preserved in the representation h i . To this end, DM is expected to have a better capability to understand the sequential dependencies. It can be formulated as follows:

where h CG i represents the coarse-grained hidden state; InterS3M denotes the coarse-grained SSM.

Fine-Grained SSM Further, to better discern the dependencies among RSA within each step, we gather the feature of each single step to obtain the fine-grained representation via a 1Dconvolution layer, and then introduce a finegrained SSM module for extracting the local pattern among RSA, as shown in the middle part and right part of Figure 1.

It can be formulated as follows:

where h FG i indicates the fine-grained hidden state; IntraS3M means the fine-grained SSM.

Fusion Module For gathering both fine-grained local trajectory patterns and coarse-grained global contextual information, we combine the h FG i with h CG i in each encoder layer and then use the layer normalization to ensure that the multi-grained features have a consistent distribution. In order to remain the important historical information, we add a residual connection. The fusion process can be formulated as follows:

Algorithm 1 Decision Mamba

where h MG i indicates the multi-grained hidden state and LN denotes the layer normalization. The forward propagation procedure of DM is presented in Algorithm 1.

3.2.2 Progressive Self-Evolution Regularization

There are typically amounts of suboptimal trajectories in RL tasks. The previous approaches usually overfit these noisy data and thus lack robustness. Fortunately, the existing literature [26] has shown that deep models learn clean samples (optimal trajectories) at the beginning of the training process, and then overfit the noisy samples (suboptimal trajectories). Inspired by this observation, we propose a progressive self-evolution regularization (PSER), which uses the knowledge of the past policy to refine the noisy labels as supervision for policy learning, thus avoiding fitting the noisy trajectories.

Specifically, we obtain a refined target by combin-

step

step

ing the ground truth and the prediction from the learned policy itself. Let ˆ a k denote the prediction about s from the current policy π k ( a | s ) at k -th iteration. The refined target at k -th can be written as follows:

To obtain more insights about the refined targets Eq. (16), we compare the gradients of the training objectives with the original label and the refined label. The standard Mean Square Error (MSE) loss function of Eq. (3) with the original label can be written as:

In comparison, the loss function with the refined target of Eq. (16) can be rewritten as:

Comparing the objectives of Eq. (17) and Eq. (18), the gradient of L SE,k with respect to the output of policy { a k,i } T i =1 can be derived by:

where ∇L k indicates the gradient of the loss function Eq. (17), and ∇R computes the difference between the past predictions and the targets. From the perspective of gradient back propagation, PSER imposes a regularization constraint on the current policy π k ( a | s ) by smoothing the original target action a k,i with the self-generated label ˆ a k -1 ,i .

Moreover, it is important to determine the value of β in Eq. (18). The β controls the learning procedure, where the policy trusts the given actions if β is set to a large value. As stated above, the policy tends to fit gradually from clean patterns to noisy patterns. Thus, we set the β to dynamically increased values. As β increases, the policy progressively gains more confidence in its own past knowledge. To maintain the learning process stable, we apply the linear growth approach and set a lower boundary. The β at the k -th iteration is computed as follows:

where K is the number of total iterations for training and β K is the hyperparameter.

We replace the original label with the refined target, leading to the objective:

We adopt the MSE loss, and then the objective 21 is converted to:

3.2.3 Training Objective

To make the training procedure more robust, we introduce the inverse training goals: predicting the next state and the next RTG. Individuals often assess the feasibility of actions by envisioning their potential outcomes. Therefore, we expect the policy to predict the post-execution state and RTG based on the predicted action, thus improving its robustness. Specifically, given the trajectory τ = ( R 0 , s 0 , a 0 , . . . , R t , s t ) , Decision Mamba originally predicts the next action ˆ a t . Further, by incorporating the action a t to the original trajectory τ t -l : t , it is also predicts the next RTG ˆ R t +1 and the next state ˆ s t +1 . Compared to the Eq. (3), the training objective of DM with the refined target can be written as follows:

where the ˜ a t is computed by Eq. (16), λ i is the weight hyperparameter, and the sum of λ i is set to 1. Note, we omit the length of context window l for simplicity.

4 Experiment

4.1 Settings

Dataset and Evaluation Metrics. We conduct our experiments on Gym-MuJoCo which is one of the mainstream benchmarks used in offline deep RL [14, 28, 59], including Hopper, HalfCheetah, Walker and Ant tasks. Each task contains medium, medium-expert, medium-replay and expert datasets. To more comprehensively evaluate our proposed method, we also adopt the AntMaze benchmark which is a navigation task of aiming to reach a fixed goal location, with the 8-DoF "Ant" quadraped robot. We evaluate Decision Mamba by using the popular suite D4RL [13]. Following the existing literature [8, 25], we normalize the score for each dataset roughly for comparison, by computing normalized score = 100 × score -random score expert score -random score . More details about dataset and implementation can be found in Appendix A.

Baselines. We compare Decision Mamba with existing SOTA offline RL approaches including Behavioral Cloning (BC), Conservative Q-Learning (CQL) [29], Decision Transformer (DT) [8], Reinforcement Learning via Supervised Learning (RvS) [11], StARformer (StAR) [44], Graph Decision Transformer (GDT) [23], Waypoint Transformer (WT) [2], Elastic Decision Transformer (EDT) [60], and Language Models for Motion Control (LaMo) [45]. Among these methods, CQL stands as a representative of value-based methods, while the other methods belong to supervised learning (SL) approaches. For most of these baselines, we cite the results from the original papers. In addition, we reimplement DT and LaMo for more comparison in different settings by using their repositories. The detailed descriptions of these baselines are presented in Appendix B.

4.2 Overall Results

For a fair comparison, we first conduct experiments on datasets commonly adopted by mainstream approaches. The overall performance is presented in Table 1. It can be observed that Decision Mamba outperforms other baselines in most datasets. On one hand, benefiting from the supervised learning objective, SL-based baselines exhibit a strong ability in the high-quality datasets (M-E), but show weakness in the suboptimal datasets (M/M-R). On the other hand, CQL performs well in the suboptimal datasets due to regularizing the Q-values during training, but struggles to perform well in the high-quality datasets.

For DM, it shows significant improvement over the other SL-based methods. Specifically, it outperforms the best of the baselines by 4%+, especially in suboptimal datasets, e.g., on the medium datasets, the performance of DM surpasses the value-based method CQL and the transformer-based method GDT by around 6% and 9% on average, respectively. This significant improvement demonstrates the robustness of DM in learning from suboptimal datasets, attributed to the multi-grained mamba encoder and PSER module in DM. Note, although DMperforms slightly worse than CQL on the Halfcheetah-M-R and HalfCheetah-M datasets, the difference is not significant. Therefore, the

proposed DM shows stronger overall performance, capable of learning from both high-quality and suboptimal datasets simultaneously.

In order to evaluate our method more comprehensively, we also conduct experiments on other datasets. We adopt representative baselines including BC, DT and LaMo, where LaMo leverages extensive additional natural language corpora and knowledge to enhance the model performance. As illustrated in Table 2, all methods perform exceptionally well on the Expert dataset. However, when it comes to mixing the suboptimal data into the training set, compared with DT, both LaMo and DM exhibit

significant superiority, while DM shows a more pronounced overall enhancement. For instance, DM outperforms LaMo by approximately 10% on the Ant-M dataset and around 6% on average. For AntMaze , it requires composing parts of suboptimal trajectories to form more optimal policies for reaching goals. 'U-D' is more difficult than 'U', and DM shows superiority in these tasks. More comparison results can be found in Appendix C.

4.3 Ablation Study

To investigate the effectiveness of each component in DM, we conduct experiments with different variants of DM. In particular, we compare 3 different implementations: (1) w/o MG removes the multi-grained branch, directly using the sequence feature original extracted from mamba; (2) w/o PSER removes the progressive self-evolution regularization, training the model with the labels in the training dataset; (3) w/o ILO removes the inverse learning objective, only predicting the action in the training procedure. As shown in Table 3, the performance of DM drops significantly without either of these components. Notably, the most substantial performance degradation with about 6% occurs when the PSER module is removed, especially in the suboptimal datasets. This observation verifies the effectiveness of this module in preventing policy from overfitting and thus enhancing its robustness. MG and ILO are also critical for offline RL tasks. Once these two modules are excluded, there is a noticeable reduction in the model's performance.

Table 3: Ablation Results. ' w/o MG/PSER/ILO ' represents removing the module of multi-grained feature extraction, the progressive self-evolution regularization, and inverse learning objectives, respectively. Best results are marked in bold.

4.4 Comparison Results with Different Context Lengths

To validate whether DM can capture the information of inter-step and intra-step, we investigate the performance of our model with different context lengths. We conduct experiments with the context length L = { 20 , 40 , 60 , 80 , 100 , 120 } . Figure 3 shows the comparison results. Regardless of different context lengths, the proposed DM consistently achieves a high score than other baselines among all datasets, showcasing its superiority in capturing the inter-step dependencies in different lengths.

It is noteworthy that BC shows a comparable performance to DT in the Hopper-M and Halfcheetah-M datasets, yet demonstrates a substantial discrepancy in other datasets. We deduce that, in addition to the inherent limitations of the imitation learning paradigm, the BC model that relies solely on MLP also has significant architectural disadvantages. Consequently, it can achieve scores of only 60-70% at most on the expert datasets. When trained on M-R data, BC evidently struggles to learn effectively, achieving only approximately 30% and less than 4% performance on Hopper-M-R and Halfcheetah-M-R, respectively. Due to the attention and SSM mechanisms, DT and DM models conspicuously exhibit a higher upper bound compared to BC. Among them, our proposed DM shows the best performance across all datasets. This indicates that the specific architecture of DM enables it to extract more useful information from the inter-step and intra-step, leading to a strong performance across different context lengths.

4.5 The Effects of β in PSER

We have shown that the proposed PSER in DM enhances the robustness of the policy significantly in learning on suboptimal trajectories. Consequently, we endeavor to delve deeper into the impact of the policy self-evolution throughout the training process. During the training procedure, β K in PSER determines the upper bound of the policy self-evolution. When β K is set to 1 , the policy has the highest dependency on self-learned knowledge; conversely, if β K is set to 0 , the policy tends to completely lose its ability to self-evolve. We conduct experiments on DM variants, by removing the lower boundary β min and selecting β K from the set { 1 , 0 . 75 , 0 . 5 , 0 . 25 , 0 } . The results are depicted in

Hopper-M

Hopper-M-E

Hopper-M-R

100

110

90

90

80

Normalized Score

001060

100

70

06

BC

60

BC

80

★

DT

★

DT

50

+

DM (Ours)

50

+

DM(Ours)

70

40

BC

40

+

★

DM (Ours)

DT

60

30

20

50

20

20

40

60

80

100

120

20

40

60

80

100

120

20

40

60

80

100

120

Context length

Contextlength

Context length

Halfcheetah-M

Halfcheetah-M-E

Halfcheetah-M-R

44.0

90

40

Normalized Score

43.5

35

42.5-

80

30

70

BC

25

BC

42.0 -

★

DT

20

★

DT

41.5

60

+

DM (Ours)

15

DM (Ours)

41.0

+

BC

10

★

DT

50

40.5

+

DM (Ours)

5

40.0

20

40

60

80

100

120

40

20

40

60

80

100

120

0

20

40

60

80

100

120

Context length

Context length

Context length

Table 4. It can be observed that if we remove the self-evolution capability of DM, i.e., setting β K to 0 , there is a notable decline in the performance across both the M and M-R correlated datasets, with the poorest overall performance. When increasing the β to 0 . 5 , the model obtains the best performance since it reaches a good balance between the ground truth and the learned knowledge. To prevent the policy from excessively relying on its past knowledge, DM additionally uses a lower boundary β min (set to 0 . 5 ), achieving the best performance.

5 Conclusion

In this paper, we study the offline reinforcement learning from the perspectives of the architecture and the learning strategy. We have accordingly proposed Decision Mamba (DM), a multi-grained state space model tailored for RL tasks with a self-evolving policy learning strategy. DM enhances the policy robustness by adapting the mamba architecture to RL tasks by capturing the fine-grained and coarse-grained information. Meanwhile, the proposed learning strategy prevents the policy from overfitting the noisy labels with a progressive self-evolution regularization. Extensive experiments demonstrate that DM outperforms other baselines by approximately 4% on the mainstream offline RL benchmarks, showing its robustness and effectiveness.

Limitations and Future Directions . According to [15, 18], the mamba structure is more friendly to long sequences than the transformer structure, not only in terms of capturing historical information, but also in terms of the computational speed. Benefiting from the structure of SSM, the computational complexity of mamba is O ( n ) , while the computational complexity of attention score in transformer is O ( n 2 ) . Thus, the computational efficiency of mamba is higher. Although the exploration of computational efficiency is an exciting direction for future research, it is not within the main scope of this paper.

Acknowledgments and Disclosure of Funding

We would like to thank the reviewers for their constructive comments. This work is supported by Shenzhen College Stability Support Plan (Grant No.GXWD20220817144428005) and National Natural Science Foundation of China (Grant No. 62236003). Additionally, it is also supported by National Natural Science Foundation of China (Grant No. 62406092), and partially supported by Research on Efficient Exploration and Self-Evolution of APP Agents & Embodied Intelligent Cerebellum Control Model and Collaborative Feedback Training Project (Grant No. TC20240403047).

References

Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-making? arXiv preprint arXiv:2211.15657, 2022.

Anirudhan Badrinath, Yannis Flet-Berliac, Allen Nie, and Emma Brunskill. Waypoint transformer: Reinforcement learning via supervised learning with intermediate targets. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 78006-78027. Curran Associates, Inc., 2023.

Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemysław D˛ ebiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Chris Hesse, et al. Dota 2 with large scale deep reinforcement learning. arXiv preprint arXiv:1912.06680, 2019.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.

Yevgen Chebotar, Karol Hausman, Yao Lu, Ted Xiao, Dmitry Kalashnikov, Jake Varley, Alex Irpan, Benjamin Eysenbach, Ryan Julian, Chelsea Finn, et al. Actionable models: Unsupervised offline reinforcement learning of robotic skills. arXiv preprint arXiv:2104.07749, 2021.

Yevgen Chebotar, Quan Vuong, Karol Hausman, Fei Xia, Yao Lu, Alex Irpan, Aviral Kumar, Tianhe Yu, Alexander Herzog, Karl Pertsch, et al. Q-transformer: Scalable offline reinforcement learning via autoregressive q-functions. In Conference on Robot Learning, pages 3909-3928. PMLR, 2023.

Huayu Chen, Cheng Lu, Chengyang Ying, Hang Su, and Jun Zhu. Offline reinforcement learning via high-fidelity generative behavior modeling. arXiv preprint arXiv:2209.14548, 2022.

Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 15084-15097. Curran Associates, Inc., 2021.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, 2019.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929, 2020.

Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. Rvs: What is essential for offline rl via supervised learning? arXiv preprint arXiv:2112.10751, 2021.

A. Feinberg. Markov decision processes: Discrete stochastic dynamic programming (martin l. puterman). SIAM Rev., 38(4):689, 1996.

Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2020.

Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning, pages 2052-2062. PMLR, 2019.

Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2023.

Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 35971-35983. Curran Associates, Inc., 2022.

Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396, 2021.

Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces, 2022.

Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state space layers. Advances in neural information processing systems, 34:572-585, 2021.

Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces. Advances in Neural Information Processing Systems, 35:22982-22994, 2022.

Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 13521361. PMLR, 06-11 Aug 2017.

Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems, 29, 2016.

Shengchao Hu, Li Shen, Ya Zhang, and Dacheng Tao. Graph decision transformer. arXiv preprint arXiv:2303.03747, 2023.

Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.09991, 2022.

Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. In Advances in Neural Information Processing Systems, 2021.

Kyungyul Kim, ByeongMoon Ji, Doyoung Yoon, and Sangheum Hwang. Self-knowledge distillation with progressive refinement of targets. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pages 6567-6576, October 2021.

Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. In International Conference on Learning Representations, 2022.

Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. Advances in neural information processing systems, 32, 2019.

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1179-1191. Curran Associates, Inc., 2020.

Kunchang Li, Xinhao Li, Yi Wang, Yinan He, Yali Wang, Limin Wang, and Yu Qiao. Videomamba: State space model for efficient video understanding. arXiv preprint arXiv:2403.06977, 2024.

Xiaojie Li, Yibo Yang, Jianlong Wu, Bernard Ghanem, Liqiang Nie, and Min Zhang. Mamba-fscil: Dynamic adaptation with selective state space model for few-shot class-incremental learning. arXiv preprint arXiv:2407.06136, 2024.

Zaijing Li, Yuquan Xie, Rui Shao, Gongwei Chen, Dongmei Jiang, and Liqiang Nie. Optimus-1: Hybrid multimodal memory empowered agents excel in long-horizon tasks. In NeurIPS, 2024.

Yue Liu, Yunjie Tian, Yuzhong Zhao, Hongtian Yu, Lingxi Xie, Yaowei Wang, Qixiang Ye, and Yunfan Liu. Vmamba: Visual state space model, 2024.

Zuxin Liu, Zijian Guo, Yihang Yao, Zhepeng Cen, Wenhao Yu, Tingnan Zhang, and Ding Zhao. Constrained decision transformer for offline safe reinforcement learning. In International Conference on Machine Learning, pages 21611-21630. PMLR, 2023.

Yi Ma, Chenjun Xiao, Hebin Liang, and HAO Jianye. Rethinking decision transformer via hierarchical reinforcement learning. 2023.

Ajay Mandlekar, Fabio Ramos, Byron Boots, Silvio Savarese, Li Fei-Fei, Animesh Garg, and Dieter Fox. Iris: Implicit reinforcement without interaction at scale for learning control from offline robot manipulation data. In 2020 IEEE International Conference on Robotics and Automation (ICRA), pages 4414-4420. IEEE, 2020.

Keiran Paster, Sheila McIlraith, and Jimmy Ba. You can't count on luck: Why decision transformers and rvs fail in stochastic environments. Advances in neural information processing systems, 35:38966-38979, 2022.

Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage-weighted regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv:1910.00177, 2019.

Matthew E. Peters, Mark Neumann, Mohit Iyyer, Matt Gardner, Christopher Clark, Kenton Lee, and Luke Zettlemoyer. Deep contextualized word representations, 2018.

Maciej Pióro, Kamil Ciebiera, Krystian Król, Jan Ludziejewski, Michał Krutul, Jakub Krajewski, Szymon Antoniak, Piotr Miło´ s, Marek Cygan, and Sebastian Jaszczur. Moe-mamba: Efficient selective state space models with mixture of experts, 2024.

Yanyuan Qiao, Zheng Yu, Longteng Guo, Sihan Chen, Zijia Zhao, Mingzhen Sun, Qi Wu, and Jing Liu. Vl-mamba: Exploring state space models for multimodal learning. arXiv preprint arXiv:2403.13600, 2024.

Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pages 8748-8763. PMLR, 2021.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019.

Jinghuan Shang, Kumara Kahatapitiya, Xiang Li, and Michael S. Ryoo. Starformer: Transformer with state-action-reward representations for visual reinforcement learning. In Computer Vision - ECCV 2022, pages 462-479. Springer Nature Switzerland, 2022.

Ruizhe Shi, Yuyao Liu, Yanjie Ze, Simon Shaolei Du, and Huazhe Xu. Unleashing the power of pre-trained language models for offline reinforcement learning. In The Twelfth International Conference on Learning Representations, 2024.

Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. arXiv preprint arXiv:2208.04933, 2022.

Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pages 5026-5033, 2012.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. V on Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.

Oriol Vinyals, Igor Babuschkin, Wojciech M Czarnecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung, David H Choi, Richard Powell, Timo Ewalds, Petko Georgiev, et al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 575(7782):350-354, 2019.

Kerong Wang, Hanye Zhao, Xufang Luo, Kan Ren, Weinan Zhang, and Dongsheng Li. Bootstrapped transformer for offline reinforcement learning. Advances in Neural Information Processing Systems, 35:34748-34761, 2022.

Yiqi Wang, Mengdi Xu, Laixi Shi, and Yuejie Chi. A trajectory is worth three sentences: multimodal transformer for offline reinforcement learning. In Robin J. Evans and Ilya Shpitser, editors, Proceedings of the Thirty-Ninth Conference on Uncertainty in Artificial Intelligence, volume 216 of Proceedings of Machine Learning Research, pages 2226-2236. PMLR, 31 Jul-04 Aug 2023.

Yuanfu Wang, Chao Yang, Ying Wen, Yu Liu, and Yu Qiao. Critic-guided decision transformer for offline reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 15706-15714, 2024.

Yuanfu Wang, Chao Yang, Ying Wen, Yu Liu, and Yu Qiao. Critic-guided decision transformer for offline reinforcement learning. In Michael J. Wooldridge, Jennifer G. Dy, and Sriraam Natarajan, editors, Thirty-Eighth AAAI Conference on Artificial Intelligence, AAAI 2024, Thirty-Sixth Conference on Innovative Applications of Artificial Intelligence, IAAI 2024, Fourteenth Symposium on Educational Advances in Artificial Intelligence, EAAI 2014, February 20-27, 2024, Vancouver, Canada, pages 1570615714. AAAI Press, 2024.

Zihan Wang, Fanheng Kong, Shi Feng, Ming Wang, Han Zhao, Daling Wang, and Yifei Zhang. Is mamba effective for time series forecasting?, 2024.

Ziyu Wang, Alexander Novikov, Konrad Zolna, Josh S Merel, Jost Tobias Springenberg, Scott E Reed, Bobak Shahriari, Noah Siegel, Caglar Gulcehre, Nicolas Heess, et al. Critic regularized regression. Advances in Neural Information Processing Systems, 33:7768-7778, 2020.

Hua Wei, Deheng Ye, Zhao Liu, Hao Wu, Bo Yuan, Qiang Fu, Wei Yang, and Zhenhui Li. Boosting offline reinforcement learning with residual generative modeling. arXiv preprint arXiv:2106.10411, 2021.

Hua Wei, Deheng Ye, Zhao Liu, Hao Wu, Bo Yuan, Qiang Fu, Wei Yang, and Zhenhui Li. Boosting offline reinforcement learning with residual generative modeling. In Zhi-Hua Zhou, editor, Proceedings of the Thirtieth International Joint Conference on Artificial Intelligence, IJCAI 2021, Virtual Event / Montreal, Canada, 19-27 August 2021, pages 3574-3580. ijcai.org, 2021.

Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.11361, 2019.

Yueh-Hua Wu, Xiaolong Wang, and Masashi Hamaya. Elastic decision transformer. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.

Taku Yamagata, Ahmed Khalil, and Raúl Santos-Rodríguez. Q-learning decision transformer: leveraging dynamic programming for conditional sequence modelling in offline rl. In Proceedings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org, 2023.

Sherry Yang, Dale Schuurmans, Pieter Abbeel, and Ofir Nachum. Dichotomy of control: Separating what you can control from what you cannot. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.

Han Zhao, Min Zhang, Wei Zhao, Pengxiang Ding, Siteng Huang, and Donglin Wang. Cobra: Extending mamba to multi-modal large language model for efficient inference, 2024.

Qinqing Zheng, Mikael Henaff, Brandon Amos, and Aditya Grover. Semi-supervised offline reinforcement learning with action-free trajectories. In International conference on machine learning, pages 42339-42362. PMLR, 2023.

Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang. Vision mamba: Efficient visual representation learning with bidirectional state space model, 2024.

Tianchen Zhu, Yue Qiu, Haoyi Zhou, and Jianxin Li. Towards long-delayed sparsity: Learning a better transformer through reward redistribution. In Edith Elkind, editor, Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence, IJCAI-23, pages 4693-4701. International Joint Conferences on Artificial Intelligence Organization, 8 2023. Main Track.

Zhengbang Zhu, Hanye Zhao, Haoran He, Yichao Zhong, Shenyu Zhang, Yong Yu, and Weinan Zhang. Diffusion models for reinforcement learning: A survey. arXiv preprint arXiv:2311.01223, 2023.

A Dataset and Implementation Details

A.1 Dataset Details.

We conduct experiments on five tasks of Mujoco and Antmaze [47] including Halfcheetah, Hopper, Walker, Ant and Antmaze, as illustrated in Figure 4. Note, all datasets we used is the v 2 version. In these tasks, there are totally 5 different datasets which are described below:

Medium: A 'medium' policy is trained by using the Soft Actor-Critic [21] with earlystopping the training, and generate 1 million timesteps, achieving about one-third the score of an expert policy.

Medium-Expert: 1 million timesteps generated by the medium policy concatenated with 1 million timesteps generated by an expert policy (a fine-tuned RL policy).

Medium-Replay: It involves recording all samples in the replay buffer observed during training until the policy achieves a 'medium' level of performance.

Umaze: It contains the trajectories where the ant to reach a specific goal from a fixed start location.

Umaze-diverse: Different from Umaze, it is a more difficult dataset where the start position is also random.

(a) Halfcheetah

(b) Hopper

(c) Walker

(d) Ant

A.2 Implement Details

For all our experiments, we utilized the default hyperparameter settings and conducted 100,000 training iterations or gradient steps. We implement our method with the official repository of the huggingface.The shared hyperparameters are set to the same as those of LaMo, including the batch size, learning rate, overall training steps, and weight decay. The setting of other specific hyperparameters, including β K , β min are presented in Table 5. The experiments are conducted on an 8*4090-24G platform, and we run each experiment with four different seeds to ensure its reliability.

A.3 Code base

The code bases employed for our evaluations are detailed below.

BC: https://github.com/kzl/decision-transformer

DT: https://github.com/kzl/decision-transformer

EDT: https://github.com/kristery/Elastic-DT

LaMo: https://github.com/srzer/LaMo-2023

B Details of Baselines

We compare our proposed Decision Mamba with previous strong baselines as follows:

Behavioral Cloning (BC): it is a representative method of imitation learning. The states and actions are collected as the training data first. Then the agent uses a classifier or regressor to replicate the trajectory when encountering the same state.

Conservative Q-Learning (CQL) [29]: it encourages policies that are less likely to choose actions with high Q-value estimates that are uncertain or unreliable, thus expecting to address overestimation bias.

Decision Transformer (DT) [8]: it flats the trajectory sequence and use conditional sequence modeling method to autoregressively predict actions.

RvS [11]: it uses the goal or reward as the condition to realize the behavior cloning. We use the reward-conditioned BC as comparison.

StARTransformer (StAR) [44]: it extracts the image state patches by self-attending mechanism, then combining the features with the whole sequence.

Graph Decision Transformer (GDT) [23]: it adopts the sequence modeling method, and models the input sequence into a causal graph to capture relationships among states, actions, and return-to-gos.

WaypointTransformer (WT) [2]:it integrates intermediate targets and proxy rewards as guidance to steer a policy to desirable outcomes.

Elastic Decision Transformer (EDT) [60]: it estimates the highest achievable value given a certain history, and inputs the traversed trajectory with a variable length to learn the stitching trajectories.

Language Models for Motion Control (LaMo) [45]: it adopts the pretrained GPT2 [43] model as the backbone, and use the additional NLP corpus to co-training the policy via the parameter-efficiently LoRA method.

C More Comparison

For extensive comparison, we compare DM with more baselines, including the diffusion-based model: Diffuser [24], Decision Diffuser (DD) [1])and more complex approaches: Trajectory Transformer (TT) [25], Critic-Guided Decision Transformer (CGDT) [53]).

As illustrated in Table 6, it can be observed DM still has the strongest overall performance, although it did not achieve the best results on some datasets. Diffusion-based models synthesize optimal trajectories from a generative perspective, showing a significant advantage on replay datasets. Although CGDT performs well, it requires complex additional training, namely Critic training, which increase convergence difficulty. Overall, DM shows superiority on average.

D The Result on the Distribution of Returns

We compare the ability of policy to understand return-to-go tokens by varying the desired target return over a wide range, especially in the out-of-distribution range. As illustrated in Figure 5, we can observed in the seen target, i.e., on the left side of the yellow dashed line, the expected target returns and the true observed returns are highly correlated. However, when it comes to the out-of-distribution target, the score of DM is consistently higher than those of DT. Among them, due to the extreme difficulty of the HalfCheetah dataset, the performance of DM under the OOD target is only slightly surpassing DT. Conversely, on the other two datasets, DM exhibits strong robustness to the OOD target, significantly outperforming DT. The experimental result has illustrated DM has a strong robustness.

HalfCheetah

Hopper

100

S

e

N

50

ON

0

25

50

75

100

125

25

50

75

100

125

25

50

75

100

125

Target Return (Normalized)

TargetReturn(Normalized)

TargetReturn(Normalized)

DM

DT

Oracle

Best Trajectory in Dataset

E Visualization of Action Distribution

We visualize the action distribution of learned policy. Specifically, we use the policies trained on different level of noisy data to predict the next action of the same trajectory, and visualize the hidden layer of the predicted action. As shown in Figure 6, the distribution obtained by DM is more concentrated. This indicates that even if the noise level in the training data varies, DM can still learn an approximate distribution, demonstrating its strong robustness.

DT

DM

DT

DM

1.5

1.0

medium-expert

medium-replay

unpeu

1.0

medium-replay

mediutm-expert

0.5

1.0

1.0

medium-replay

wn(piku

medium-expert

0.5

5'0

0.5

0.0

0.0

0.0

0.0

0.5

-0.5

0.5

0.5

-1.0

-1.0

1.5

1.0

0.5

0.0

1.0

1.5

2.0

1.0

medium

badxa-wngerw

medium-replay

-1.5

1.0

S°0

0.0

0.5

1.0

-1.0

0.5

0.0

0.5

1.0

1.5

(a)Halfcheetah

(b) Hopper

F Impact/Safegard Statements

Impact Statement In this study, we propose DM, which effectively extracts historical information, fuses multi-grained information to predict action, and enhances the effectiveness of conditional sequence modeling for offline RL tasks. In addition, we introduce a self-evolving policy learning strategy to effectively prevent the policy from overfitting the noisy trajectories, and further enhance the robustness of the policy. To this end, this technology is expected to advance the offline RL agent which can assist human beings in working under the dangerous circumstances. There are many potential societal consequences of developing advanced RL algorithms, none which we feel must be specifically highlighted here.

Safegrad Statement In the paper, we have taken rigorous steps to ensure the responsible release of our new offline RL algorithm and any associated models or data. Given the potential for misuse or

dual-use of such technology, we have implemented several safeguards to mitigate these risks. Use of our algorithms is subject to the CC BY-NC-SA 4.0 agreement. In addition, we are committed to ongoing monitoring and evaluation of our models' usage to identify any emerging risks or patterns of misuse. We will take prompt action if we detect any unauthorized or inappropriate use of our technology. Finally, we are open to working with regulators, researchers, and industry partners to further refine our safeguards and ensure the safe and ethical use of our offline RL algorithm.

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: The contributions are summarized clearly and accurately in the introduction. More details please refer to Section 1.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it 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: We discuss the limitations of our work in the conclusion. More details please refer to Section 5.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section 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 on 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 a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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 a complete (and correct) proof?

Answer: [NA]

Justification: This work does not involve hypotheses and proofs of theory. In addition, all formulas are labeled and cross-referenced normally.

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

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 a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: All information about reproducing the main experimental results are presented in the Section 4.1 and Appendix A.2

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to 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 is a specific model and empirical evaluation, it 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 all submissions 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 to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case 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 to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: We have released our code in https://github.com/aopolin-lv/ DecisionMamba .

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/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on 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. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: Details about all experimental settings are presented in the Section 4.1.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a 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, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We conduct experiments with four different random seeds and report the mean and variance of each result.

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, confidence intervals, or statistical significance tests, at 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 run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

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 a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide details about compute resources in Appendix A.2.

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 as 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 or failed experiments that didn'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 https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We strictly adhere to 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 circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We have stated the societal impact in the Appendix F.

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 or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness 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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as 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 or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe 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)?

Answer: [Yes]

Justification: We have stated the safeguard in the Appendix F.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We cite the original paper and provide URLs of the code we used.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset'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: We do 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/model 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 is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. 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: This 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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This 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.

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 NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a novel application of state space model (SSM) into the offline RL problem. Prior to this, transformer-based architectures were used heavily, but the authors claim that they did not process “historical information,” which is often an important requirement in real world scenarios. In this work, a Decision Mamba encoder is composed of three stages. First, it embeds its trajectory using three types of MLPs along with time embeddings. Second, coarse-grained SSM directly uses hidden states from a structured SSM for understand sequential dependency. Thirdly, fine-grained use an 1D-convolutional layer for encompassing past few embeddings to precisely understand context. For training, this work proposes a progressive self-evolution regularization (PSER). In experiments, DM archives very promising scores in various offline RL MuJoCo tasks.\n\nSTRENGTHS:\nSSMs are a new class of models that is reported to be successful for various temporal problems; this reviewer believe they need to be tested for various subfields of ML that deal with sequential tasks (presumably equipped with the Markov assumption).  Since SSMs always produce linear-time sequential representation of their states, I think it was reasonable for authors to consider this characteristic as an advantage to solve decision problems. Notably, this paper have following strengths\n- This is a good application of SSM into the offline RL tasks. \n- The proposed architecture was straightforward to understand\n- The experimental results are evaluted in various environments and the proposed model achieved solid numbers.\n\nWEAKNESSES:\n- I find the lack of detailed reasoning of the manuscript why Conv1D module brings “fine-grained” control.\n- From my perspective, the proposed PSER is related to PPO and DPO methods in terms of regularization. It would be great there was similar theoretical justification (formally or informally) for this proposed loss term.\n\nQUESTIONS:\nI do not see tendency in Fig. 3. Is there results with very few context lengths? eg. {1, 5, 10}\n\nLIMITATIONS:\nThe authors authors adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-20"}

comments: SUMMARY:
This paper tackles the sequential decision-making problem in an offline RL setting. The authors propose Decision Mamba (DM), an extension of Mamba to adapt to the problem. There are 3 main technical contributions: (a) DM architecture (a mix of fine-grained and coarse-grained SSM modules), (b) progressive self-evolution regularization (PSER) for label robustness, and (c) self-supervised loss (predicting states and rewards as well as action labels). The experiments show that DM outperforms baselines, including transformer-based approaches.

STRENGTHS:
1. Applying Mamba-like models to the offline RL problem sounds natural and an interesting direction.
1. Empirical results look nice. DM significantly beats baselines. 
1. The paper is clearly written and easy to follow.

WEAKNESSES:
1. It is somewhat unclear why that design choice on DM was made. 
1. The technical contributions (a,b,c) would be independent and not tightly connected. I mean, (b) and (c) are possibly also effective to other methods like decision transformer.

QUESTIONS:
Regarding fine-grained and coarse-grained modeling, h^FG contains a conv layer, but h^CG doesn't. Why does the difference come out? What's the motivation?

Are PSER and self-supervised loss also effective against baselines such as decision transformers? Have you conducted any experiments on that?

LIMITATIONS:
Limitations are addressed.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-06-28

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a robust method based on Mamba for Offline Reinforcement Learning. Additionally, the paper using the knowledge of the past policy to refine the noisy labels as supervision avoids the model fitting the noisy trajectories. To better train the model, the paper introduce the inverse training goals which simultaneously predict action, state, and RTG for better robustness. Extensive experiments demonstrated the effectiveness of proposed model and training strategy.\n\nSTRENGTHS:\n1.\tIntroducing Mamba from the perspective of capturing historical information is interesting, and there are also some modifications to the Mamba architecture itself.\n\n2.\tUsing easy-to-implement methods to refine noise labels using past policies to avoid overfitting onto suboptimal paths.\n\n3.\tA novel training method that enhances the robustness of the model by incorporating the next stage's state and RTG into the prediction.\n\nWEAKNESSES:\n1.\tLack of explanation for extracting intra step relationships: The paper lacks a specific explanation and verification for how intra-step relationships are extracted, as mentioned in the contributions. \n\n2.\tLack of visual experiments on action changes affected by Formula16 and whether overfitting has truly been avoided: The paper only demonstrates the effectiveness of Formula 16 from the perspective of metrics. Can the effect of Formula 16 be visualized to prove that it preventS the model from overfitting to the suboptimal path and to show the impact it had on the actions compared with actions without Formula 16.\n\n3.\tLack of specific experimental settings for the three hyperparameters in equation 23: The paper mentions improving the robustness of the model by changing the training objectives during the training phase, but there is a lack of experiments on the specific coefficients of the losses for the three parts and the impact of different hyperparameter values.\n\nQUESTIONS:\n1.\tHave you considered applying the self-evolution regulation method and inverse training to other transformer-based methods to improve performance? Are these two methods universally applicable?\n\n2.\tThe formula of refined target at k-th is too simple which is a linear formula. Can it be made more complex? Additionally, the parameters of linear formulas are generally learned through learning. Can they also be learned in your work.\n\n3.\tAs mentioned in lines 183-187, there is a lack of explanation as to why Formula 16 does not introduce more historical information such as ak-2, ak-3, etc.\n\nLIMITATIONS:\nDespite not utilizing Mamba's computationally efficient capability to capture long-range dependencies, a window length was still set.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces Decision Mamba, an offline RL backbone based on State Space Models. It enhances policy robustness by integrating a fine-grained SSM module alongside the original coarse-grained SSM in Mamba. Meanwhile, it adopts a progressive self-evolution regularization to prevent the policy from overfitting the noisy labels. Extensive experiments against a broad baseline demonstrate its effectiveness.\n\n-----------Post Rebuttal-----------------\n\nMain concerns resolved. Update rating from BA to WC.\n\nSTRENGTHS:\n- Mamba have shown good performance across NLP and CV tasks and applying Mamba to offline reinforcement learning is an interesting direction.\n- The paper conducts a comprehensive comparison with SOTA offline RL methods.\n\nWEAKNESSES:\n- The advantages of DM appear somewhat constrained when considering variance. It appears that performance enhancements stem more from the PSER rather than the Mamba architecture. Combining PSER with other architectures, such as transformers, could potentially yield superior performance outcomes.\n- In the experiments investigating various context lengths, the range examined is rather restricted, and there appears to be no notable difference in performance as the context length increases in either DT or DM. It may be beneficial to conduct experiments using longer context lengths to assess potential performance variations.\n- It is not appropriate to assume that readers are very familiar with Mamba related terminology. Enhancing the methodology section with additional details to explain such terms would improve readability. For instance, the term \"group branch\" used in the ablation experiments lacks an earlier clear definition within the text.\n\nQUESTIONS:\n- Table 3 indicates a significant impact of PSER on performance. Given its compatibility with other proposed baselines, it would be worth testing whether adding that into them would improve their performance results.\n- The context length in Mujoco environments is relatively short, minimizing the need for extensive historical temporal information. Can you provide experimental data on Atari environments to assess DM's capability in handling longer contextual sequences?\n- More experiments on Multi-Grained SSM modules are requested to bolster the supporting evidence for their effectiveness.\n\nLIMITATIONS:\nUnable to surpass the performance of the behavior policy, the model performs poorly on a suboptimal dataset.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}