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

Authors: qi  lv
xiang  deng
gongwei  chen
michael yu wang
liqiang  nie
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 03:53:12

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 states, actions and return-to-gos (RTGs), (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 state-action-RTG 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.

Main paper: 
 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 state, action, and reward/return-to-go (RTG) 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 states, actions, and RTGs (SARs). A RL policy typically predicts the next action given the current state based on the RTG. Thus, this kind of finegrained intrinsic connection among SARs 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 Decision Mamba (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 SAR 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 WorkQuestion: 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. 
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.Recently, State Space Models (SSMs) show high potentials in various domains, including natural language processing [15, 16, 17, 19, 20, 40, 46] , computer vision [30, 31, 33, 41, 65, 63] 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 spatio-temporal 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 MethodQuestion: 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. 
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):EQUATIONEQUATIONWhen 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:EQUATIONwhere τ 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):EQUATIONy(t) = Ch(t).(EQUATIONDecision Mamba Encoder 0 0 0 1 1 1 2 2 2GroupCoarse-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) .EQUATIONEQUATIONwhere 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.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.Trajectory Embeddings Following the sequence modeling, we first use multilayer perceptrons (MLPs) to embed the SARs from the given trajectory τ = (s 0 , R 0 , a 0 , . . . , s T , R 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:EQUATIONEQUATIONwhere 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:EQUATIONwhere h CG i represents the coarse-grained hidden state; InterS3M denotes the coarse-grained SSM. h f i : (B, L, D) ← SiLU(Conv1d f i (h f ))13:A f i : (D, N) ← Parameter14:B f i : (B, L, N) ← Linear f B i (h f i )15:C f i : (B, L, N) ← Linear f C i (h f i )16:∆ f i : (B, L, D) ← log(1 + exp(Linear f ∆ i (h f i ) + Parameter f ∆ i ))17:A f i , B f i : (B, L, D, N) ← discretize(∆ f i A f i , B f i )18:h f i : (B, L, D) ← SSM(A f i , B f i , C f i )(h f i )19: end for 20:h CG i : (B, L, D) ← h CG ⊙ SiLU(z) 21: h FG i : (B, L, D) ← hFG ⊙ SiLU(z)22: /* fusion of multi-grained features */ 23:h MG i : (B, L, D) ← LayerNorm(h CG i + h FG i ) 24: hi : (B, L, D) ← Linear(h MG i + hi-1)25: end for 26: at : (B, L) ← MLP(h) 27: return at Fine-Grained SSM Further, to better discern the dependencies among SARs 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 SAR, as shown in the middle part and right part of Figure 1 .It can be formulated as follows:EQUATIONEQUATIONEQUATIONwhere 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:EQUATIONEQUATIONwhere 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. 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 combining the ground truth and the prediction from the learned policy itself. Let â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:EQUATIONwhere âk-1 ∼ π k-1 (•|s) and β is the trade-off weight.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:EQUATION)In comparison, the loss function with the refined target of Eq. ( 16) can be rewritten as:EQUATION)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:EQUATIONwhere ∇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 â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:EQUATION)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:EQUATIONWe adopt the MSE loss, and then the objective 21 is converted to:EQUATIONTo 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 τ = (s 0 , R 0 , a 0 , . . . , s t , R t ), Decision Mamba originally predicts the next action ât . Further, by incorporating the action a t to the original trajectory τ t-l:t , it is also predicts the next state ŝt and the next RTG Rt . Compared to the Eq. ( 3), the training objective of DM with the refined target can be written as follows:EQUATIONwhere the ã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.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?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 will release our code.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
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. Dataset BC CQL † DT RvS † StAR † GDT † WT † EDTFor a fair comparison, we first conduct experiments on datasets commonly adopted by mainstream approaches. The overall performance is presented in Table 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 DM performs 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.To 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.Avg. EQUATIONTo 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.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 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.Avg. EQUATION 
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 [18, 15] , 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. 
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. 
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. • 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). 
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. 
 
 Appendix: 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 v2 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• 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. 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. The code bases employed for our evaluations are detailed below.• BC: https://github.com/kzl/decision-transformer 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.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. EQUATIONWe 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.Figure 5 : The normalized scores of DT and DM when conditioned on the specified target returns.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. 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.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 ChecklistCorresponding Author. 38th Conference on Neural Information Processing Systems (NeurIPS2024). 


