Multi-Agent Reinforcement Learning with Selective State-Space Models

TL;DR

This paper introduces Multi-Agent Mamba, a scalable and efficient alternative to Transformer-based models in Multi-Agent Reinforcement Learning, matching performance while significantly improving computational efficiency and scalability to larger agent populations.

Contribution

The paper demonstrates that State-Space Models, specifically Mamba, can replace Transformers in MARL, maintaining performance and enhancing scalability.

Findings

MAM matches MAT's performance in multiple environments

MAM offers superior scalability to larger agent groups

SSMs can effectively replace Transformers in MARL

Abstract

The Transformer model has demonstrated success across a wide range of domains, including in Multi-Agent Reinforcement Learning (MARL) where the Multi-Agent Transformer (MAT) has emerged as a leading algorithm in the field. However, a significant drawback of Transformer models is their quadratic computational complexity relative to input size, making them computationally expensive when scaling to larger inputs. This limitation restricts MAT's scalability in environments with many agents. Recently, State-Space Models (SSMs) have gained attention due to their computational efficiency, but their application in MARL remains unexplored. In this work, we investigate the use of Mamba, a recent SSM, in MARL and assess whether it can match the performance of MAT while providing significant improvements in efficiency. We introduce a modified version of MAT that incorporates standard and bi-directional Mamba blocks, as well as a novel "cross-attention" Mamba block. Extensive testing shows that our Multi-Agent Mamba (MAM) matches the performance of MAT across multiple standard multi-agent environments, while offering superior scalability to larger agent scenarios. This is significant for the MARL community, because it indicates that SSMs could replace Transformers without compromising performance, whilst also supporting more effective scaling to higher numbers of agents. Our project page is available at https://sites.google.com/view/multi-agent-mamba .