Mamba as Decision Maker: Exploring Multi-scale Sequence Modeling in Offline Reinforcement Learning

TL;DR

This paper introduces MambaDM, a novel sequence modeling approach for offline reinforcement learning that effectively captures multi-scale dependencies, achieving state-of-the-art results and providing insights into dataset scaling effects.

Contribution

We propose MambaDM with a new mixer module for multi-scale sequence modeling in RL, outperforming existing methods and analyzing the impact of dataset scaling on performance.

Findings

MambaDM achieves state-of-the-art performance on Atari and OpenAI Gym datasets.

Scaling dataset size improves performance significantly, while increasing model size does not.

The mixer module effectively captures both global and local sequence features.

Abstract

Sequential modeling has demonstrated remarkable capabilities in offline reinforcement learning (RL), with Decision Transformer (DT) being one of the most notable representatives, achieving significant success. However, RL trajectories possess unique properties to be distinguished from the conventional sequence (e.g., text or audio): (1) local correlation, where the next states in RL are theoretically determined solely by current states and actions based on the Markov Decision Process (MDP), and (2) global correlation, where each step's features are related to long-term historical information due to the time-continuous nature of trajectories. In this paper, we propose a novel action sequence predictor, named Mamba Decision Maker (MambaDM), where Mamba is expected to be a promising alternative for sequence modeling paradigms, owing to its efficient modeling of multi-scale dependencies. In particular, we introduce a novel mixer module that proficiently extracts and integrates both global and local features of the input sequence, effectively capturing interrelationships in RL datasets. Extensive experiments demonstrate that MambaDM achieves state-of-the-art performance in Atari and OpenAI Gym datasets. Furthermore, we empirically investigate the scaling laws of MambaDM, finding that increasing model size does not bring performance improvement, but scaling the dataset amount by 2x for MambaDM can obtain up to 33.7% score improvement on Atari dataset. This paper delves into the sequence modeling capabilities of MambaDM in the RL domain, paving the way for future advancements in robust and efficient decision-making systems.