Decision: This paper proposes "Decision Mamba," a novel multi-grained state space model (SSM) with self-evolution regularization tailored for offline reinforcement learning (RL) tasks. The paper addresses three key challenges in offline RL: (1) insufficient utilization of historical temporal information, (2) overlooking intra-step relationships, and (3) overfitting suboptimal trajectories. The proposed method demonstrates robust performance improvements across various tasks by integrating fine-grained and coarse-grained SSM modules and employing a self-evolving policy learning strategy.
This paper received all positive recommendations (three weak accepts and one borderline accept). The major concern from reviewers was about the learning strategies, which was addressed in the rebuttal. The AC agrees with the reviewers’ assessments and recommends accepting the paper.

Reviews: review: paper_summary: paper_summary: This 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.

-----------Post Rebuttal-----------------

Main concerns resolved. Update rating from BA to WC.
strength_weakness: strength_weakness: - 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.
- 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.
- 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.
questions: questions: - 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.
- 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?
- More experiments on Multi-Grained SSM modules are requested to bolster the supporting evidence for their effectiveness.
limitations: limitations: Unable to surpass the performance of the behavior policy, the model performs poorly on a suboptimal dataset.
paper_summary: paper_summary: The 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.
strength_weakness: strength_weakness: 1.	Lack 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. 

2.	Lack 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.

3.	Lack 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.
questions: questions: 1.	Have 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?

2.	The 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.

3.	As 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.
limitations: limitations: Despite not utilizing Mamba's computationally efficient capability to capture long-range dependencies, a window length was still set.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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: Limitations are addressed.
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: - I find the lack of detailed reasoning of the manuscript why Conv1D module brings “fine-grained” control.
- 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.
questions: questions: I do not see tendency in Fig. 3. Is there results with very few context lengths? eg. {1, 5, 10}
limitations: limitations: The authors authors adequately addressed the limitations.

Comments: title: To authors
Kind reminder
Kind reminder
Kind reminder
Post-rebuttal comments
comment: We sincerely thank Reviewer vt4f for the positive feedback and the decision to raise the score. We will make these points raised in the review more clear in the next version of this paper. Thank you again.
Thank you for your response. My concerns have been addressed, and the additional experiments are convincing. I will raise my score.
Dear Reviewers, Area Chairs, Senior Area Chairs, and Program Chairs,

We are grateful to all the reviewers for the insightful comments and constructive suggestions and to the area chairs, senior area chairs and program chairs for the great service to the community. We have received the responses of Reviewer K8tX and 2Wf9, and hope our responses have adequately addressed the concerns of Reviewer vt4f and qjYZ. Please let us know if you have further questions. We take this as a great opportunity to improve our work and shall be grateful for any additional feedback you could give us.

Thanks,

Authors of Submission#13292
Thank Reviewer qjYZ again for the valuable comments. We have provided the response to each of the concerns raised in the review,  and we are eager to continue the conversation. As the interactive discussion window will close soon, we kindly invite the reviewer to read our response to see if there are any further questions.
Thank Reviewer vt4f again for the valuable comments. We have provided the response to each of the concerns raised in the review,  and we are eager to continue the conversation. As the interactive discussion window will close soon, we kindly invite the reviewer to read our response to see if there are any further questions.
Thank you for your encouraging feedback and for recognizing our efforts to address your concerns.

We are dedicated to continuously improving our work and would greatly appreciate any additional suggestions you may have to further enhance the quality of our work. If there are specific aspects where you believe further improvements could positively impact the score, we would be grateful for your advice.

Thank you once again for your time and valuable insights.
We greatly appreciate Reviewer K8tX for the insightful feedback and the decision to raise the score. In the next version, we will offer more background about mamba, and we will include additional experiments to further strengthen our work.
Thanks for the responses! I would like to keep the rating and looking forward to the improved revision.
I thank the authors for providing detailed responses to my comments. Most of my concerns are solved after reading the rebuttal materials. I also read the comments from other reviewers and decide to update my rating from Boardline Accept to Weak Accept.

Good luck to the authors for the final decision of this work.

Best wishes,