Mamba in Mamba: Offline Reinforcement Learning via Sequence Modeling with Inner and Outer Selective State Spaces
摘要
Decision Transformer demonstrates effective performance by performing sequence modeling with Transformers. However, there are issues with the high computational costs of Transformer and inaccurate model predictions. To address the above issues, this paper proposes an offline reinforcement learning via sequence modeling with inner and outer selective state spaces (Mamba in Mamba). It consists of three components: element dependency representation, inner mamba module, and outer mamba module. Firstly, the model constructs element dependency representations of adjacent time step trajectory elements based on the Markovian property in the input trajectory sequence. Then, the inner mamba module models these element dependency representations to capture short-term causal relationships within the trajectory. These short-term causal relationships are subsequently incorporated into the trajectory sequence as supplementary information to improve the accuracy of model predictions by utilizing the Markovian property. Finally, the outer mamba module models the trajectory sequence to capture long-term temporal relationships. By using Mamba with linear complexity instead of Transformer for trajectory sequence modeling, the computational cost is reduced. Experimental results show that Mamba in Mamba achieves high average normalized returns in the D4RL benchmark environment.