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.

错误:搜索内容不能为空,请输入英文关键词
错误:关键词超出字数限制,请精简
高级检索

Mamba in Mamba: Offline Reinforcement Learning via Sequence Modeling with Inner and Outer Selective State Spaces

  • Qiang Han,
  • Xiwen Wang,
  • Lifang Wang,
  • Wei Guo,
  • Kaixin Jin,
  • Xiaoqin Yu

摘要

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.