Decision Mamba: Reinforcement Learning via Sequence Modeling with Selective State Spaces

Introduction

The paper explores the integration of the Mamba framework, a novel sequence modeling architecture, into the Decision Transformer model for reinforcement learning. Decision Transformer introduced a paradigm shift by replacing traditional reinforcement learning's reliance on value functions with a direct mapping of state-action-reward sequences to optimal actions using causal self-attention mechanisms.

While causal self-attention has been instrumental in achieving remarkable results, the paper investigates the potential performance improvements that could be gained by substituting the Mamba framework for the causal self-attention mechanism within the Decision Transformer model. Mamba introduces a data-dependent selection mechanism and efficient hardware-aware design to address the data- and time-invariant issues of prior state space models.

The paper hypothesizes that Mamba's architecture design could offer a novel way to encode and exploit the temporal dependencies and intricate patterns present in sequential decision-making tasks, potentially resulting in more accurate, robust, and nuanced decision-making outputs. The primary focus of the paper is on the possible performance gains that Mamba's architecture design might confer on the Decision Transformer's sequence modeling capabilities, rather than solely on efficiency considerations.

By examining the performance implications of integrating Mamba into Decision Transformer, the work aims to contribute to the broader discourse on how best to architect and train models for the nuanced demands of sequential decision-making, offering potential pathways to significant advancements in the field of reinforcement learning.

Preliminaries

This section provides an overview of offline reinforcement learning (RL) and state space models (SSMs) in deep learning.

In offline RL, also known as batch RL, the agent learns from a fixed dataset without interacting with the environment, which is more challenging than online RL. One approach is behavior cloning (BC), which learns the mapping from states to actions directly from the dataset. However, the offline dataset often lacks sufficient expert demonstrations, so return-conditioned BC can be used to incorporate reward information and target future returns.

Decision Transformer (DT) is a promising approach that reconceptualizes RL as a sequence modeling task, taking a sequence of returns-to-go, states, and actions as input and autoregressively learning to predict the optimal action for a given target return.

SSMs in deep learning are a class of sequence modeling frameworks based on linear ordinary differential equations. They map an input signal to an output signal via a latent state. The discretized SSM can be expressed as a sequence-to-sequence mapping, and imposing structural conditions on the state matrix gives rise to the structured state space model (S4). Mamba builds upon the S4 framework, introducing a data-dependent selection mechanism and a hardware-aware parallel algorithm to effectively capture contextual information, especially for long sequences, while maintaining computational efficiency.

Decision Mamba

Decision Mamba (DMamba) utilizes the Mamba block as a token-mixing module instead of the self-attention module of Transformer networks. The paper provides an overview of the DMamba architecture, describing it as a variant of Transformer-type networks.

The network architecture of DMamba adopts the basic Transformer-type neural network, similar to GPT architectures. The main module, the Mamba layer, consists of a token-mixing layer and a channel-mixing layer. These layers include layer normalization, residual connections, and corresponding mixing blocks.

The input trajectory is transformed into token embeddings by an embedding layer, which can be a linear layer or a 2D convolutional layer depending on the environment. These token embeddings are then processed by a stack of Mamba layers.

The Mamba block operates linear projections to obtain hidden states, which are transformed using a SiLU activation function. This produces the core of the block, the selective state-space model (SSM). The output is then generated by applying another linear and SiLU functions.

During training, the model is trained on a dataset of offline trajectories. Sub-trajectories of length 3K are sampled, and the model predicts the next action given the current state and return-to-go (RTG). The loss function depends on the environment, using either mean squared error for continuous actions or cross-entropy loss for discrete actions.

In the inference phase, the model receives a current trajectory and generates an action. The next state and reward are then used to update the RTG for the previous timestep.

Experiments

The paper compares the effectiveness of the DMamba sequence modeling framework for reinforcement learning (RL) against various DT variants, including the original DT, Decision S4 (DS4), and Decision ConvFormer (DC). The experiments were conducted on continuous OpenAI Gym tasks and discrete Atari control tasks.

For the OpenAI Gym tasks, the results show that DMamba performs comparably to the other methods across various medium-quality and medium-expert datasets. An ablation study on the channel-mixing layers in DMamba indicates that the core Mamba block is sufficient for RL sequence modeling tasks within the given scope.

In the Atari domain, which poses challenges due to long-term credit assignment, the paper examines the effect of the context length K in DMamba. The results suggest that longer context lengths can improve performance on Breakout but degrade it on Qbert, potentially indicating limitations in the Mamba block's selection mechanism for certain tasks.

Overall, the paper provides a comprehensive evaluation of DMamba against state-of-the-art sequence modeling approaches for RL, highlighting its competitive performance and the role of the core Mamba block and context length in different task domains.

Conclusion and Discussion

The paper examines the capabilities of the recently proposed Mamba model for sequence modeling in reinforcement learning (RL). It introduces the Decision Mamba (DMamba) model, which incorporates the Mamba block based on the selective SSM into a DT-type neural network architecture. The empirical study shows that DMamba is competitive with existing DT-type models, suggesting the effectiveness of Mamba for RL tasks.

However, the paper notes that merely applying the Mamba block to DT-type networks does not enhance efficiency due to the numerous interactions between CPUs and GPUs for the RL tasks considered. The paper adheres to the DT experimental setup for fair comparison but suggests reconsidering the implementation to leverage Mamba's advantages efficiently.

Another limitation is the absence of a hyperparameter search and an analysis of how to use the Mamba block more effectively to reflect the data structure of RL tasks. The paper proposes adapting the network architecture to better suit the RL data structure or preprocessing trajectory datasets into a format more compatible with Mamba. Applying Mamba to non-Markov decision processes is also suggested as a future direction.

The paper concludes that a comprehensive study with an exploration of improved Transformer-type network architectures for RL will be released elsewhere.

Appendix A Implementation Details of DMamba

The provided section describes the implementation details of the Mamba layer and the hyperparameters used for the experiments on the OpenAI Gym and Atari domains.

For the Mamba layer, the key aspects are:

• It is implemented using the Mamba module, as shown in the provided pseudo-code.
• The Mamba layer is incorporated into a Block module that also includes layer normalization and an MLP-based channel transformation.

For the OpenAI Gym tasks:

• The hyperparameters are adopted from previous studies, with some adaptations for each dataset.
• Key hyperparameters include the number of layers, batch size, context length, return-to-go conditioning, dropout, nonlinearity function, gradient norm clipping, and learning rate decay.

For the Atari domain:

• The hyperparameters also follow previous work, with specific values for different Atari games.
• Hyperparameters include the number of layers, embedding dimension, batch size, context length, return-to-go conditioning, nonlinearity functions, maximum epochs, dropout, learning rate, and learning rate decay.

The full set of code for the experiments is available in the provided GitHub repository.

Appendix B Task Scores and Normalization

The provided text describes the formula used to calculate the normalized scores for each domain. The normalized score is calculated as 100 times the raw score minus the random score, divided by the expert score minus the random score.

For OpenAI Gym tasks, the protocol from Fu et al. (2020) is used. For Atari games, the raw and baseline scores are provided in Table 7, which are used for normalization as described in Ye et al. (2021), Table 1. The column of DMamba in Table 7 shows the raw scores corresponding to Table 2 in the main text.