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

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.

Strengths

SSMs 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 - This is a good application of SSM into the offline RL tasks. - The proposed architecture was straightforward to understand - The experimental results are evaluted in various environments and the proposed model achieved solid numbers.

Weaknesses

- 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

I do not see tendency in Fig. 3. Is there results with very few context lengths? eg. {1, 5, 10}

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors authors adequately addressed the limitations.

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?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

Limitations are addressed.

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.

Strengths

1. Introducing Mamba from the perspective of capturing historical information is interesting, and there are also some modifications to the Mamba architecture itself. 2. Using easy-to-implement methods to refine noise labels using past policies to avoid overfitting onto suboptimal paths. 3. A novel training method that enhances the robustness of the model by incorporating the next stage's state and RTG into the prediction.

Weaknesses

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

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.

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Despite not utilizing Mamba's computationally efficient capability to capture long-range dependencies, a window length was still set.

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.

Strengths

- Mamba have shown good performance across NLP and CV tasks and applying Mamba to offline reinforcement learning is an interesting direction. - The paper conducts a comprehensive comparison with SOTA offline RL methods.

Weaknesses

- 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

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

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

Unable to surpass the performance of the behavior policy, the model performs poorly on a suboptimal dataset.

Post-rebuttal comments

Thanks for the responses! I would like to keep the rating and looking forward to the improved revision.

Kind reminder

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

Code not available on URL

Dear authors the link you provided is not available on the GitHub repo. Can you please correct that?