STAIRS-Former: Spatio-Temporal Attention with Interleaved Recursive Structure Transformer for Offline Multi-task Multi-agent Reinforcement Learning
Organizations: School of Electrical Engineering Korea Advanced Institute of Science and Technology (KAIST) Daejeon 34141, Republic of Korea
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 & tables
| Tasks | Expert | Medium | ||||||
| UPDeT-m | ODIS | HiSSD | STAIRS (Ours) | UPDeT-m | ODIS | HiSSD | STAIRS (Ours) | |
| Source Tasks | ||||||||
| 3m | ||||||||
| 5m6m | ||||||||
| 9m10m | ||||||||
| Unseen Tasks | ||||||||
| Tasks | Expert | Medium | ||||||
| UPDeT-m | ODIS | HiSSD | STAIRS (Ours) | UPDeT-m | ODIS | HiSSD | STAIRS (Ours) | |
| Source Tasks | ||||||||
| 2s3z | ||||||||
| 2s4z | ||||||||
| 3s5z | ||||||||
| Unseen Tasks | ||||||||
| Tasks | UPDeT-m | ODIS | HiSSD | STAIRS (Ours) | |
| Seen | Marine-Hard | 21.2 | 47.9 | 64.6 | 79.0 |
| Marine-Easy | 44.3 | 59.3 | 83.9 | 91.2 | |
| Stalker-Zealot | 20.3 | 34.8 | 45.9 | 63.4 | |
| Mean | 28.6 | 47.3 | 64.8 | 77.9 | |
| Unseen | Marine-Hard | 21.1 | 31.8 | 52.7 | 57.0 |
| Marine-Easy | 29.9 | 42.5 | 79.8 | 86.7 |
| Tasks | UPDeT-m | ODIS | HiSSD | STAIRS (Ours) | |
| Seen | Terran | 10.9 | 15.2 | 24.9 | 31.3 |
| Protoss | 9.9 | 12.7 | 29.5 | 32.5 | |
| Zerg | 6.6 | 10.4 | 21.0 | 29.2 | |
| Mean | 9.1 | 12.7 | 25.1 | 31.0 | |
| Unseen | Terran | 7.0 | 13.2 | 25.1 | 32.3 |
| Protoss | 8.1 | 11.6 | 28.5 | 32.8 |
| Tasks | STAIRS | w/o Temporal | w/o Spatial | w/o Dropout | w/o ST | w/o STD | |
| Seen | Marine-Hard | 79.0 | 77.5 | 71.0 | 75.7 | 69.9 | 66.0 |
| Marine-Easy | 91.2 | 88.1 | 87.2 | 89.6 | 86.9 | 87.9 | |
| Stalker-Zealot | 63.4 | 63.0 | 59.0 | 62.6 | 50.3 | 55.1 | |
| Mean | 77.9 | 76.2 | 72.4 | 76.0 | 69.0 | 69.6 | |
| Unseen | Marine-Hard | 57.0 | 57.4 | 54.7 | 56.0 | 54.1 | 40.1 |
| Marine-Easy | 86.7 | 78.5 | 79.0 | 83.0 | 78.0 | 79.7 |
Appendix figures & tables33 assets
Supplementary material from the paper’s appendix.
Appendix
| Task | Quality | # Trajectories | Average return | Average win rate |
| 3m | expert | 2000 | 19.89 | 0.99 |
| medium | 2000 | 13.99 | 0.54 | |
| medium-expert | 4000 | 16.94 | 0.77 | |
| medium-replay | 3630 | N/A | N/A | |
| 5m | expert | 2000 | 19.94 | 0.99 |
| medium | 2000 | 17.33 | 0.74 |
| Task type | Task | Ally units | Enemy units | Properties |
| Source | 3m | 3 Marines | 3 Marines | homogeneous & symmetric |
| 5m | 5 Marines | 5 Marines | homogeneous & symmetric | |
| 10m | 10 Marines | 10 Marines | homogeneous & symmetric | |
| Unseen | 4m | 4 Marines | 4 Marines | homogeneous & symmetric |
| 6m | 6 Marines | 6 Marines | homogeneous & symmetric | |
| 7m | 7 Marines | 7 Marines | homogeneous & symmetric |
| Task type | Task | Ally units | Enemy units | Properties |
| Source | 3m | 3 Marines | 3 Marines | homogeneous & symmetric |
| 5m_vs_6m | 5 Marines | 6 Marines | homogeneous & asymmetric | |
| 9m_vs_10m | 9 Marines | 10 Marines | homogeneous & asymmetric | |
| Unseen | 4m | 4 Marines | 4 Marines | homogeneous & symmetric |
| 5m | 5 Marines | 5 Marines | homogeneous & symmetric | |
| 10m | 10 Marines | 10 Marines | homogeneous & symmetric |
| Task type | Task | Ally units | Enemy units | Properties |
| Source | 2s3z | 2 Stalkers, 3 Zealots | 2 Stalkers, 3 Zealots | heterogeneous & symmetric |
| 2s4z | 2 Stalkers, 4 Zealots | 2 Stalkers, 4 Zealots | heterogeneous & symmetric | |
| 3s5z | 3 Stalkers, 5 Zealots | 3 Stalkers, 5 Zealots | heterogeneous & symmetric | |
| Unseen | 1s3z | 1 Stalker, 3 Zealots | 1 Stalker, 3 Zealots | heterogeneous & symmetric |
| 1s4z | 1 Stalker, 4 Zealots | 1 Stalker, 4 Zealots | heterogeneous & symmetric | |
| 1s5z | 1 Stalker, 5 Zealots | 1 Stalker, 5 Zealots | heterogeneous & symmetric |
| Tasks | Expert | Medium | ||||||
| UPDeT-m | ODIS | HiSSD | STAIRS (Ours) | UPDeT-m | ODIS | HiSSD | STAIRS (Ours) | |
| Source Tasks | ||||||||
| 3m | ||||||||
| 5m | ||||||||
| 10m | ||||||||
| Unseen Tasks | ||||||||
| Task | Quality | # Trajectories | Average return | Average win rate |
| CN-2 | expert | 2000 | 1.0000 | 1.0000 |
| medium | 2000 | 0.6152 | 0.6152 | |
| CN-4 | expert | 2000 | 0.7173 | 0.7173 |
| medium | 2000 | 0.4273 | 0.4273 |
| Tasks | Expert | |||
| UPDeT-m | ODIS | HiSSD | STAIRS (Ours) | |
| Source Tasks | ||||
| CN-2 | ||||
| CN-4 | ||||
| Unseen Tasks | ||||
| CN-3 | ||||
| Hyper-parameter | Value |
| hidden layer dimension | 64 |
| attention dimension | 64 |
| 1.0 | |
| optimizer | Adam |
| learning rate | 0.0005 |
| Number of layers |
| Algorithms | UpDeT-m | ODIS | HiSSD | STAIRS |
| # Parameters | 79,095 | 138,573 | 679,335 | 220,023 |
| GPU Memory Usage (MiB) | 7,046 | 7,020 | 17,492 | 14,370 |
| Temporal Interval | |||||||
| Token Dropout Rate | |||||||
| 0.05 | 0.1 | 0.05 | 0.1 | 0 | 0.05 | 0.1 | |
| Average Performance | 65.9 | 65.3 | 64.5 | 66.2 | 63.6 | 66.2 | 65.6 |
| Task | Quality | # Trajectories | Average return | Average win rate |
| Terran 3_vs_3 | medium | 2000 | 11.6 | 0.44 |
| medium-replay | 2000 | N/A | N/A | |
| Terran 5_vs_5 | medium | 2000 | 13.09 | 0.4 |
| medium-replay | 2000 | N/A | N/A | |
| Terran 10_vs_10 | medium | 2000 | 12.58 | 0.42 |
| medium-replay | 2000 | N/A | N/A |
| Race | Unit Types (weights) | Start Pos. Dist. |
| Terran | marine (0.45), marauder (0.45), medivac (0.10) | surrounded_and_reflect ( ) |
| Protoss | stalker (0.45), zealot (0.45), colossus (0.10) | surrounded_and_reflect ( ) |
| Zerg | zergling (0.45), baneling (0.10), hydralisk (0.45) | surrounded_and_reflect ( ) |
| Tasks | Medium | Medium-replay | ||||||
| UPDeT-m | ODIS | HiSSD | STAIRS (Ours) | UPDeT-m | ODIS | HiSSD | STAIRS (Ours) | |
| Terran Source Tasks | ||||||||
| Terran 3 | ||||||||
| Terran 5 | ||||||||
| Terran 10 | ||||||||
| Terran Unseen Tasks | ||||||||
| Task type | Task | Number of Agents | Observation Space | Action Space |
| Source | (3,3) | 2 | [(8,), (8,)] | [(3,), (3,)] |
| (2,2,2) | 3 | [(6,), (10,), (8,)] | [(2,), (2,), (2,)] | |
| (1,1,1,1,1,1) | 6 | [(4,), (6,), (6,), (4,), (6,), (6,)] | [(1,)] 6 | |
| Unseen | (6) | 1 | [(12,)] | [(6,)] |
| (2,4) | 2 | [(6,), (10,)] | [(2,), (4,)] | |
| (1,2,3) | 3 | [(4,), (8,), (8,)] | [(1,), (2,), (3,)] |
| Task | Quality | # Trajectories | Average return |
| (3,3) | medium | 100 | 5043.32 |
| (2,2,2) | medium | 100 | 5074.4 |
| (1,1,1,1,1,1) | medium | 100 | 4076.55 |
| Tasks | HalfCheetah | ||
| UPDeT-BC | ODIS | STAIRS (Ours) | |
| Source Tasks | |||
| (3,3) | |||
| (2,2,2) | |||
| (1,1,1,1,1,1) | |||
| Unseen Tasks | |||
| Tasks | Expert | Medium | ||||||
| UPDeT-m | ODIS | HiSSD | STAIRS (Ours) | UPDeT-m | ODIS | HiSSD | STAIRS (Ours) | |
| Source Tasks | ||||||||
| 3m | ||||||||
| 5m6m | ||||||||
| 9m10m | ||||||||
| Unseen Tasks | ||||||||
| Tasks | Expert | Medium | ||||||
| UPDeT-m | ODIS | HiSSD | STAIRS (Ours) | UPDeT-m | ODIS | HiSSD | STAIRS (Ours) | |
| Source Tasks | ||||||||
| 2s3z | ||||||||
| 2s4z | ||||||||
| 3s5z | ||||||||
| Unseen Tasks | ||||||||
| Tasks | Expert | Medium | ||||||
| UPDeT-m | ODIS | HiSSD | STAIRS (Ours) | UPDeT-m | ODIS | HiSSD | STAIRS (Ours) | |
| Source Tasks | ||||||||
| 3m | ||||||||
| 5m | ||||||||
| 10m | ||||||||
| Unseen Tasks | ||||||||
| Task / Dataset | UpDeT-m | UpDeT-m + GRU | UpDeT-bc | UpDeT-bc + GRU | Ours |
| Marine-Hard | 21.1 | 20.7 | 46.6 | 49.8 | 62.5 |
| Stalker–Zealot | 15.3 | 16.7 | 43.3 | 42.9 | 51.7 |
| Marine-Easy | 34.2 | 31.6 | 82.1 | 85.7 | 88.1 |
Offline MARL
Offline multi-agent reinforcement learning (MARL) learns coordinated policies from fixed interaction datasets without further environment exploration.
See all 30 papers