STAIRS-Former: Spatio-Temporal Attention with Interleaved Recursive Structure Transformer for Offline Multi-task Multi-agent Reinforcement Learning

TL;DR

STAIRS-Former is proposed, a transformer architecture augmented with spatial and temporal hierarchies that enables effective attention over critical tokens while capturing long interaction histories and introduces token dropout to enhance robustness and generalization across varying agent populations.

Abstract

Offline multi-agent reinforcement learning (MARL) with multi-task datasets is challenging due to varying numbers of agents across tasks and the need to generalize to unseen scenarios. Prior works employ transformers with observation tokenization and hierarchical skill learning to address these issues. However, they underutilize the transformer attention mechanism for inter-agent coordination and rely on a single history token, which limits their ability to capture long-horizon temporal dependencies in partially observable MARL settings. In this paper, we propose STAIRS-Former, a transformer architecture augmented with spatial and temporal hierarchies that enables effective attention over critical tokens while capturing long interaction histories. We further introduce token dropout to enhance robustness and generalization across varying agent populations. Extensive experiments on diverse multi-agent benchmarks, including SMAC, SMAC-v2, MPE, and MaMuJoCo, with multi-task datasets demonstrate that STAIRS-Former consistently outperforms prior methods and achieves new state-of-the-art performance.

Figures (19)

• Figure 1: Overall Proposed Q Structure
• Figure 2: Attention map on both seen and unseen task with basic transformer in HiSSD
• Figure 3: Overview of STAIRS-Former architecture.
• Figure 4: Attention map on both seen and unseen task with basic transformer in Ours.
• Figure 5: Temporal attention map in a SMAC 3m scenario. The attention maps from STAIRS-Former (top) and HiSSD (bottom) illustrate how attention shifts over time. Lighter-colored regions indicate eliminated agents. A detailed explanation of these heatmaps is provided in Appendix \ref{['appendix:attention_map']}