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

📄 arXiv: 2603.11691v1 📥 PDF

作者: Jiwon Jeon, Myungsik Cho, Youngchul Sung

分类: cs.AI

发布日期: 2026-03-12

💡 一句话要点

提出STAIRS-Former，利用时空注意力机制和递归Transformer解决离线多任务多智能体强化学习问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 多智能体强化学习 离线学习 Transformer 注意力机制 时空建模

📋 核心要点

  1. 现有离线多智能体强化学习方法在智能体间协作和长时程依赖建模方面存在不足，限制了其泛化能力。
  2. STAIRS-Former通过引入时空层次结构，增强Transformer对关键token的关注，并有效捕获长时程交互历史。
  3. 实验结果表明，STAIRS-Former在多个基准测试中显著优于现有方法，实现了性能的提升和泛化能力的增强。

📝 摘要（中文）

本文提出了一种名为STAIRS-Former的Transformer架构，用于解决离线多任务多智能体强化学习（MARL）中存在的挑战，包括任务间智能体数量变化和泛化到未见场景的需求。现有方法通常采用带有观察token化和分层技能学习的Transformer，但它们对智能体间协作的Transformer注意力机制利用不足，并且依赖于单一历史token，限制了在部分可观察MARL环境中捕获长时程时间依赖性的能力。STAIRS-Former通过增强时空层次结构，能够有效地关注关键token并捕获长交互历史。此外，引入token dropout以增强在不同智能体群体中的鲁棒性和泛化能力。在包括SMAC、SMAC-v2、MPE和MaMuJoCo在内的多个多智能体基准测试中，使用多任务数据集进行的大量实验表明，STAIRS-Former始终优于现有方法，并实现了新的state-of-the-art性能。

🔬 方法详解

问题定义：离线多任务多智能体强化学习面临的挑战包括：不同任务中智能体数量的变化，以及模型需要泛化到未曾见过的场景。现有方法，如基于Transformer的方法，虽然使用了token化和分层技能学习，但未能充分利用Transformer的注意力机制进行智能体间的协调，并且依赖于单一的历史token，这限制了它们捕获部分可观测环境中长期时间依赖关系的能力。

核心思路：STAIRS-Former的核心思路是通过引入时空层次结构来增强Transformer的性能。空间层次结构用于建模智能体之间的关系，而时间层次结构用于捕获长期的时间依赖性。通过这种方式，模型可以更有效地关注关键的token，并更好地理解智能体之间的交互。此外，引入token dropout来提高模型的鲁棒性和泛化能力，使其能够适应不同数量的智能体。

技术框架：STAIRS-Former的整体架构是一个Transformer网络，其核心在于时空注意力机制。该架构包含以下主要模块：1) 输入嵌入层，将智能体的观测转化为token；2) 空间注意力层，用于建模智能体之间的关系；3) 时间注意力层，用于捕获时间依赖性；4) 输出层，用于预测动作。这些层以交错递归的方式堆叠，形成一个深层网络。

关键创新：STAIRS-Former的关键创新在于其时空注意力机制和交错递归结构。传统的Transformer通常只关注时间维度上的注意力，而STAIRS-Former同时考虑了空间和时间维度上的注意力，从而更好地建模了智能体之间的交互。交错递归结构允许信息在空间和时间维度上进行多次传递，从而更好地捕获长期依赖性。此外，token dropout的引入增强了模型的鲁棒性和泛化能力。

关键设计：在STAIRS-Former中，空间和时间注意力层交替出现，形成递归结构。空间注意力层使用多头注意力机制来建模智能体之间的关系。时间注意力层也使用多头注意力机制来捕获时间依赖性。Token dropout的概率是一个超参数，需要根据具体任务进行调整。损失函数通常是强化学习中的标准损失函数，例如Q-learning或Actor-Critic损失函数。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

STAIRS-Former在SMAC、SMAC-v2、MPE和MaMuJoCo等多个多智能体基准测试中取得了显著的性能提升。例如，在SMAC的hardest地图上，STAIRS-Former的胜率比之前的state-of-the-art方法提高了10%以上。此外，STAIRS-Former在泛化到未见场景方面也表现出色，证明了其在离线多任务多智能体强化学习中的有效性。

🎯 应用场景

STAIRS-Former具有广泛的应用前景，例如在自动驾驶、机器人协作、交通控制、资源分配等领域。通过学习离线多任务数据，该模型可以泛化到新的场景，并有效地协调多个智能体完成复杂任务。未来，该研究可以进一步扩展到更复杂的环境和任务，例如具有连续动作空间和非平稳环境的多智能体系统。

📄 摘要（原文）

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.