STAIRS-Former: Spatio-Temporal Attention with Interleaved Recursive Structure TransFormer for Offline Mulit-task Multi-agent Reinforcement Learning
Authors: Jiwon Jeon
Myungsik Cho
Youngchul Sung
Venue: ICLR 2026 Poster (2026)
Keywords: Reinforcement Learning
multi-agent
multi-task
transformer
offline learning
Primary area: reinforcement learning

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.

Published as a conference paper at ICLR 2026
STATRS-FORMER: SPATIO-TEMPORAL ATTENTION WITH INTERLEAVED RECURSIVE STRUCTURE TRANSFORMER FOR OFFLINE MULTI-TASK MULTI- AGENT REINFORCEMENT LEARNING
Jiwon Jeon;" Myungsik Cho,
Youngchul Sung'
School of Electrical Engineering
Korea Advanced Institute of Science and Technology (KAIST)
Daejeon 34141, Republic of Korea {j [jiwon. jeon, ma. cho, ycsung) www.st.ac..........
ABSTRACT
Offline multi-agent reinforcement learning (MARL) with multi-task datasets is challenging due to varying numbers 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 a single history token, which limits their ability to capture long-horizon temporal dependencies in partially observable MARL settings. In this paper, propose T'AIRS-Formed 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 TTAIRS-Formen consistently outperforms prior methods and achieves new state-of-the-arr performance.
INTRODUCTION
Offline multi-agent reinforcement learning (MARL) has emerged as a promising approach to training many practical multi-agent systems such as connected vehicle and collaborative drones reduce costly and sometimes unsafe online interactions. Existing offline MARL works address overestimation bias, distributional shift, and out-of-distribution errors through conservative value estimation, hybrid optimization, or regularization strategies (Pan et al. 2022 Shao et al., 2023 Wang et al., 2023b Yang et al. 2021a). These advances are sigmificant but most results are limited to single-task settings. Real-world multi-agent applications demand agents that can master diverse skills, transfer knowledge across tasks, adapt the varying number of agents, and remain robust under heterogeneous conditions (Kaufmann et al.) 2023 Tang al. 2024). These requirements necessitate offline MARL methods that are not only stable but also generalizable to complex multi-task scenarios.
Multi-task (MT) RL (Caruana 1997 provides pathway to realize such generalization, but extending the conventional single-agent MT framework to MARL poses a unique challenge particularly due to the need to cover the varying number of agents. For example, for a collaborative task under the assumption of seven agents to still operate well even one, two three drones are missing- One possible solution to handle such multi-agent variability is to transformer-basee architectures with scalability, proposed for on-line transfer learning for MARL (Hu)
wants local drone policy trained
al. 202 Zhou al. (221). In this vein, ODIS and HiSSD (Zhang et al. 2023 Liu al. 2025) adopted UPDeT (Hu et al. 202T). transformer-based scalable architecture developed for on-line
"Equal contribution. Code available
Youngchul Sung the corresponding author.

---

Published as a conference paper at ICLR 2026
transfer learning, and leveraged hierarchical skill learning to extract transferable coordination patterns for offline MT-MARL, yielding promising results. While these works demonstrate the effectiveness of transformers for MT-MARL, they primarily use transformers to handle task-dependent variability in observation dimensions, rather than to fully exploit transformers' capacity for modeling sequential history and complex token relationships (Vaswani al. 2017b). As result, much of their potential to capture long-range dependencies and relational structures remains underutilized, as we shall see shortly.
que
To address this limitation, in this paper, we propose Spatio- Temporal Attention with Interleaved Recursive Structure which extends the transformer
Dat Preduct
Q Qatten
Sutteee
Transformer (STAIRS-Forme) architecture with spatial and temporal hierarchies to better model entity correlations and historical dependencies, and in- troduce an overall Q decomposition architecture for offline MT-MARL, shown in Fig. [ STAIRS-Former consists of three key components: 1) a spatial hierarchy that directs atten- tion toward the most relevant entities, 2) temporal hierarchy Figure 1: Overall Proposed Q Structure that strengthens the use of long-range past information, crucial in the partially-observabel setting of MARL, and 3) token dropout, which improves generalization across the varying number of agents.
Q Stairs
Q f
Stairs Formar
Stairs Former
(ol,a),A))
(o) .g
The contributions of this work are: 1) A novel transformer architecture for offline multi-agent reinforcement learning in multi-task scenarios, which selectively focuses attention across tokens to better capture critical information. 2) Introduction of spatial and temporal hierarchies within the transformer, highlighting their importance for handling varying agent populations and history dependence in multi-task settings. 3) Empirical evaluation on multi-task scenarios, demonstrating significant gains over baselines and setting new satte-of-th--as art performance.
2 RELATED WORKS
Offline MARL Offline MARL faces several key challenges, including coordination under partial observability, distributional shift, and convergence to sub-optimal. Recent works tackle these issues through conservative estimation regularization. For example, CFCQL (Shao et al. (023) introduces counterfactual regularization, OMAR (Pan et al. (002)) combines policy gradients with population- based search, OMIGA (Wang et al. 2023b) integrates value decomposition with offline policy learning, B3C (Kim & Sycara, 2025) incorporates behavior cloning with critic clipping, and MA- ICQ (Yang et 200T))) extends implicit Q-learning to multi-agent settings. While these methods enhance stability in offline training. they remain limited to single-task regimes and fail to address generalization and adaptability in dynamic multi-task settings.
Generalization and MT-MARL Although MT-RL (Hendawy et al. 2024 Cho et al.; Hendawy jtt al. 2023) and MARL (Jeon et al., 2022] Yang et al., 2020 Rashid et al.) 2020 Peng el al. 2021) have been extensively studied, their integration remains relatively underexplored. One research direction emphasizes architectural ficxibility. UPDeT (Hu et al. 202T) uses cansformer-basee value networks that adapt to dynamic agent populations and variable observation structures, providing a scalable inductive bias for cooperative MARL. Multi-Tass Multi-Agent Shared Layers (Wang) [et al. 2023a) shows that combining shared decision layers with task-specific perception modules enables concurrent training across tasks and supports transfer to unseen environments. DT2GG (Tian) al. 2023) further advances MT-MARL by decomposing complex tasks into transferable subtasks, reducing interference and improving cross-task generalization.
Beyond architecture advances, recent works explore representation learning and modularization. M3 (Meng et al.] (203)) introduces an offline pre-training framework that disentangles agent-invariant and agent-specific representations, improving few-shot and zero-shot transfer. HyGen (Zhang et al., 2024) combines offline multi-task data with limited online fine-tuning to extract generalizable skills, ODIS (Zhang et al. 2023) learns task-invariant coordination strategies, and HiSSD (Liu et al. 2025 decomposes cooperative knowledge into shared and task-specific components for structured transfer. Together, these methods highlight the promise of architectural design, modularization, and skill discovery for enhancing generalization in MARL. However, they mainly emphasize skill and representation transfer, without addressing how to attend to critical factors such as historical context or changing agent interactions, which crucial for robust policy learning under partial observability.
2

---

Published as a conference paper at ICLR 2026
3 PRELIMINARIES
MARI A cooperative partially-observable Markov game with N' agents is modeled as a Dec- POMDP T S, { Ω, O, P.n) (Oliehoek & Amato 2016). The state space is denoted by S, and each agent has its own action space A'. The observation space is represented by St, and the initial state is drawn from the distribution p : S + [o. 11. A discount factor y E [o, 1) controls how future rewards are valued. Although the agents interact with the same environment state E S, they each receive individual observations O; 52 which are produced by the observation function O SxN + S. At every timestep, each agent selects an action a' E A', and together these form the joint action to the transition dynamics s' ~ P(-(s,a), and all agents receive a shared reward r(s, For clarity, bold symbols are used for joint variables, such as o (o'.....,0" for observations and To = (T....., T ) for the collection of agent trajectories, i.c. Tt = (o at 10:--:700t-1). The goal is to optimize collection of decentralized policies π (a'la" )O) that maximize the expected cumulative reward, expressed as ( = E t=0 7're
(a 3....... ..... a²) The environment then updates its state according
In this paper, to capture generalizable behaviors across diverse tasks, we consider MT-MARL framework (Omidshatici et al. 2017). In this setting, we have a set of her where each task T, is modeled as an aforementioned Dec-POMDP: The goal is to learn a universal decentralized policy т over the training set Crrain that generalizes to an unseen test set Crest {Tj number of agents, state spaces, and action spaces differ across the tasks.
Offiine MT-MARL
training tasks Crrain {T))/21
Tj = S, (A)) Dg. P... 7).
Lise Here, the
To handle such heterogeneity, Hu et al. (202T) proposed UPDeTT, a transforreerbased unified pol- icy network (Vaswani et al., (017a). The key idea of UPDeT is that it decomposes each agent's observation o' according to characteristics, c.g., of is decomposed into three groups of entities: (I) own information ooun. (2) information about other agents { K,, and (3) environment information { i.e., o'=(o o Io..: O doa. "'n. ".,, Oen.K.) ). Then, each element in this decomposition is tokenized with a linear transform according to its characteristics, i.e., -wwn-Woww d'un +bown, ed wea.y-Woooo +bo and = "nn These tokens are ap- pended by a history token ei the appended overall tokens are fed to a transformer, and the local Q value for each discrete action is obtained from the output layer of the transformer. Note that the tokenizing matrices Wown, W0 and W the query, key and value generation matrices Wº, W and W in attention and the up and down projection matrices W'" and W7down in MLP of the transformer are independent of the context length, i.e., the number of tokens, once they are learned, and are common to all tokens. Hence, as more agents are added, the added elements in o² are just decomposed according to their characteristics, and all the previously-learne transformer parameters can be used again to cover this new setup with a different number of agents. Building on UP recent offline MT-MARL methods such as ODIS and HiSSD (Zhang al. 2023: Liu et al. 2025) train policies from fixed offline datasets Dj for each task T, without further environment interaction.
Limitations of UPDeT in MT-MARL To investigate how UPDeT, the central structure of previous offline MT-MARL algorithms, actually works, we conducted an experiment on the Marine- Easy task set in SMAC (Samvelyan et al.) 2019), which includes three training tasks ('3m', '5m', '10m'). We ran HiSSD, the current state-of the offline MT-MARL algorithm, and analyzed the resulting attention weights over the observation and history tokens of each agent for a seen task ('3m') and an unseen task (4m'). Fig. shows the attention maps of individual agents, where the rows correspond to queries and the columns to keys. It is seen that attention is distributed nearly uniformly across tokens in both seen and unseen tasks, failing to capture important entities in spatial domain. However, both HiSSD and ODIS use UPDeT-style transformers with only a single layer (depth 1) to model skills and actions. It limits the model's 's expressiveness, as a one-layer transformer cannot capture the diverse relations among agents, entities, and history. This explains the nearly uniform attention maps we observed in Fig. 2 Furthermore, the history token, which is important for partially-observable environments, is not heavily used. Note that in UPDeT, the attention output at the history position at time step t is basically given by a linear combination of of and history token input ehs, and this linear combination is fed to the MLP part of the transformer, yielding
ha
Wdown a (W"P (A + BioE))
(1)
3

---

Published as a conference paper at ICLR 2026
Agen
Ageenn
Agen
(a) Seen Task (3m)
(b) Unseen Task (4m)
Figure Atlention map on both seen and unseen task with basic transformer in HiSSD
for some matrices A, and Br, where a--) is the nonlinear activation of MLP. Hence, UPDeT's operation on the history token is simple RNN processing, which cannot incorporate long-term history information essential for partially-obeeraable environments. Thus, is seen that this information- lacking history token is not heavily attended in other output positions. In summary, the UPDeT structure does not possesses the capability long-term history preservation and does not fully exploit the strength of transformers, particularly their ability to model rich correlations between tokens.
This observation raises our central question: "How can we enhance the transformer architecture to capture richer correlations between entities while effectively leveraging historical information, for offline MT-MARL?" In the next section, we aim to provide one solution to this question.
METHOD
We propose Former), a new architecture designed for offline multi-task multi-agent reinforcement learning (MT-MARL). STAIRS-Former enhances both the modeling of inter-entity relationships and the utilization of historical information by integrating three key components:
Attention with Interleaved Recursive Structure Transformer (STAIRS-
Spatial Recursive Module: A recursive transformer that strengthens relational reasoning among entities within local observations.
hierarchical temporal structure with both step-wise and periodic updates, enabling agents to capture both short-term and long-term dependencies under partial observability.
Temporal Module:
Token-Dropout Mechanism: A stochastic regularization strategy that drops entity tokens during training, improving generalization to unseen tasks with different numbers of entities.
As illustrated in Fig. 3 STAIRS-Formin consists two trainable networks: a spatial- f(:0s) and GRU g(:: v). Together, they define the local Q-networks, which are then aggregated through the Qatten mixing network (Yang et al., 2020) which can adapt to varying number inputs, the following subsections, describe each component in detail.
4.1 SPATIAL RECURSIVE MODULE
In MARL, it is crucial to model diverse relationships among entities so that agents can prioritize the most relevant parts of their observations and generalize policies more effectively to unseen tasks. Prior methods, such as HiSSD, rely on shallow transformer layers that struggle to capture this diversity (see Fig. (b)). To address this, STAIRS-Former employs a recursive deep transformer, called Spatial-Forman Former, which refines relational reasoning through recursive steps for each layer.
Entity Embeddings For agent i, entity- observations given by o'=(o (0 O 1..1: o² 1m.1: where K₂ and Ke denote the numbers of other agents and environment entities. We embed entities as
e [eouns fou,1..... C'a, K. e.m,,1.... ....,,e'.. E.n.kk, Kxd
(2)
with e Coun =Wown oonn +Down, e'o,--Woon -Woo om,, aoa and e -Wen oi oen., +ben and parameters 0. (Wown 1own, Wra 10a Wen 15n).
4

---

Published as a conference paper at ICLR 2026
STAIRS-Former
Spatial-Former
Temporal-Focus Layer
Outout
Output
Spatial Former
Spatial Former
Spatial Former
Layer
UniToken
VMx
Add Norm & Dropout
ob8TH
obeyy++
obBeTw
History
Jewr 7y-1
low 2r2
Dropout
Layer
GRU
Riah
GRU
V1x
Attention
Layer
UniTaken
Jobs
EEemm
Eneerty
UniToken
Figure 3: Overview of STAIRS-Formn architecture.
Recursive Spatial Updates Let the Spatial-Former have M distinct layers. Each layer has weights 8: and is applied vi times with shared parameters for robust feature extraction with nominal by 1. Let z denote the recursive latent state at step j in layer l. For initialization = 0), the input the token sequence concatenated with history tokens (defined in 4.2 20 = [e', h², hh"..
At layer the recursive state is initialized as 2 0 (shape as 2 and then updated recursively using the previous state of together with the final state from the preceding layer, 2'
+ [(]] ++=-==00))
j= 0, =,....,,44 1.
The final state of layer I is obtained as 2² == z/ which is then passed to the next layer. Once all M layers are applied, the spatial representation is given by Asp 2M Per-agent action values are then obtained through an output head fo: Q(o',-) fo fo). This recursive design enables deeper relational reasoning while controlling parameter costs through weight sharing.
4.2 TEMPORAL MODULE
Partial observability is central challenge in MARL, each agent i only has access to local observations of rather than the global state 8t. Existing approaches, such as UPDeT, augment embeddings et with history token h; forming the input set [c), h However, these methods struggle to capture long-range dependencies (see Fig. 2 To address this limitation, we introduces a hierarchical temporal process that maintains two history states with different update frequencies.
Hierarchical Temporal Updates Each agent i maintains low level history h;-² updated every step, and a high level history h; hi updated every Tи steps by a GRU Chung et al. (2014) g(: v)).
At time f, the transformer input is the token set {ef, A/-1, A.H.H of length K, KK. From the Spatial module output Zsp. we read the history position and update
(9 A i.L. ;4), 6.
t= 0 mod THr otherwise.
h = 2pp 2,:1
h
(4)
h--1
Both histories are initialized to zero at t = O. The arrangement enables immediate responsiveness via 4 and long-range summarization via 1,
Temporal Feature Learning Entity tokens (spatial relational content) and history tokens (temporal context) play distinct roles; yet a single position-wiss FFN (or MLP) after attention tends to blur them. Note that the two-layer FFN after attention performs feature matching test with key vectors stored the first layer and vector reconstruction with the value vectors (stored in the second layer) with nonnegative key vector correlation with the input. To enable distinct feature extraction and test for spatial and temporal tokens, we attach two independent FFNs after each attention block inside the Spatial-Former: one specialized for spatial entity tokens and for history tokens capturing time evolution. Formally, let the attention output at recursive step in layer / be 2 = [ z ).ss. where z! oobs the updated entity tokens and x/ tts are the updated history tokens. Instead of sending both through single shared MLP, we apply two position-wise FFNs with disjoint parameters:
z/ j.oss
FFN odtn (1 (b))))
Ad hiiis FFN (2),)
(5)
5

---

Published as a conference paper at ICLR 2026
and concatenate form the post-FFN state 2/ - j.obes 2 jjiis This ensures that relational reasoning over entities and temporal abstraction through history tokens are refined along distinct pathways, en- couraging specialization while preventing interference between spatial and temporal representations.
4.3
MECHANISM
Generalization to unseen tasks is challenging because the number of entities varies across en- vironments according the number of agents and enemies. Although transformers can handle variable-length inputs, training is restricted to entity counts observed in the training set Ciain- As result, performance may drop on unseen tasks with new entity configurations. To reduce overfitting. STAIRS-Formed employs a token-dropost strategy.
During training, each entity embedding in e² (eowns e o,, K e I:K. is randomly dropped with probability Purop, except for:
(1) the agent's own entity w critical for stable learning. (2) both history tokens h and h..
and, when the policy head associates actions with per-entity outputs as in UPDeT, the entity token linked the dataset action to respect offline regularization.
This mechanism exposes the model to variable token lengths during training, improving robustness to unseen entity configurations by reducing overfitting to Crain-
4.4 TRAINING
We train STAIRS-Former with TD3+BC-styi objective (Fujimoto & Gu (021) adapted for discrete action spaces. The objective integrates temporal-differene (TD) learning with behavior cloning (BC) regularization, balancing value optimization with stability in the offline regime.
STAIRS-Former Loss
For each agent i, STAIRS-Formin outputs Q-value, Qi = Q(o o 0) given the observation and action sequences (oi... ). With each agent's individual Q-value Qi we adopt the Qatten mixing network (Yang et al., 2020) to obtain the global Q-value in MARL. Quo((Te, St, ag 0, o) from the set of individual Q-values ..!.... Q.", Here, let to, } denote the set of all parameters for STAIRS-Forma and & denote the parameters for mixing network. The target for TD learning is defined
=
+ Y max Qtee (Te+1; 8t+1, 0';0, O)
where 0, 6 are target parameters. The STAIRS-Former loss then jointly optimizes TD learning and BC regularization:
X
CSTAIRS(O, (
(Qeoc(Te. St, aci 0)
V1
Qobrr abre: 0)
(7)
BC loss
TD loss
where the first term fits TD targets the second encourages higher Q-values for dataset actions. The coefficient À controls the strength of policy regularization. Token-dropout ($ 4.3) applied during training to improve robustness, and target networks are updated at each target update interval.
In summary, STAIRS-Former integrates recur- sive spatial transformer for richer inter-entity rea- soning, a dual-timescale temporal module for both short and long horizons, and token dropout for robustness varying entity counts. Fig. shows STAIRS-Former results for the same setup as Fig. s 3).
seen that STAIRS-Formed consistently emphasizes critical entities and his- tory tokens, leading to more robust, generalizable policies across and unseen tasks. Additional visualizations the attention maps for the vari- tasks are provided in Appendix G
(a) Seen Task (3m)
(b) Unseen Task (4m)
Figure 4: Attention map on both seen and unseen task with basic transformer in Ours.

---

Published as a conference paper at ICLR 2026
Table
Comparison of average and per-task performances on the Marine-Harr task set across four dataset qualities. We report standast deviation, with the best shown in bold.
Tasks
Expert
Medium
UPDeT-m
ODIS
HiSSD
STAIRS (Ours)
UPDeT-m
ODIS
HiSSD
STAIRS (Ours)
Source Tasks
99.4
33.8
4.4
5m66m
6.9 199=111
43.1 + 28.2
72.5 10.7
70.6 10.5 =
19
22.5 10.0
35.6 9.5 68.1
50.0 12.5
10m
24.4 24.4
05.6 28.2
99.4
31.2
57.5
37.5-:21.5
86.9 7.5
4m
51.9 1: 15.6
88.1 8.7
100.0
1 0.0
97.5 1:
27.5 30.3
T1.3 18.0
78.1 21.1
89.4
13.9
91.9
83.1
100.0
0.0
80.6
94.4
100.0
10m
46.3 1: 19.3
43.1 1 1.311
98.8 2.8
100.0 1: 0.0
49.4
40.8
65.6 1.211
5.6 96.3
97.5
4.1
12m
16.3
7m8m
0.6 = 1.4
12.5 10.1
43.1
15.4
26.0 22.0
1.3 1.7
6.3
5.6
10.6
8.7
8m9m
4.4 +
9.4
49.4
14.8
1.3 1.7
10.6
14.4++9.
15.6
8.0
10m1 Im
8.8
29.1
S0.6
18.9
87.5 4.9
4.
15.2
46.3
61.3
18.2
10m 2m
0.0 0.0
0.0
11.3
13.6
6.6 7.5
0.0
0.0
1.3
1.3 1.7
2.6
0.0
Avg
26.5
39.9
69.7
68.4
20.3
37.5
49.5
$7.9
Mediam Expert
Medium Replay
Source Tasks
3m
32.3 1: 20.3
44.4 25.7
88.8 1 12.8
98.8 1: 1.7
40.0 i: 19.8
81.9 7.5
75.6 7.8
78.1 17.1
5móm
5.6
29.4 13.6
32.5 1: 22.6
57.5 13.9
0.0
11.9
11.9:111 13.0
2444 17.0
50.6 5.1
9m110m
10.6 = 14.8
55.6 31.7
69.1 35.8
94.1 4.1
0.6 0.6
15.0 + 13.3
45.0 13.0
78.1 + 16.1
Unseen Tusks
5m
72.5 = 35.3
31.6 68.8
93.8
14.0
100.0 0.0
64.4 23.2
55.6 21.6::1..
71.3 29.3
100.0
0.0
10m
49.4
71.3+
95.0
- 4.8
12.8
47.9
93.1 4.6
97.5
5.6
12m
20.0 1: 13.0
45.8
40.0
85.6
14.1
94.4
6.4
25.6
22.0
51.3 47.33
9..9
94.4
6.0
7m8m
0.6 = 1.4
11.9 15.1
40.0
20.4
15.0 4.1
0.6
6.1
12.
23.1
15.1
10.8 13.0
33.1
26.9
10m1 Im
3.8 1: 5.1
25.6
13.6.122 21.1
62.5
16.4
80.6
18.1
0.6 1.6114 1.4
12.5 15.5
33.8 11.6
66.9
11.2
10m1
13m 5m
0.0 1. 0.0
0.0
0.0
8.8 9.2
0.6
0.0 L. 0.0
1,3 2.8
6.3 11 1.4
4.4 1 4.7
Avg
5 EXPERIMENTS
5.1 BENCHMARKS
We evaluate the proposed method the offline MT-MARL setting primarily on task sets from the StarCraft Multi-Agent Challenge (SMAC) (Samvelyan et al. 2019). following the setup Zhang et al. (023)). Each task set is split into disjoint training and testing tasks to evaluate generalization to both and unseen tasks. We consider three task Marine- Easy, Marine-Hard, and Stalker-Zealot. Tasks within each set share the same unit types but differ in unit counts. Following D4RL (Fu et al. 2020), each task is associated with four datasets of varying quality: Expert, Medium, Medium-Expeer and Medium-Replar D All results report the mean and standard deviation of final performance over five random seeds. In addition to SMAC, evaluate our method on SMAC-v2 (Ellis et al. (003)) to test performance in more diverse and challenging scenarios; details are given in Appendix [ Results on other benchmarks, including the Multi-Agent Particle Environment (MPE) (Lowe et al.) 2017) and Multi-Agent MuJoCo
Additional details on benchmark construction are provided in Appendix B
(MaMuJoCo) (Peng et al. 2021) are provided in Appendix and Appendix| respectively.
5.2 BASELINES
We compared our method, STAIRS-Former, against the several offline MT-MARL approaches: 1) UPDeT-m: An offline variant of UPDeT (Hu et al. 2021), using a Qatten (Yang et al. 2020 mixing network trained with the CQL (Kumar et al. 2020) loss. 2) ODIS (Zhang et al. 2023): Discovers task-invariant coordination skills from offline multi task data and learns a coordination policy that selects skills under the CTDE paradigm generalize to unseen tasks. 3) HiSSD (Liu al. (025):
Uses hierarchical framework to jointly learn common cooperative skills and task-specific skills, enabling effective policy transfer and fine-grained action execution across tasks.

---

Published as a conference paper at ICLR 2026
Table 2: Comparison of average and per-task performances on the Balker-Zealde task set across four dataset qualities. We report mtanddst deviation, with the best shown in bold.
Tasks
Expert
Medium
UPDeT-m
ODIS
HiSDD
STAIRS (Ours)
UPDeT-m
ODIS
HESSD
STAIRS (Ours)
Source Tusks
263x
40.6 4: 33.9
56.9 29.3
90.6 7.7
95.6
5.2
27.5 15.5
44.4 16.1
46.9 15.5
56.9 1: 10.3
16.3 +: 14.6
52.5 # 18.8
80.0 = 10.5
77.5 + 11.6
21.
21.3
11.3
15.0 15.08.1
60.0 + 16.1
3652
23.8 4: 27.2
65.6 37.8
90.6 + S.8
87.5 + 10.6
13.1 +:
15.6 8.0
28. 20.1
52.5
Unseen Tasks
1832
25.6
20.3
23.1 + 26.5
82.5 25.4
78.1 12:7
39.4
37.0
31.9
36.2
16.3 16.1
38.8
1sf2
26.9 + 27.6
18.8 8.0
39.4 33.2
76.3
+ 21.0
20.6 21.
26.3
16.6
18.8
11.9
25.6 9.7
165a
10.6
16.9
10.6
18.8 23.0
55.6
+ 23.5
8.8 9.7
26.3
10.6
10.0 16
$1.9
10.5
36.3
363x
35.0
80.0
18.1
22.7
24.4
69.4
14.1
3644
32.5
+ 57.5
60.0 1000=550 35.3
88.5
0.35299.2 9.3
92.5
3.6
25.0 23.1
21.4 15.4
27.5 10.0
59.4
24.7
463d
11.9
+ 15.1
43.8 + 36.0
72.5
28.9
70.0
11.8
5.6
21.9
21.1
28.8 13.9
41.9
17.9
4647
10.6
33.8
26.8
17.5
8.1
465
2.5
2.6
82.5
19.8
16.3
19.9
53.1
18.9
4.4
8.1 6.1
1.9
11.3
7-8
Jshi
3.8
+ 6.1
26.3
+: 28.1
47.5
22.4
69.4 + 17.5
0.6
388 2.6
4.4 2.8
11.9
6.6
66.1
75.0
19.4
Tasks
Medium- Expert
Medium-Repha
Source Tisks
2632
34.4 23.4
60.3 + 21.6
76.9 20.2
92.5 + 10.3
388 1.1
10.6 23.8
5.6 6.1 22.5 15.8
20.6 + 10.0
27.5
3652
16.9 + 14.6
38.1 = 12.2
51.3 20.1
85.0 + 15.8
5.0 7.2
12.5 15.8
28.8 10.2
Unseen Tasks
183g
30.0 + 22.2
70.6 + 34.7
65.6
24.9 6.3
63.1
15.2
6.9 8.7
3.1 T.0
45.6
32.3
12.5 = 14.5
26.3
4.4
0.0
185x
13.8
19.4
82.9
0.0
36.3
265g
34.4 + 18.1
26.3
12:2
19.4
10.9
62.5
21.2
5.0
4A= 6.5
244
27.5
11.4
3337
30.6
+ 20.0
12.8 38.8
46.3
52.5
10.9
81.9
11.6
1.3 1.7
9.4 16.2
56.3 + 15.9
3647
29.4
75.0
16.4
95.6
1:
1.9
8.1
10.4 18.8
10.4
463
15.0
22.7
12.5 17.0
62.5
12.5
61.3
15.7
3.8
5-4
28.1
20.4
A64x
10.0
19.1
10.6 18.8
31.3
10.4
59.4
4: 14.8
0.6
5.0
T.8
13.1
8.8
15.0
2.6
465z.
2.5
4.1
11.9
1 13.1
11.9
4.1
53.8
1 21.7
1.3 1.7
1.3 2.8
5.0
4.7
3.8
466z
5.0 7.2
3.8
13.8
14.8
40.0
1 15.5
0.0 0.0
5.0
7.5
6.8
Avg
21.5
33.7
38.7
69.4
5.0
15.0
24.3
5.3 MAIN RESULTS
Evaluation SMAC Benchmark The results for the Marine-Hard and Stalker-Zealor task are presented in Tables and while the Marine-Easy results are provided in Appendix E due to space limits. Across all task sets, STAIRS-Forman demonstrates outstanding performance. and Stalker-Zealor it consistently achieves the best average results on both train and test tasks, with only a minor gap on the Expert dataset Marine-Hart These results show that STAIRS-Fore Former is not only effective in-distribution but also highly robust unseen tasks, highlighting its strong generalization ability. This robustness arises directly from proposed hierarchical spatial-temporal process with token dropout, which enables the model to capture richer dependencies across entities and leverage historical information more effectively, thereby maintaining robustness on unseen tasks.
Marine-Hard
Compared to the previous state of the art, HiSSD, the advantages of STAIRS- dataset quality. Best in bold, second-best underlined. Former become optimal datasets (Medium, Medium-Expert, and Mepium-Replay) in Marine- and Stalker-Zealot, STAIRS-Former achieves large gains-imprrovng average performance by 39.5%. 36.6%. and 40.5%, respectively. On the challenging Stalker-Zealot task set, which requires complex heterogeneous unit interac- tions, STAIRS- outperforms HiSSD by remarkable 48.6% average. Table 3 further illustrates STAIRS-Fomerr's superiority. It achieves the highest mean win rates on both seen tasks (77.9% vs. 64.8% for HiSSD) and unseen tasks (64.0% 54.7%), resulting in overall mean of 67.47 67.4% compared to HiSSD's 57.2%. These results high- lights that the our spatial temiooall reasoning, reinforced with token dropout, gives STAIRS-Foreer decisive advantage in both exploiting seen tasks and generalizing unseen scenarios.
Table Results on seen and unseen tasks, averaged over
clearer. On sub-
Marine Hard
21.2
47.9
64.6
79.0
Marine Easy
41.3
59.3
83.9
91.2
Seen
Stalkur-Zaaol
20.3
34,8
45.9
63.4
Murine-Huu
21.1
31.8
52.7
57.0
Unsween
Murine Fasy
Staker Zcalot
13.7
20.9
42,5
22.5
31.5
70.8
86.7
48.2
$2.2
8

---

Published as a conference paper at ICLR 2026
2114
Figure 5: Temporal attention map in 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 G.1
Evaluation on SMAC-v2 Benchmark Build- ing on the strong performance observed on SMAC benchmark, we further evaluate STAIRS- Former on the more challenging SMAC-v2 benchmark, which features increased stochastic- ity, and more diverse unit interactions. As shown in Table 4 TTAIRS-Formn consistently outper- forms prior methods across all task sets (Terran, Protoss, and Zerg). Compared to the previous state of the art, HiSSD, it improves average per- formance on seen and unseen tasks by 23.5% and 24.5%, respectively, achieving the highest overall average win rate of 30.3%. These results demon- strate that STAIRS-Former effectively extends its advantages to significantly more complex and stochastic environments. Full results are provided Table [8 Appendix
Table 4: Results on SMAC-v2 benchmark. Best in bold, second-best underlined.
Tacks
UPDeT-m
ODIS
STAIRS (Dars)
Terran
10,9
15.2
24.9
31.3
Seen
Zorg
Prosoes
6.6
99
12.7
10.4
20.3
21.0
32.5
29.2
Mean
31.0
Ternun
7.0
13.2
25.1
32.3
Unsenn
Zorg
Protoss
5.0
8.1
11.6
8.0
28.5
1888
32.8
25.0
Mean
30,0
Total Mean
7.4
11.5
30,3
5.4 ATTENTION DYNAMICS OVER TIME
In this section, analyze how attention maps evolve during SMAC episode in the 3m environment, using trajectories generated from our trained STAIRS-Former policy (Fig. 5)
At the beginning (t=0), all agents mainly attend their tokens, stabilizing local information under partial observability. By 1-4, agents O and 2, who first encounter enemies, shift attention to enemy tokens, while agent 1 maintains focus itself and leverages history tokens infer hidden state. At 8- all agents still focus on the enemy tokens, while agents and also attend to agent 0 protect the weakened ally. At 1-9, agent 0 successfully retreats and emphasizes history tokens to decide between counterattack and withdrawal, whereas agents and continue attacking while monitoring agent O's status. At t=14, agent becomes critically weak, agents and attend to each other while sustaining fire on enemy 1, and enemy 2. Finally, at (=21, agent is eliminated, and the surviving agents ( and 2 concentrate fire on enemy 2, demonstrating adaptive reallocation of attention between protective and offensive strategies. For a detailed explanation of the attention maps, please refer to Appendix G.I We also identify a complementary strategy, termed focus fire, which is discussed further in Appendix G.2
In contrast, attention maps from a basic transformer remain nearly uniform across tokens and time steps, regardless of context. This lack of selectivity shows its inability to prioritize critical tokens such as enemies or history, causing it to miss the temporal and relational structures. By comparison, STAIRS-Forme not only captures immediate interactions but also learns higher-level strategies such as focus fire and kiting, with attention dynamics closely aligned to observed tactical behaviors. This alignment highlights both its effectiveness and its interpretability in multi-agent decision making.
5.5 ABLATION STUDY
We conducted ablation studies three core components: (1) spatial recursive module, (2) temporal module, and (3) loken-dropout mechanism. Each was removed individually ("w/o"", where "w/o STD" excludes both spatial, temporal, and dropout.
Seen Tasks The spatial hierarchy is most critical for seen tasks, with performance dropping sharply when removed (77.9% 72.4%). In contrast, dropout and temporal abstraction yield little improvement in performance. This highlights that the rich correlation with entities is essential to capture the structured interactions within known environment.

---

Published as a conference paper at ICLR 2026
Table 5: Ablation results on Seen and Unseen tasks, "ST" = Spatial & Temporal, "STD" = ST Dropout. The best performance is shown in bold, and the second-best performance is underlined.
Tasks
STAIRS
w/o Temporal
w/o Spatial
w/o Dropout
w/o ST
w/o STD
Marine-Har
79.0
77.5
71.0
75.7
69.9
66.0
Marine-Fass
91.2
88.1
87.2
89.6
86.9
87.9
Seen
Kalker-Zealoo
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
Marine-Hard
57.0
$7.4
54.7
56.0
54.1
40.1
Marine-Easy
86.7
78.5
79.0
83.0
78.0
79.7
Unseen
Stalker-Zcalot
48.2
46.1
47.0
46.5
44.0
39.7
Mean
64.0
60.6
60.2
61.8
58.7
53.2
Total Mean
67.4
64.6
63.1
65.4
61.4
57.3
Unseen Tasks On unseen tasks, all components are essential. Removing dropout, spatial, or temporal modules lowers performance. Dropout improves generalization by mitigating overfitting, the temporal hierarchy captures long-term information crucial under partial observability, and the spatial hierarchy helps identify critical tokens for adapting new configurations. With all three, STAIRS achieves the best performance (64.0%), showing their joint importance for generalization to novel environments.
Considering both seen and unseen tasks, STAIRS consistently outperforms all ablations, achieving the highest overall mean (67.4%). The results clearly show that while the spatial hierarchy dominates performance on seen environments, the synergy of spatial, temporal, and dropout modules is essential generalization to unseen scenarios. Additional ablation results are provided in Appendix !
5.6 UNDERSTANDING ABLATION RESULTS THROUGH DORMANT NEURON ANALYSIS
To further assess the impact of structural components, analyze the proportion of dormant neurons across all 12 Marine-Hard tasks. Following Sokar et al. (2023), the score
STAIR
Doemam Meune Bamos
Erep (x(x)) of neuron i in layer f defined as HT Eten (x)
2.37
and neuron regarded as 7-dormant if s S r, with 0.05. Dormant neurons indicate under-utilized capacity. compute the average dormant neuron ratios for STAIRS and compare them against the ablated variants obtained by removing the two most influential components identified in
2.00
3.00
(a) Average dormant neuron ratios of STAIRS vs. ablations.
the ablation study (temporal and spatial attention).
As shown Fig. 6(a) both temporal and spatial modules reduce dormant neuron ratios, with the temporal module hav- ing the stronger effect. To examine this further, we ablate the GRU and the Temporal Focus Layer (TFL) within the tempo- ral module. Fig. 6(b) shows that TFL substantially reduces dormant neurons observation tokens, which drive Q-value estimation. By mitigating redundancy, increasing neuron ac- tivation, and improving the effective use of model capacity, TFL. plays central role in achieving better performance.
(b) Dormant ratios with token type of ab- lating GRU, TFL, TeRporalGRU++TFL)
Figure Dormant neuron ratios on marine-hard tasks: STAIRS vs. ablations.
CONCLUSION
In this work, we addressed the limitations of offline multi-agent reinforcement learning in multi- task settings, where transformers underutilize historical dependencies and relational structures. We proposed STAIRS-Former transformer with spatial and temporal hierarchies for selective attention to critical tokens and effective history use, while token dropout improves robustness across agent populations. Experiments the SMAC benchmark show that STAIRS STAIRS-Former Former achieves state-of-the art performance, underscoring the value structured attention for scalable and generalizable offline MARI.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
This work was supported in part by Institute of Information Communications Technology Plan- ning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No.RS-2022-1122) Development of Artificial Intelligence Technology for lelf-impoo Competency-Awwrc Learning Capabilities, 50%) and in part by the Institute of Information & Communications Technology Plan- ning Evaluation (ITTP) grant funded by the Korea government (MSIT) (No.RS-2022) -...R--2022-12220469 Development of Core Technologies for Tosk-oriented Reinforcement Learning for Commercialization of Autonomous Drones, 50%)
ETHICS STATEMENT
This study relies solely on publicly available benchmark environments (e.g., SMAC) and does involve human subjects, personal data, or sensitive information. All experiments and results adhere the ICLR Code of Ethics.
REPRODUCIBILITY STATEMENT
We provide detailed descriptions of the proposed models, training protocols, and evaluation procedures the main text and appendix. All datasets are publicly available, and anonymized source code with scripts is included in the supplementary materials.
REFERENCES
Rich Caruana. Multitask learning, Mach. Learn., 28(1):41-75, 1997.
Myungsik Cho, Jongeui Park, Jeonghye Kim, and Youngchul Sung. Ars: Adaptive reward scaling for multi-task reinforcement learning. In Forty-second International Conference on Machine Learning.
Junyoung Chung, Çaglar Gülçehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. CoRR, abs/1412.3555 2014.
Benjamin Ellis, Jonathan Cook, Skander Moalla, Mikayel Samvelyan, Mingfei Sun, Anuj Mahajan, Jakob Foerster, and Shimon Whiteson. Smacv2: An improved benchmark for cooperative multi- agent reinforcement learning. Advances in Neural Information Processing Systems, 36:37567- 37593, 2023.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: datasets for deep data-driven reinforcement learning. CoRR, abs/2004.07219, 2020.
Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NewrIPS 2021, December 6-14, 2021, virtual, pp. 20132-201455 2021.
Ahmed Hendawy, Jan Peters, and Carlo D D'Eramo. Multi-task reinforcement learning with mixture of orthogonal experts. arXiv preprint arXiv:2311.11333 2023.
Ahmed Hendawy, Jan Peters, and Carlo D'Eramo. Multi-task reinforcement learning with mixture of orthogonal experts. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024.
Siyi Hu, Fengda Zhu, Xiaojun Chang, and Xiaodan Liang. Updet: Universal multi-agent RI via policy decoupling with transformers. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OvenReview.net, 2021.
Jeewon Jeon, Woojun Kim, Whiyoung Jung, and Youngchul Sung. Maser: Multi-agent reinforcement learning with subgoals generated from experience replay buffer. In International conference on machine learning, pp. 10041-10052. PMLR, 2022.

---

Published as a conference paper at ICLR 2026
Elia Kaufmann, Leonard Bauersfeld, Antonio Loquercio, Matthias Müller, Vladlen Koltun, and Davide Scaramuzza. Champion-level drone racing using deep reinforcement learning. Nat., 620 (7976):982998))) 2023. doi: 10. 100.1088//11186-02--66666/
Woojun Kim and Katia P. Sycara. B3C: A minimalist approach to offline multi-agent reinforcement learning. CoRR, abs/2501.18135 2025. doi: 10.48500ARXI22500.1888.
Jakub Grudzien Kuba, Ruiqing Chen, Muning Wen, Ying Wen, Fanglei Sun, Jun Wang, and Yaodong Yang. Trust region policy optimisation in -aget reinforcement learning. In 10th International Conference on Learning Representations, ICLR 2022. OpenReview.net, 2022.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-TTen Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference Neural Information Processing Systems 2020, NewrIPS 2020, December 6-12, 2020, virtual, 2020.
Sicong Liu, Yang Shu, Chenjuan Guo, and Bin Yang. Learning generalizable skills from offline multi-task data for multi-agent cooperation. arXiv preprint arXiv:2503.21200 2025.
Ryan Lowe, Yi Wu, Aviv Tamar, Jean Harb, OpenAl Pieter Abbeel, and Igor Mordatch. Multi-agent actor-critic for mixed cooperaiiee-oompetiivee environments. Advances in neural information processing systems, 30, 2017.
Linghui Meng, Jingqing Ruan, Xuantang Xiong, Xiyun Li, Xi Zhang, Dengpeng Xing. and Bo Xu. M3: Modularization for multi-task and multi-agent offline pre-training. In Proceedings of the 2023 International Conference Autonomous Agents and Multiagent Systems, pp. 1624-1633, 2023.
Frans A. Oliehoek and Christopher Amato. A Concise Introduction to Decentralized POMDPs. Springer Briefs in Intelligent Systems. Springer, 2016. ISBN 978 978-3-319-28927-44 doi: 10.1007/ 978-3-319-28929-88
Shayegan Omidshafiei, Jason Pazis, Christopher Amato, Jonathan P. How, and John Vian. Deep decentralized multi-task multi-agent reinforcement learning under partial observability. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pp. 2681 1-900. PMLR, 2017.
Ling Pan, Longbo Huang, Tengyu Ma, and Huazhe Xu. Plan better amid conservatism: Offline multi-agent reinforcement learning with actor rectification. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (eds.), International Conference on Machine Learning, ICML 2022, 17-23 July 2022. Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 17221-17237 PMLR, 2022.
Bei Peng, Tabish Rashid, Christian Schroeder de Witt, Plerre-Alexandre Kamienny, Philip Torr, Wendelin Böhmer, and Shimon Whiteson. Facmac: Factored multi-agent centralised policy gradients. Advances in Neural Information Processing Systems, 34: 12208-1222, 2021.
Tabish Rashid, Mikayel Samvelyan, Christian Schroeder De Witt, Gregory Farquhar, Jakob Foerster, and Shimon Whiteson. Monotonic value function factorisation for deep multi-agent reinforcement learning. Journal of Machine Learning Research, 21(178):1-51, 2020.
Mikayel Samvelyan, Tabish Rashid, Christian Schroeder De Witt, Gregory Farquhar, Nantas Nardelli, Tim GJ Rudner, Chia-Man Hung, Philip HS Torr, Jakob Foerster, and Shimon Whiteson. The stareraft multi-agent challenge. arXiv preprint arXiv: 1902iv:902.0044 040404 2019.
Jianzhun Shao, Yun Qu, Chen Chen, Hongchang Zhang, and Xiangyang Ji. Counterfactual conserva- tive Q learning for offline multi-agent reinforcement learning. Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Anmual Conference on Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
12

---

Published as a conference paper at ICLR 2026
Ghada Sokar, Rishabh Agarwal, Pablo Samuel Castro, and Utku Evci. The dormant neuron phe- nomenon in deep reinforcement learning. In International Conference on Machine Learning, pp. 32145-32168. PMLR. 2023.
Chen Tang, Ben Abbatematteo, Jiaheng Hu, Rohan Chandra, Roberto Martin-Martin, and Peter Stone. Deep reinforcement learning for robotics: A survey of real-world successes. Annual Review of Control, Robotics, and Autonomous Systems, 8, 2024.
Zikang Tian, Ruizhi Chen, Xing Hu, Ling Li, Rui Zhang. Fan Wu, Shaohui Peng, Jiaming Guo, Zidong Du, Qi Guo, et al. Decompose task into generalizable subtasks in multi-agent reinforcement learning. Advances in Neural Information Processing Systems, 16:78514-78532 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017. December 4-9, 2017, Long Beach, CA, USA, pp.
2017a.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017b.
Jiawei Wang, Jian Zhao, Zhengtao Cao, Ruili Feng, Rongjun Qin, and Yang Yu. Multi-task multi- agent shared layers are universal cognition of multi-agent coordination, 2023a. URL https
/ 1/arxiv. http://bb/2331...11
Xiangsen Wang, Haoran Xu, Yinan Zheng, and Xianyuan Zhan. Offline multi-agent reinforcement learning with implicit global-to-locse value regularization. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023b.
Yaodong Yang, Jianye Hao, Ben Liao, Kun Shao, Guangyong Chen, Wulong Liu, and Hongyao Tang. Qatten: A general framework for cooperative multiagent reinforcement learning. arXiv preprint arXiv:2002.03939 2020.
Yigin Yang, Xiaoteng Ma, Chenghao Li, Zewu Zheng, Qiyuan Zhang, Gao Huang, Jun Yang, and Qianchuan Zhao. Believe what you see: Implicit constraint approach for offline multi-agent reinforcement learning. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 10299-103 2, 2021a.
Yigin Yang, Xiaoteng Ma, Chenghao Li, Zewu Zheng, Qiyuan Zhang, Gao Huang, Jun Yang, and Qianchuan Zhao. Believe what you see: Implicit constraint approach for offline multi-agent agent reinforcement learning. In M. Ranzato, A. Beeygllzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processsig Systems, volume 34, pp. 10399-10312
Curran Associates, Inc., 2021b. URL httpp:://pppooeeiiss...e..iioois co/paper http:///2.....a222....1..11631aa
Fuxiang Zhang, Chengxing Jia, Yi-Chen Li, Lei Yuan, Yang Yu, and Zongzhang Zhang. Discover- ing generalizable multi-agent coordination skills from multi-task offline data. In The Eleventh International Conference on Learning Representations, 2023.
Mingliang Zhang, Sichang Su, Chengyang He, and Guillaume Sartoretti. Hybrid training for enhanced multi-task generalization in multi-agent reinforcement learning. arXiv preprint arXiv:2408. 13667, 2024.
Tianze Zhou, Fubiao Zhang, Kun Shao, Kai Li, Wenhan Huang, Jun Luo, Weixun Wang, Yaodong Yang. Hangyu Mao, Bin Wang. Dong Li, Wulong Liu, and Jianye Hao. Cooperative multi-agent transfer learning with level-adaptive credit assignment. CoRR, abs/2106.000 2021.
13

---

Published as a conference paper at ICLR 2026
A THE USE OF LARGE LANGUAGE MODELS (LLMS)
Large language models were employed as auxiliary tools improve readability, refine phrasing, and perform grammar checks. They were not used for research ideation, methodological design, or experimental analysis, and did not contribute the generation of research results.
B SMAC BENCHMARKS
We adopt two widely used benchmarks to evaluate approach for offline MARL with multi-task datasets proposed by Zhang et al. (2023). Our primary benchmark is the StarCraft Multi-Agent Challenge (SMAC) (Samvelyan et al. 2019), which has become standard testbed for cooperative multi-agent reinforcement learning. SMAC provides a variety of micromanagementt scenarios where agents must coordinate under partial observability, facing both homogeneous and heterogeneous unit dynamics. These characteristics make SMAC challenging and realistic environment, well-suited for assessing the robustness and scalability of offline MARL methods complex domains.
14

---

Published as a conference paper at ICLR 2026
C DATASET PROPERTIES
We use the offline datasets provided by ODIS (Zhang et al., 2023). which are based on the PyMARL implementation of QMIX (Rashid et 2020). Similar the D4RL benchmark (Fu al. 2020
four dataset qualities are defined:
Expert: trajectories collected by QMIX policy trained for 2M environment steps, achieving high test win rates.
Medium: trajectories collected by weaker QMIX policy whose win rate is roughly half of the expert policy.
Medium-Expert a mixture of the expert and medium datasets, providing increased diver- sity.
Medium-Replay: the replay buffer of the medium policy, containing lower-quality trajecto- ries sampled during training.
For each source task, the expert and medium datasets contain 2,000 trajectories each, while the medium-expert dataset includes 4,000 trajectories as their union. The size of the medium-replay dataset depends the number of trajectories collected before the medium policy terminates training. During multi-task training, we use up to 2,000 trajectories per task (or all available trajectories when fewer exist), and merge them across tasks form unified multi-task dataset. The detailed statistics of these datasets are summarized in Table
Table 6: Properties of offline datasets with different qualities.
Task
Quality
# Trajectories
Average return
Average win rate
expert
2000
19.89
0.99
medium
2000
13.99
0.54
3m
medium-expert
4000
16.94
0.77
medium-replay
3630
N/A
N/A
expert
2000
19.94
0.99
medium
2000
17.33
0.74
5m
medium-expens
4000
18.63
0.87
medium-replay
771
N/A
N/A
expert
2000
19.94
0.99
medium
2000
16.63
0.54
10m
medium-cxpert
4000
18.26
0.76
medium-replay
571
N/A
N/A
expert
2000
17.34
0.72
medium
2000
12.64
0.28
5m_v jm_
medium-expert
4000
14.99
0.50
medium-replay
32607
N/A
N/A
expert
2000
19.61
0.94
medium
2000
15.50
0.41
9m_ 9m_vs_ 10m
medium-expent
4000
17.56
0.68
medium-replay
13731
N/A
N/A
expert
2000
19.77
0.96
medium
2000
16.63
0.45
2s3z
medium-cxpert
4000
18.20
0.70
medium-replay
4505
N/A
N/A
expert
2000
19.74
0.95
medium
2000
16.87
0.50
2s4z
medium-expert
4000
18.31
0.72
medium-replay
6172
N/A
N/A
expert
2000
19.79
0.95
medium
2000
16.31
0.31
3s5z
medium- medium-expert
4000
18.05
0.63
medium-replay
11528
N/A
N/A
15

---

Published as a conference paper at ICLR 2026
D DETAIL DESCRIPTIONS OF TASK
In our experiments, we select three representative tasks from SMAC: marine-easy, marine-hard, and stalker-zealot. The marine-easy and marine-hare tasks both involve homogeneous units marines: marine-casy is balanced setting with equal allied and enemy counts, while marine-hard introduces more difficult cases where enemies are equal to or greater than allies. The stalker-zealot task, in contrast, is heterogeneous with two unit types (stalkers and zealots) distributed identically across both sides. Together, these tasks span different levels of difficulty and unit diversity, providing comprehensive testbed for evaluating coordination strategies under varying conditions, with detailed specifications given in Tables 7
Table 7: Descriptions of marine-easy tasks
Task type
Task
Ally units
Enemy units
Properties
3m
Marines
Marines
homogeneous & symmetric
Source
5m
5 Marines
Marines
homogeneous & symmetric
10m
10 Marines
10 Marines
homogeneous symmetric
4m
Marines
Marines
homogencous symmetric homogeneous & symmetric
6m
6 Marines
Marines
7m
Marines
Marines
homogeneous symmetric
Unseen
8m
8 Marines
Marines
homogeneous & symmetric homogeneous & symmetric
9m
9 Marines
Marines
11m
11 Marines
11 Marines
homogeneous symmetric
12m
12 Marines
12 Marines
homogeneous & symmetric
Table 8: Descriptions of marine-hane tasks
Task type
Task
Ally units
Enemy units
Properties
3m
Marines
Marines
homogeneous & symmetric
Source
$m_vs_6n
Marines
6 Marines
homogeneous & asymmetric
9m_vs_ vs_10m
Marines
10 Marines
homogeneous asymmetric
4m
Marines
Marines
homogeneous & symmetric
5m
Marines
5 Marines
homogeneous & symmetric
10m
10 Marines
10 Marines
homogeneous & symmetric
12m
12 Marines
12 Marines
homogeneous & symmetric homogeneous asymmetric homogencous asymmetric homogeneous & asymmetric
Unseen
7m_vs_8m
Marines
8 Marines
8m_vs. 9m
Marines
Marines
10m_vs_ vs_ 11
10 Marines
11 Marines
10m_vs vs_ 12m
10 Marines
12 Marines
homogencous asymmetric
13m vs 15m
13 Marines
15 Marines
homogeneous & asymmetric
Table 9: Descriptions of stalker-zealot tasks
Task type
Task
Ally units
Enemy units
Properties
2s3z
Stalkers, 3 Zealots
2 Stalkers, 3 Zealots
heterogeneous & symmetric heterogeneous & symmetric
Source
2s4z
2 Stalkers, 4 Zealots
2 Stalkers, 4 Zealots
3s5z
Stalkers, 5 Zealots
Stalkers, 5 Zealots
heterogeneous & symmetric
1s3z
Stalker, Zealots
Stalker, Zealots Stalker, 4 Zealots
heterogeneous & symmetric
1s4z
Stalker, 4 Zealots
heterogeneous & symmetric heterogeneous & symmetric heterogeneous & symmetric heterogeneous & symmetric heterogeneous & symmetric heterogeneous & symmetric heterogeneous & symmetric heterogeneous & symmetric
1s5z
Stalker, Zealots
Stalker, Zealots
2s5z
Stalkers, 5 Zealots
2 Stalkers, 5 Zealots
Unseen
3s3z
3 Stalkers, 3 Zealots
3 Stalkers, 3 Zealots
3s4z
Stalkers, 4 Zealots
3 Stalkers, 4 Zealots
4s3z
4 Stalkers, 3 Zealots
4 Stalkers, 3 Zealots
4s4z
Stalkers, 4 Zealots
4 Stalkers, 4 Zealots
4857.
4 Stalkers, 5 Zealots
4 Stalkers, 5 Zealots
16

---

Published as a conference paper at ICLR 2026
E RESULTS ON MARINE-EASY TASK SET
The results for the Marine-Easy task set are presented in Table Across both source and unseen tasks, STAIRS-Former consistently delivers strong performance. On the Expert dataset, it matches HiSSD by achieving nearly perfect success rates on average, demonstrating that our model can fully exploit high-quality data. On the Medium and Medium-Experr datasets, SEAIRS-Former significantly ouiperforms prior methods, showing clear advantages in handling sub-optima data. For example, compared to HiSSD, STAIRS-Former improves average performance by + 1-16.6 0% on Medium, +26.6% on Medium-Fxpeer However, on the Medium-Replay dataset, STAIRS- underperforms compared to HiSSD. We attribute this to the relatively small size of the Medium- Replay dataset in Marine- -asy (see Table which limits trajectory diversity and thus reduces the effectiveness of offline reinforcement learning. To further mitigate overfitting caused by the limited trajectories, results for this task reported at 10K time steps.
Table 10: Comparison of average and per-task performances the Marine-Eat -asy task set across four dataset qualities. We report - deviation, with the best shown in bold.
Expert
Medium
UPDeT-m
ODIS
HiSSD
STAIRS (Ours)
UPDeT-m
ODIS
HiSSD
STAIRS (Ours)
3m
58.8 20.1
814
100.0
99.1
55.6 28.8
00.0
67.5 10.5
85.6
5m
48.8 35.7
79:1 29.0
99.1 14
99.1
719
79.1 10.3
80.6
85.0 9.2
46.3
99.4
77.5
66.3
Unseen Tasks
22.5 17.6
48.1
97.5 3.4
9.1
55.0
73.8
6m
32.5 39322
44.4
44.7
1.0
96.9
70.6 31.1
87.5 17.3
86.3 8-4
$2.5 9.3
7m
36.92
42.5
47.2
1.0
100.0
53.8 37.9
81.3 22.9
93.8 12.3
98.1
8m
31.3
39.5
59.4 43.2
1.0
1..4
78.8 12.8
68.8
9..6
96.9 1
9m
45.6
35.7
64.4
37.6
100.0
100.0 0.0 0.0
52.5 17.9
79.4
76.9 1.47
93.1 5.1
11m
38.1
+ 32.8
31.9
99.1
100.0 0.0
26.9 9.9
51.3
11.8
48.1 7.2
65.6
14.5
12m
33.1
+ 28.2
69.4
39.1
96.3 4.1
98.1 1.7
13.9 20.0
3133 - 16.3
40.6 17.8
65.6
6.6
Avg
63.4
99.2
99.0
51.3
69.2
72.5
Medium-Bee Medium-Expert
Medium Medium-Replay Replay
Tasks
UPD:T-m
ODIS
HiSSD
STAIRS (Ours)
UPDeT-m
ODIS
HisSD
STAIRS (Ours)
Source Tasks
3m
48.1 34.33
76.3 20.4
81.3 17.3
98.8
25.8 31.9
50.0 33.1
87.5 6.6
96.9 6.8
5m
66.3 19.2
84.4 1.9.9 9.9
80.0 1. 16.0
98.8 1.7
0.0 1 0.0
1 0.0
85.0 8.4
89.4 7.8
10m
60.6 37.8
51.9 + 30.1
74.4 +: 19.4
100.0 1 0.0
0.0 0.0
0.0
85.6 8.4
56.9 + 18.7
Unseen Tasks
4m
24.4 1. 17.6
64.4 1. 29.7
75.6 1: 6.0
60.0 25.4
0.0 1 D.D
15.6 1. 34.9
66.9 10.0
79.4 1: 13.0
6m
46.9 1:00::000 33.9
67.5
33.6
75.6
14.6
94.4 2: 4.6
0.0 = 0.0
0.0 +: 0.0
100.0 1: 0.0
91.3 : 6.4
7m
37.5
62.5
22.0
73.1
12.9
96.9 + 3.1
0.0 = 0.D
0.0 0.0
99.4 1.4
90.6 5.8
8m
10.6 1: 7.5
45.6
13.7
71.9
6.3
84.4 18.8
0.8 = 1.6
0.0 0.0
96.9 2.2
83.1 8.1
9m
48.1
40.2
62.5
30.8
73.1
18.6
100.0 +: 0.0
0.0 0.D
0.0 0.0.10 0.0
80.6 5.1
82.5 5.7
11m
55.6
30.9
19:1
37.8
57.5
19.0
98.1 + 1.7
0.0 0.0
0.0 10:200 0.0
53.1
18.9
55.0 + 23.7
12m
36.9
31.5
10.6
28.5
69.1
09. 9.7
95.6 4.7
0.0
0.0 0.0
36.3 11.2
19.4 + 30.0
Avg
43.5
60.5
73.2
92.7
2.7
6.6
79.1
76.5
17

---

Published as a conference paper at ICLR 2026
F COOPERATIVE NAVIGATION TASK
In addition to SMAC, we also include the Cooperative Navigation (CN) task from the Multi-Agent Particle Environment (MPE) (Lowe al. 2017) supplementary benchmark. CN provides a simpler but complementary setting, where multiple agents must coordinate to occupy distinct landmarks while avoiding collisions. While less complex than SMAC, this environment emphasizes pure cooperation, making it useful supplement to our main benchmark and enabling us to test the generality of our approach across different types of multi-agent scenarios. The detailed specifications of CN-tasks are in T1
Table 11: Properties of offline datasets with different qualities.
Task
Quality
Trajectories
Average return
Average win rate
expert
2000
1.0000
1.0000
CN-2
medium
2000
0.6152
0.6152
expert
2000
0.7173
0.7173
CN-4
medium
2000
0.4273
0,4273
Relative to the recent state-of-the-art HiSSD, STAIRS-Former achieves higher scores in the Expert setting, improving from 49.1 to 51.3, and also shows gains in the Medium setting, increasing from 13.2 to 14.3. These results, obtained in the MPE domain in addition to our main SMAC experiments, indicate that STAIRS-Former provides modest but consistent improvements over HiSSD across different environments.
Table 12: Comparison of average and per-task performances the Cooperative navigation task set across two dataset qualities. We report standrd deviation, with the best shown in bold.
Expert
Tasks
UPDeT-m
ODIS
HSSS
STAIRS (Ours)
Source Tasks
CN-2
68.8 + 19.4
78.8 44.1
100.0 0.0
100.0 + 0.0 30.0 + 12.6
CN-4
13.8 = 12.6
21.9 +
15.1
24.4 = 1.4
Unseen Tasks
CN-3 CN-5
34.4 + 19.1
48.8 + 27.8
65.0
10.2
64.4 5.7
1.9 = 2.8
5.6 + 3.4
6.9 = 7.5
10.6 + 1.7
Average
29.8
49.1
51.3
Medium
Tasks
UPDeT-m
ODIS
HiSSD
STAIRS (Ours)
Source Tasks
CN-2 CN-4
11.1
16.7 14.4
38.8 11.2
45.0.5 13.6
1.9 1.7
2.1 = 3.6
2.5 - 2.6
1.9 2.8
Unseen Tasks
CN-3
3.1
2.2
5.2 4.8
8.8 3.4
8.8 + 2.6
CN-5
0.0 = 0.0
0.0 + 0.0
2.5 = 2.6
1.3 1 1.7
Average
3.5
13.2
14.3
18

---

Published as a conference paper at ICLR 2026
G VISUALIZATION OF ATTENTION MAP
G.1 ATTENTION HEATMAP
We describe how the attention map constructed. As shown in Fig. 7 the vertical axis corresponds to queries and the horizontal axis keys. Each entry denotes an attention weight, computed as
Key
Own
HH
As shown Fig. 7 queries (vertical) attend to keys (horizontal). with each entry attention weight: softmax (QK?) /vdx)
Own
Query
Attention(Q. K) softmax (OKT//QQ/T// /v.s).
(8)
Here, Q (queries) and K (keys) are linear projections of the input tokens that determine, respectively, what a token attends to and what it provides to others. The scaling factor Vdg normalizes the dot-product by the dimensionality of the key vectors, preventing excessively large values that could saturate the softmax. Once the attention weights are computed, they are applied to the corresponding values V. another linear projection Figure 7: Structue of attention map of the tokens containing the actual information to be aggregated. In this way, the attention mechanism produces weighted sum of the values, where the weights specify how strongly each token attends to others.
L4
Building on this formulation, the SMAC 3m task constructs tokens by first decomposing the agent's observation into three categories: the agent's 's own token, enemy tokens (E0, E1, E2), and ally tokens (A0, A1). In addition these observation-derived tokens, the model also incorporates two history state tokens, namely the low-level history token (LH) and the high-level history token (HH). The LH token functions as short-term memory that captures fine-grained temporal dependencies, while the HH token provides a more abstract representation that summarizes longer-horizon information.
G.2 ALTERNATIVE TRAJECTORY EVOLUTION IN A SMAC 3m EPISODE
Figure 8: Another temporal evolution in SMAC 3m episode: STAIRS-Former (Above) vs. HiSSD (Below)
Compared to Figure 5.4 we present alternative trajectory that illustrates the focus-fire strategy. As shown in Figure g the agents behave almost identically to those in Figure 5 up to 1=4, and thus the attention distributions at this stage remain largely the same. However, notable divergence occurs at t=8, when all agents collectively direct their attention toward enemy 2. This coordinated decision to execute focus-fire results strong emphasis on enemy 2's token, and consequently, enemy is quickly eliminated from the battlefield. Unlike the main trajectory in Figure 5, where agent 0 with the lowest health retreated to preserve survivability, here agent 0 remains engaged in the fight and is eliminated immediately after enemy 2's death at 1= Following this loss, agents I and 2 shift their focus toward history tokens, reflecting a period of reassessment as they deliberate between possible countermeasures under partial observability. Ultimately, at 222 the two surviving agents reestablish coordination and concentrate their attention on the remaining enemy
G.3 SUPPLEMENTARY TEMPORAL VISUALIZATION: ATTENTION MAPS IN OTHER TASKS
In Section 5.4 and G.2 we analyzed the evolution of attention maps and real trajectories on the SMAC 3m scenario. To further assess the generality of our method, we extend these visualizations
19

---

Published as a conference paper at ICLR 2026
to additional tasks, including both seen and unseen scenarios every five timesteps. Specifically, report results on the challenging marine-hard setting, adding three tasks: two unseen tasks (4m, 8m 9m) and one seen task (5m), in addition to the 3m task previously shown.
Across tasks, our attention maps consistently highlight critical tokens while adaptively leveraging historical information, demonstrating the ability to capture both local interactions and temporal dependencies. contrast, HiSSD fails to attend to critical tokens and exhibits little utilization of history, limiting its ability to model long-term coordination. These results confirm that our method generalizes well to diverse tasks, including those unseen during training, and robustly captures essential spatio-temporal dynamics.
4m (Unseen)
030
0.1
Sm.rii8m (Seen)
0.00
8mLvs_9m (Unseen)
(a) 0
(0) 5
(c) 10
(d) t 15
(e) t 20
(e) t 25
Figure 9: Attention maps and trajectories on other tasks: STAIRS-Former (Above) vs. HiSSD (Below)
G.4 AVERAGE ATTENTION MAPS OVER WHOLE EPISODES WITH HISSD)
While the previous subsections focused trajectory-level analyses selected timesteps, to aggregated statistics over entire episodes. In particular, we compute the average attention maps of the HiSSD transformer across all timesteps and episodes for each of the benchmark tasks, including marine-hard, marine-easy, and stalker-zealot. These averaged maps reveal the characteristic behavior of HiSSD: attention distributions diffuse and fail to consistently concentrate on critical tokens, suggesting limited ability to capture task-relevant structures over long horizons.
turn
G.5 AVERAGE ATTENTION MAPS OVER WHOLE EPISODES WITH STAIRS (OURS)
We conduct the same analysis with STAIRS-Former averaging attention maps across full episodes for the same set of tasks (marine-hard, marine-eusy, and stalker-zealot). contrast to HiSSD, our method exhibits sharper token-level concentration, consistently highlighting important entities while incorporating historical tokens when necessary. Since these maps are averaged over all timesteps, the degree of focus critical tokens appears less pronounced than in the timestep-speciffc visualizations. Nevertheless, "TAIRS-Former still demonstrates clearer emphasis on task-relevant tokens compared HiSSD, maintaining more coherent and interpretable attention allocation throughout entire episodes. Overall, these results reinforce the advantage of STAIRS-Formin in modeling complex multi-agent coordination.
20

---

Published as a conference paper at ICLR 2026
Seen Tasks
3m
5m_vs_6m
9m_vs_10m
Unseen Tasks
4m
5m
10m
12m
7m_vs_8m
8m_vs_9m
10m_vs_11m
10m_vs_12m
13m_vs_15m
Figure 10: Average attention map on marine mardne-hard task with HiSSD
21

---

Published as a conference paper at ICLR 2026
Seen Tasks
2s3z
2s4z
3s5z
Unseen Tasks
1s3z
1s4z
1s5z
2s5z
3s3z
3s4z
4s3z
4s4z
4s5z
4s6z
Figure 11: Average attention map on stalker-zealot task with HiSSD
22

---

Published as a conference paper at ICLR 2026
Seen Tasks
3m
5m
10m
Unseen Tasks
4m
6m
7m
8m
9m
11m
12m
Figure 12: Average attention map on marine-easy task with HiSSD
23

---

Published as a conference paper at ICLR 2026
Seen Tasks
3m
5m_vs_6m
9m_vs_10m
Unseen Tasks
4m
5m
10m
12m
7m_vs_8m
8m_vs_9m
10m_vs_11m
10m_vs_12m
13m_vs_15m
Figure 13: Average attention map on marine-hard task with STAIRS
24

---

Published as a conference paper at ICLR 2026
Seen Tasks
2s3z
2s4z
3s5z
Unseen Tasks
1s3z
1s4z
1s5z
2s5z
3s3z
3s4z
4s3z
4s4z
4s5z
4s6z
Figure 14: Average attention map on stalker-zealot task with STAIRS
25

---

Published as a conference paper at ICLR 2026
Seen Tasks
3m
5m
10m
Unseen Tasks
4m
6m
7m
8m
9m
11m
12m
Figure 15: Average attention map on marine-easy task with STAIRS
26

---

Published as a conference paper at ICLR 2026
H TRAINING DETAILS
H.1 HYPER-PARAMETERS
In this section, provide the hyperparameters for STAIRS-Former used in the SMAC offline MT-MARL benchmarks in Table 13 Across all tasks in benchmarks, we use same hyperparmeters.
Table 13: Haper-parameters of STAIRS-Fome
Hyper-parameter
Value
hidden layer dimension
64
atlention dimension
1.0
optimizer
Adam
learning rate
0.0005
Number of layers M
2
Recursive steps 1/1
Recursive steps 1/2
1
Temporal interval Tμ
Dropout ratio Pdrop
0.1
Training timesteps
30,000
H.2 TRAINING COST
To measure computational cost, we used a single NVIDIA RTX 4090 GPU (24.565 MiB memory usage). Training for 50K steps took approximately hours 20 minutes for HiSSD, hours for ODIS, whereas our method required only 4 hours under the same setup.
H.3 PARAMETERS AND MEMORY USAGE
On the Marine-Hares -adadd-Medium task, UpDeT-m, ODIS, method, and HiSSD 79,095, 138,573, 220,023, and 679,335 parameters.
In terms of GPU memory consumption, UpDeT-m and ODIS require 7,046 MiB and 7,020 MiB, HiSSD requires 17,492 MiB, and our method uses 14,370 MiB. While slightly heavier than ODIS and UpDeT-m,, our model remains substantially more efficient than HiSSD and achieves significantly higher performance.
Table 14: The parameter counts and GPU memory footprint for UpDeT-m, ODIS, HiSSD, and our method (STAIRS) on the Marine-Hard-Menn task
Algorithms
UpDeTTm
ODIS
HiSSD
STAIRS
Parameters
79.095
138,573
679,335
220,023
GPU Memory Usage (MiB)
7,046
7,020
17,492
14,370
27

---

Published as a conference paper at ICLR 2026
I ADDITIONAL ABLATION STUDIES
The main hyperparameters of our STAIRS-Former are the temporal interval Tи for long-term dependency, and the token dropout ratio pårop- this section, we conduct ablation studies on each hyperparameter to examine their effect on performance. Note that we conduct all ablation studies without Temporal Focus Layer (TFL).
1.1 ABLATION STUDY ON THE HYPERPARAMETER Tи AND PDROP
First, we conducted ablation studies on T'H and Pdrop- Table 15 shows the average performance for different values of Tи and Parop- The results show that performance remains robust across various settings, except when token dropout is not used (Purop 0). These results highlight that our token dropout mechanism is essential for enhancing performance.
Table 15: Comparison of average performances over all task set and dataset qualities. Best perfor- mance shown in bold.
Temporal Interval TH
3
5
Token Dropout Rate Paop
0.05
0.1
0.05
0.1
0.05
(1.1
Average Performance
65.9
65.3
64.5
66.2
63.6
66.2
65.6
28

---

Published as a conference paper at ICLR 2026
J BMAC-V2
In addition SMAC, also include SMAC-v2 Bllis al. (2023) a supplementary benchmark, which is a more complex and realistic environment compared to SMAC-vI. SMAC-v2 introduces significantly higher stochasticity due to randomized initial unit placements with dynamic team compositions and unit types. These changes make the environment less deterministic and substantially more challenging than SMAC-v1, especially for offline RI. algorithms.
We generated the SMAC-v2 offline datasets using QMIX Rashid et al. (2020) implemented in PyMARL, collecting 2, 2,000 trajectories for each task. The average return and win rate across all tasks are summarized in Table
Table 16: Properties of offline datasets with different qualities.
Task
Quality
# Trajectories
Average return
Average win rate
medium
2000
11.6
0.44
Terran 3_vs_
medium-replay
2000
N/A
N/A
medium
2000
13.09
0.4
Terran 5_ $_
medium-replay
2000
N/A
N/A
medium
2000
12.58
0.42
Terran 10 vs_10
medium-replay
2000
N/A
N/A
medium
2000
16.44
0.42
Protoss 3_vs_3
medium-replay
2000
N/A
N/A
medium
2000
17.98
0.41
Protoss 5_vs_5
medium-replay
2000
N/A
N/A
medium
2000
19.12
0.42
Protoss 10_vs_10
medium-replay
2000
N/A
N/A
medium
2000
11.6
0.44
Terran 3_vs_3
medium-replay
2000
N/A
N/A
medium
2000
13.09
0.4
Terran 5_vs_5
medium-replay
2000
N/A
N/A
medium
2000
12.58
0.42
Terran 10_ 10_vs_10
medium-replay
2000
N/A
N/A
medium
2000
9.75
0.43
Zerg 3_vs_3
medium-replay
2000
N/A
N/A
medium
2000
13.53
0.41
Zerg 5_vs_5
medium-replay
2000
N/A
N/A
medium
2000
13.56
0.4
Zerg 10_vs_10
medium-replay
2000
N/A
N/A
Since SMAC-v2 stochastic environment, the map configuration for each race determined probabilistically. For instance, in the Terran race, units are sampled according to predefined weights-mare (0.45), marauder (0.45), and medivac (0.10). Similarly, the starting formation is sampled from the surrout huurounnded_dd laclled__nd______ re treflett distribution, where the agents are surrounded with probability 0.5 and placed a reflected configuration with probability 0.5. Because SMAC-v2 substantially more challenging than SMAC-v1, we evaluate our method on settings with equal num- bers of allied and enemy units (e.g., 3 vs 5,5 vs_ 5), similar spirit to the classic mar ineeeasy and stalker-zealot zealot scenarios.
The probabilistic generation rules for each race are summarized Table 17
The results for the complete SMAC-V2 task suite are presented in Table 18 Our method achieves substantial performance improvements across all races. Terran tasks, it improves performance by approximately 292% over UpDeT-m, 132% over ODIS, and 28% over HiSSD. Similarly, in Protoss tasks, we observe gains 280%. 175%. and 14%, respectively. Zerg tasks exhibit comparable improvements, with increases of 381 1% over UpDeTTmm, 201 % over ODIS, and 35% over HiSSD.
29

---

Published as a conference paper at ICLR 2026
Table 17: Unit generation and start-position configuration for each SMAC-v2 race.
Race
Unit Types (weights)
Start Pos. Dist.
Terran
marine (0.45), marauder (0,45), medivac (0.10)
surrounded_ __d__refeeectttttt
reflect (p=0.5)
Protoss
stalker (0.45), zealot (0.45), colossus (0.10)
prrrunded__nddreeee (p=0.5)
Zerg
zergling (0.45), baneling (0.10), hydralisk (0.45)
(p=0.5)
Aggregated over all SMAC-V2 tasks, our approach outperforms UpDeTTm,, ODIS, and HiSSD by roughly 310%, 164%, and 24%, respectively. These results demonstrate that our method generalizes effectively across all races, maps, and unit compositions, even under the high stochasticity inherent SMAC-V2.
While the improvement over HiSSD is relatively smaller compared to the other baselines, it is important to note that HiSSD requires more than twice the number parameters (679,335 vs. our 220,023) and nearly double the training time. Thus, the comparison remains strongly favorable to our method in terms both performance and efficiency.
Table 18: Comparison of average and per-task performances on the SMAC- V2 task set. We report standardd deviation, with the best result shown in bold. For brevity, we abbreviate task names such as 3 3_v3_ to Terran 3, rotoss 3, and so on.
Tasks
Medium
Medium-rephyy
UPDeT-m
ODIS
HiSSD
STAIRS (Ours)
UPDeT-m
ODIS
HiSSD
STAIRS (Ours)
Terran Source Tisks
Terran
15.0
18.1 9.2
31.3 11.0
37.5 + 5.8
106 5.2
19.4 12.2
31.3 8.6
28.1 15.9
Terran
6.0
16.3 1 6.8
18.8 3,1
26.9 5.7
8.1
11.9 10.0
23.1 11.0
3133
Tacks
Terran
11.94:26
19.4 1:
33.8
35.6 6 9.5
10.6
16.3 1: 12.8
25.0 8.8
36.3
5.7
Teran
10.6
144
26.9 10.5
26.9
6.8
13.1 10.0
214= 5.6
33.8
10.5
Terra
10.6 6.8
5.7
35.8 + 11.0
35.0
8.1 4.7
15.0
10.2
25.0 2 9.4
28.8
7.8
Terrnn
4.1
14.8
26.9
38.8 4.7
5.4
16.3
15.1
9.2
12.2
Terran
5.7
12.5 8.6
28.9
35.6
13.2
3.6
19.1
9.2
19.4
8.9
25.0
10.4
Terrnn 11
2.5 5.6
6.3 4.9
20.6 5.7
37.5 = 11.0
5.0 7.2
5.6 + 5.1
23.1
23.1-6.1
24.4 + 15.1
12
3.1
4.2
Terrun Avg
9.3
14.7
28.6
35.0
7.0
12.9
23.4
28.9
Protoss Source Tasks
Protoss
1.90
14.4 15.1
30.6 1 7.1
28.1 6.6
8.1
14.4 14.4.19 9.8
28.1 1. 12.5
28.1 8.8
Protoss
13.1:18..
9.1 7.3
42.5 9.3
39.4 11.8
5.177.
16.9 :: 12.2
28.1 =: 7.3
43.1
4.1
Protoss 10
10.0 9.7
11.9 a 14.0
27.5 à: 7.1
31.3 1: 6.3
6.3 4: 9.6
8.8 1: 11.6
20.0 a: 8.1
25.0 1: 6.3
Protoss Unseen Tasks
Protoss
20.0 12.2
13.1 12.2
35.0 6.0
35.1 13.9
4.9
19.4 16.4
33.1
37.5 9.1
Protoss
12.5
10.6 9.5
35.0
11.4
40.0
9.2
5.0 3: 7.8
11.3
41.9
12.0
12.5
11.2
Protnss 8
14.4 8.4
15.0 19.1
25.6 6.0
36.9 7.1
4.7
11.3
23.1
7.2
31.9
8.9
Protoss 9
13.1 16.7
40.0 8.1
31.3 14.8
1.7
11.3
20.0=81
33.1
6.5
6.5
11.8
33.8
35.6
8.8
29.4
Protoss 12
3.8 4.1
8.8 + 10.2
20.0 + 3.6
23.8 = 15.4
3.8 2.6
5.0 = 3.6
15.6 16.6211 4.4
14.4
4.2
Arg
Zerg Source Tasks
Zerg 3 Zerg
11.3 8.1
13.1
28.1 1: 8.3
33.1 6.1
5.1
11.3
27.5 11.1
37.5 14.8
10.6 4.2
11.3 2.8
15.6 + 8.0
28.8 8.7
5.+
11.9 9.5
17.6 =:3
20.6 + 7.5
Zerg
3.8
19.3
Zerg Tacks
Zerg 4
11.9 10.0
15.0 13.0
18.1 5.6
33.8 8.1
6.3 4:
8.1 1111:55
19:4 20 4.1
23.8 4.7
Zerg
10.5
4.7
Zorg.
8826.
12.54 9.9
26.3
9.8
28.8
1 7-1
3.8 4.1
3.1 4.4
13.1 3.4
23.1
9.8
Zerg
6.1
8.8
25.0
13.4
35.0
16.1
2.5 +
13.1
10.4
Zog
4.7
15.0 15.2
23.1
10.3
31.9
3.6
3.1
13.1 =8
20.0
6.5
Zorg 11
2,5 2.4
11.3
12.8
25.0
4.9 4.9
23.5
10.3
1.9 2.8
1.9 4.2
1133
15.6
7.3
Zong
3.8
Zerg Avg
7.1
12.2
23.3
30.3
3.8
5.2
15.6
22.1
30

---

Published as a conference paper at ICLR 2026
K MAMUJoCo
In addition to SMAC, we also include Multi-Agent MuJoCo (MaMuJoCo) (Peng et al. 2021 as an additional supplementary benchmark, which is more complex and realistic robotic environment with continuous aciton space. MAMuJoCo models single robot as multiple cooperating agents. Each agent is responsible for controlling a designated group of joints, and the agents must collaborate and align their actions to achieve the robot's overall goals.
HISSD Liu et al. 2025) introduced the MAMuJoCo benchmark for offline multi task multi agent (MAMA) reinforcement learning (RL) to demonstrate the performance of their method on a realistic robotic system with continuous control. Their task set is built using the HalffChettah-22 environment with six agents in MAMuloCo and each task is formed by disabling one agent. The offine dataset for each task is collected using a HAPPO trained policy (Kuba al. 2022). However the dataset is not publicly available and the observations are based on the full state of "HalsCheetah-v2 rather than agent specific local observations. This makes the dataset unsuitable for evaluating STAIRS because STAIRS focuses on leveraging history tokens to mitigate partial observability in the offline MTMA setting. Furthermore the task configuration in HISSD (Liu et al. 2025) uses the same number of agents and identical observation spaces except for the non disabled case which limits its ability to test robustness under varying agent configurations.
To accommodate the offline MTMA learning setting, we construct customized multi-task dataset in' HalfCheetah-v22 following the general procedure of Wang et al. (2023a). Unlike the original task configuration (Liu et al., 2025), where each task is defined by disabling single agent, our framework introduces tasks with varying joint partitioning schemes. Specifically, the six joints of the robot ('bfoot', 'bshin', 'bthigh, 'ffoot', 'fshin', 'fthigh') are grouped into different agent configurations, such as (2,2,2), (3.3). (1,2,3), or (1,1,4), where each tuple represents the number of joints observable and controllable by each agent. The hyperparameter 'agent obsk', which specifies how far agents can observe in terms of connection distance, is set to 1. Models are trained using multiple source partitions and evaluated on previously unseen configurations without relying on additional interaction data. Further implementation details are provided in Tables 19
We generated the MAMuJoCo offline datasets using HAPPO (ubba et al., 2022) collecting 100 trajectories for each task. The average return across all tasks are summarized in Table 20
In our setting each agent has different observation dimension across tasks, which requires obser- vation decomposition similar to SMAC. A single joint in HalfCheetah provides a two dimensional observation consisting of its qpos and qvel values. Therefore the observation is segmented in multiples of two. For example if an agent observes a 10 dimensional vector it is decomposed into five tokens represented as (2,2,2,2.2). The first tokens up to the number of joints assigned to the agent are treated as the agent's ' own observations and the remaining tokens correspond to observations of other agents.
Using this tokenization scheme we train STAIRS with the TD3+BC algorithm (Fujimoto & Gu 202 I) for one million timesteps. We compare our approach with two baselines UpDeT (Hu et al. 2021 combined with TD3+BC and ODIS (Zhang et al.) 2023). We do not include HISSD (Liu et al. 2025) in comparison due to the complexity of its architecture which relies multiple transformer modules for skill and action extraction. The hyperparameters used in the MAMuJoCo benchmark are the same as those used in SMAC.
Table 19: Descriptions of 'HalfCheetah'
Task type
Task
Number of Agents
Observation Space
Action Space
(3,3)
2
[(8.), (8.)] [(6.). (10.), (8.)
[(3.), (3.)
Source
(2,2,2)
3
[(2.). (2.). (2.)]
11,1,1,1,1.I)
6
[(4.). (6.). (6.). (4.). (6.). (6.)
[(1.] x 6
(6)
1
[(12.)]
[(6.)]
(2,4)
2
((6.), (10.)
((2.), (4.)
(1,2,3)
3
[(4,), (8.). (8.)
I(1), (2.). (3.)
Unseen
(1,1,4)
3
[(4.). (10.)
[(1.). (1.). (4.))
(1,1,2,2)
4
[(4.), (10.), (8.)
I(.)). (1.), (2.), (2.)
11,1,1,3)
4
[(4.), (6.). (6.). (8.)
[(1), (L), (3.1
31

---

Published as a conference paper at ICLR 2026
Table 20: Properties of offline datasets on "HallCheetah"
Task
Quality
''rajectories
Average return
(3,3)
medium
100
5043.32
(2,2,2)
medium
100
5074.4
(1III.I1III)
medium
100
4076.55
The results for the complete 'HalfCheetah' task suite are presented in Table [1] Our method achieves substantial performance improvements across all tasks and improves performance by 129% over ODIS.
Table 21: Comparison of average and per-task performances the HalfCheetah task set in MAMu- JoCo. We report mean standard deviation, with the best result shown bold.
HalfCheetah
Tasks
UPD:T-BC
ODIS
STAIRS (Ours)
Source Tasks
(3,3)
148.2 + 307.9
970.4 416.9
1459.0 + 400.3
(2.2.2)
-66.4 124.1
537.8 318.8
1410.6 537.6
(1.1I.111.))
262.6 + 200.7
727.6 662.0
1006.0 + 420.4
Unseen Tasks
(6)
-990.9 + 296.5
-18.1 167.2
256.7 297.9
(2,4)
-38.2 203.9
137.3 + 77.2
249.1 151.7
(1.2.3)
127.6 189.4
48.7 130.5
627.2 732.6
(1.1.4)
-104.4 258.4
25.1
141.4 1: 31.3
(1.1,2,2)
-17.00 96.4
222.9
399.6 + 222.9
1078.4 + 849.5
(1,1,13)
2.3 1: 252.2
178.9 271.1
606.2 700.6
Terran Avg
-6.1
331.4
759.4
32

---

Published as a conference paper at ICLR 2026
L ATTENTION SHARPENING
To test whether simple attention sharpening can mitigate the "uniform attention" behavior observed in other baselines, we performed ablation study that modifies the softmax temperature in the attention module. Specifically, we adjust the attention computation as:
QK TVd=
Attn(Q, Q,K, K.K,)) softmax
V.
where smaller values of produce sharper attention distributions. We compare models trained with 7 1.0, 0.5, 0.1, representing progressively stronger sharpening. As shown in Figure T6 and Table 22 mild sharpening provides slight improvements over the baseline. However, applying excessive sharpening (i.e., using small leads to notable performance degradation. When the temperature becomes too low, the attention distribution approaches nearly deterministic selection, preventing the model from flexibly capturing relationships among tokens and ultimately reducing overall performance.
Performence ice with on starpening
Table 22: Performance with Bold indicate the best performance among the sharpening variants (excluding ours).
STMR5 1hamd)
Datasel
1.0
7=0.5
Marine-Hard
Expert
26.5
20.7
26.1
68.4
Modium
20.3
25.4
22.4
$7.9
Medium- Expeerr
20.3
18.7
16.6
63.9
Medium Replay
17.3
19.0
15.8
59.7
Stalker-Zealot
Expert
19.9
24.4
24.2
75.0
Mediam
16.4
15.3
38.2
Medium Expert
69.4
Sespenstvr
Modiam- Reallm
3,6
15.7
24.3
Marine-Easy
Figure 16: Average test win rate across all tasks (marine-hard, stalker-zealot, and marine-easy) under different temperatures T.
Expert
40.6
Medium
84.1
Medium-Expee
43.5
53.7
12.9
92.7
Medium-Replay Replay
Mediumn
2.7
4.4
2.0
76.5
In addition, to understand why simple attention sharpening cannot achieve the performance of STAIRS, examine the attention maps produced under different temperature settings. As shown in Figure 17 decreasing (i.e., applying stronger sharpening) causes the model to place increasingly higher attention the history token in both the seen (3m) and unseen (4m) tasks. At first glance, this tendency might appear desirable, since attending to history can help mitigate partial observability.
However, when we visualize the attention maps (Figure [8], different pattern emerges: with strong sharpening, the model attends almost exclusively to the history token at every timestep. In contrast, STAIRS attends to history only when necessary; depending on the situation, may instead focus on enemy tokens, ally tokens, or history tokens. This adaptive behavior enables more effective reasoning under partial observability.
These results reveal that forcing the model to focus on the history token at all timesteps is detrimental. Effective policies require situation- aware attention, not uniformly sharpened attention distributions.
33

---

Published as a conference paper at ICLR 2026
Seen (3m)
Unseen (4m)
7=0.5
Sharper
T=0.1
Figure 17: Average attention maps of models trained the Marine-Hard-MM task, evaluated the seen (3m) and unseen (4m) tasks across the entire trajectory for different values of r.
T=1
T-0.5
T=01
Sharper
Figure 18: Attention maps of models trained on the Marine-Hann Mard-Medium task, evaluated on the 3m task and sampled every 2-3 timesteps for different values r.
34

---

Published as a conference paper at ICLR 2026
M COMPARISON WITH SAME DEPTH (2-LAYER TRANSFORMERS)
Since STAIRS employs hierarchical spatial structure, we additionally compare all baseline methods under comparable number of transformer parameters. Specifically, reconfigure each baseline (UpDeT-m, ODIS, and HiSSD) to use a 2-layer transformer, matching the transformer-level parameter count of our model and ensuring a fair architectural comparison. (For reference, the total parameter counts are: UpDeTTmm 79,095; ODIS 138,573; HiSSD 679,335; and our model 220,023.)
Across the tables 23 24 and below, we observe that increasing the transformer depth does not consistently improve baseline performance. In fact, performance degradation is observed in both the Marine-hard and Stalker-Zealot benchmarks. For Marine-hard, performance decreases by 2.5 % (UpDeTTmm), 14.5 % (ODIS), and 2.2% (HiSSD). For Stalker-Zealot, the degradation is even more pronounced: 2.9 (UpDeT-m), 21.5% (ODIS), and 3.2% (HiSSD), Only in the Marine-easy benchmark does deeper architecture provide improvements: UpDeT-m increases by 29.2%. ODIS by 17.1% % and HiSSD by 4.27%. %.
Even when all methods use deeper transformer backbones, STAIRSFormer consistently outperforms all baselines across every task group. With 2-layer transformers, the improvements are substantial:
Marine-hard: +203.6% vs. UpDe'T-m, +104% vs. ODIS, +14.6% vs. HiSSD
Stalker-Zealot: +248. 9% vs. UpDeT-m, +160.6% vs. ODIS, +53.5% vs. HiSSD
Marine-easy: +99.2% vs. UpDeT-m, +50.7% vs. ODIS, +3.8% vs. HiSSD
Overall, these results demonstrate that simply increasing transformer depth does not close the performance gap for existing baselines, while STAIRSFormer continues to provide strong gains, highlighting that its advantages arise from its architectural design rather than depth alone.
35

---

(Ours) Published as a conference paper at ICLR 2026
Table 23: Comparison of average and per-task performances on the Marine-harte task set across four dataset qualitieswith all transformer backbones using depth 2. We report standrdd deviation, with the best shown in bold.
Expert
Medium
Tasks
UPDeT.m
(Ouns)
Source Tasks
3m
78.8
91.1 10.9
99.1 1.4
99. 4 = 144
41.9 29.0
61.9 + 18.5
62.5 + 10.4
84. 44
5m6m
5.0 4.7
38.1 13.1
72.5 6.8
70.6 10.5
5.0 11.2
28.1 + 9.1
33.8 13.3
50.0 12.5
9m 10m
18.1 + 19.3
71.9+ 15.9
97.5 1.4
99.4 + 1.4
15.0 16.4
51.9
- 25.1
68.8 13.4
86.9
Unseen Tasks
4m
47.5
27.0
76.9 29.6
100.0 0.0
97.5
39.1 23.7
65.6 21.8
72.5 10.9
89.4 13.9
5m
88.8 1: 11.2
86.9 13.0
100.0 0.0
100.0 0.0
81. 21.7
86.3 1: 21.4
90.6 15.8
100.0 0.0
10m
40.0 42.2
61.3 36.3
95.0 11.2
100.0
45.6 31.3
50.0 29.1
87.5 9.1
97.5 4.1
12m
12.5 26.2
18.1:281
48.1 36.7
99.4#1 1.4
20.0 23.0
31.9 19.7
88.8 2.8
95.6 288
7m88
13#
8.8 6.4
32.5 15.4
25.0
9.
7,5 1 11.8
8.1 118
10.6 8.7
8m9m
2.5
15.0
40.0 20.9
35.6 14.8
0.6 1.4
4.4 2.8
8.1 3.1+-.111
15.6 8.0
10ml Im
8.8 12.8
15.6 15.5
+
87.5 4.9
6.4
10.6
28.8-.22
61.3 18.2
10m 2m
0.0:0.00
0.0
15.6 18.6
5.6 7.5
0.0
0.0
0.0
1.3 1.7
13m1 55
0.0
0.0 0.0
2.4
0.6 1.4
0.0
0.0
0.0
1.9 1.9
Avg
25.3
40.6
64.6
68.4
21.7
33.2
45.8
57.9
Tasks
Eedium-Expert
Medium-Repla
Source Tasks
3m
48.1 1 33.8
81.9
88.8 25.2
98.8 1.7
40.6 29.0
63.8 L. 29.4
84.4
78.1 17.1
5móm
3.8
18.1 19.1
40.6 26.0
57.5 13.9
0.0
1 7.2
30.0 9.3
50.6 5.1
9m 10m
5.0 + 6.5
41.3 33.3
65.0 23.6
94.4 4.1
4:70
8.8 12.0
41.9 23.1
78.1 16.1
Unseen Tasks
43.8
33.7
40.1
80.6
67.5
41.3
32.5
99.4
90.0 12.8
23.8 31.6
46.9
95.0 5.7
42.1
17.1
33.1 16.6
7.5
10ml
2.6
15.0+:26.5
71.3
18.1
1.9
1.9
13ml
0.0
Avg
20.5
29.4
61.9
63.9
14.8
19.3
45.7
59.7
36

---

Published as a conference paper at ICLR 2026
Table 24: Comparison of average and per-task performances on the Stalker-Zeale task set across four dataset qualitieswith all transformer backbones using depth 2. We report standard deviation, with the best shown in bold.
Expert
Medium
Tasks
LPDeT-m
ODIS
HESSD
STAIRS (Ours)
UPDeT-m
ODIS
HISSD
STAIRS (Ours)
Source Tasks
283z
33.8 11.1
70.0 38.6
92.5
95.6 + 5.2
25.0 12.1
43.1 19.2
39.1 12.1
56.9 = 10.5
294z
13.8 = 113
58.8 37.3
6.1
77.5 = 11.6
25.0 20.6
7,3 3.6
9.4
60.0 = 16.1
3:52
28.1
18.12215
66.9 33.0
88.8 + 5.7
87.3 10.6
20.6 12.2
24.4 =:
12.0 26.8
52.5 3.4
Unseen Tasks
1832
13.8 20.6
36.0 + 37.7
63.8 19.1
78.1 + 12.7
22.5 + 10.5
5.0 7.8
25.6 + 27.7
38.8 1 34.0
1s4z
4.4
1.4+60 5.2
21.9
26.4
41.3
19.3
76.3 2 21.0
20.6 + 21.6
1.9 2.8
6.9 12.2
25.6 + 9.7
1552
2.5 4.1
9.4 12.9
20.6
141
14.1
55.6
23.5
11.9 10.2
0.0 1.0::00 0.0
3.8 1.26
31.9 + 10.5
2652
7.5 1.00
42.5
92.9
78.8 17.0
84.4 4: 7.0
16.9 14.8
8.1
11.4
15.6 10.6
25.6 +: 8.7
3532
20.6
1: 21.0
58.8
35.2
74.4 7.1
86.3 1 8.4
18.8 1 18.1
26.9
13.7
25.6 1: 17.0
59.4 1: 14.1
3s4z.
24.4 + 26.0
65.0
38.2
83.8 8.1
92.5 3.6
3.6
32.5 + 14.1
47.5
23.4
20.4
12.6
59.4 1 24.7
4632
21.3
2-.2
46.9
31.3
81.3
12.1
70.0 11.8
11.9 13.9
22.5 + 14.4
21.9 11.0
41.9 - 17.9
4s4z
15.0 + 15.2
28.1 24.7
68.8
16.8
68.1 20.8
10.0 11.1
8.1 +.1
13.1
21.3
18.0
465z
9.4
13.3
16.3 17.2
24.7 40.6
53.1 = 18.9
4.6
0,6 1.4
5.0
11.3 7.8
4662
3.8
1. 5.6
9.4 11.7
35.6
22.6
59.4 1: 17.5
1.3 1. 1.7
0.6 1: 1.4
1.9 1.91228
11.9 5.6
Avg
15.3
40.T
64.3
75.0
17.1
15.1
17.3
38.2
Tasks
Medium- Expert
Medium-Replay
Source Tasks
253z
35.0 9.2
41.3 26.6
78.8
92.5 = 10.3
16.3
10.0 13.7
7.5
20.6 10.0
264z
33.8 10.0
21.3
25.9
74.4 85.0 15.8
9.8
8.1
5.0
28.8 15.8
365z
16.9 20.0
34.4 30.0
58.8 24.8
0.0 0.0:+:0.0
5.6 5.1
11.3
28.8 10.2
Unseen Tasks
1s3x
27.5 22.0
21.3 32.0
73.8 28.9
63.1 15.2
32.5 1 35.3
5.4
39.4 38.2
12.5 14.5
194z
114 9.5
1.9422
5.0 6.5
80.6 2 21.8
18.8 21.6
8.1 11.2
7.7
10.6
1s52
6.3
0.3:50 5.8
1.9288
2.
51.9 - 32.9
11.3 11.
1.9 2.
7.5 10.5
23.1 + 36.3
2552
14.4
1.442:10 12.2
25.0 21.9
8.1
62.5 21.2
10.8
8.8 1: 13.7
7.7
27.5 - 11.4
3:32
23.8
19.3
213
21.
85.0 6.0
81.9 11.6
10.0 1: 13.9
6.9 13.7
10.0 14.2
56.3 - 15.9
3,4,
26.9
19.1
41.3
37.8
74.4
95.6 +: 4.2
5.6 166115.5 7.8
11.9 11.1
21.9
14.5
53.1 + 10.4
463z
20.0
33.4
20.3 11.3
51.3
61.3 15.7
3.8 1.8.1.8.4
7.5 16.8
23.1
24.1
28.1
20.4
454z
4.4
4.6 3.6
5.6-+-95
30.6 15.8
59.4 14.3
0.0 1.00000.0
4.4
10.6 10.6++8.11
15.0 2.6
4s52
5.7
1.711
9.4 9.1
53.8
21.7
1.9 4.2
1.9
8.8
3.8
4s64
1.717
1.7
6.9.1.4. 4.1
40.0 15.5
0.6
0.0
3.0
7.5 1 6.8
Avg
18.0
17.6
8.9
12.7
37

---

Published as a conference paper at ICLR 2026
Table 25: Comparison of average and per-task performances on the Marine easy task set across four dataset qualitieswith all transformer backbones using depth 2. We report stadaard deviation, with the best shown in bold.
Expert
Medium
UPDeT-m
ODIS
HiSSD
STAIRS (Ours)
UPDeT-m
ODIS
HiSDD
STAIRS (Ours)
Source Tasks
71.9
70.6.15.77
97.5
0.0
7.66
914
Unseen Tasks
4m
46.3 22.4
69.4 27.7
95.6 1.6...5..
96.9
48.1 24.5
71.9 24.7
65.6 18.6
73.8 1: 13.4
6m
51.9 36.6
18.8 85.6
100.0 0.0
9.89
72.5 = 12.8
91.3 10.5
819 17.3
82.5 9.3
7m
48.1
31.9
75.6 31.1
98
100.0 0.0
81.9 19.4
94.4
86.9 18.1
98.1 4.2
8m
60.0
+0.05:24.55
91.9 11.8
99.4 1.4
99.4 1.4
83.1 = 12.6
95.0
96.9 5.4
96.9 3.1
9m
50.0 18.4
98.8
100.0
100.0 0.0
58.1 21.7
85.0 7.5
80.6
93.1 + 5.1
11m
58.1 23.9
96.3
99.1 144
100.0 = 0.0
30.6 = 10.7
13.1 = 16.7
5.88
65.6 14.5
12m
50.0 + 19.1
89.4 12.2
98.8 1 1.7
98.1 1.7
17.5 16.9
30.6
16.0
42.5 7.2
65.6 6.6
Avg
53.9
89.7
99.2
99.0
58.5
69.6
73.2
84.1
Medium-Experi
Medium-Repla
Tasks
UPDeT-m
ODIS
HiSSD
STAIRS (Ours)
UPDeT-m
ODIS
HiSSD
STAIRS (Ours)
Source Tasks
3m
47.5 37.2
53.1 1 20.1
90.6 7.7
98.8 1.7
45.6 = 23.8
61.9 1: 36.6
88.8 : 2.8
86.9 6.8
5m
81.3 23.5
77.5 21.6
100.0 1:
98.8 1.7
0.0 1 0.0
21.9 30.0
90.6 1 7.3
89.4 7.8
10m
78.1 23.5
9.2
91.9 14.8
100.0 + 0.0
0.0 0.0
0.0 0.0
88.1 + 7.5
56.9 + 18.7
Unseen Tasks
4m
48.8
+ 11.0
58.8 23.1
95.0 +: 4.2
60.0 25.1
0.0 = 0.0
22.5 30.4
70.6 5.2
79.4 + 13.0
6m
16.9
11.3
15.6
29.1
86.3 16.2
94.1 1.6
0.0 0.0
18.8 10.2
99.4
91.3 6.1
7m
67.5
18.0
58.8
41.2
814 13.6
96.9 3.1
0.0 0.0
20.0 11.7
100.0
90.6 5.8
8m
77.5
+ 14.6
36.00
77.5 11.6
84.4 16.8
0.0 - 0.0
3.1 7.0
96.3 1.4
$3.1 8.1
9m
51.3
19.0
61.9
17.2
73.1 15.1
100.0 0.0
0.0
13
87.5 6.3
$2.5 5.7
11m
61.9
1..11.20 21.9
58.5
20.1
82.5 19.5
98.1 1.7
0.0 0.0
191
54.4 9.3
55.0 1: 23.7
12m
38.1
4: 23.5
39.4
20.9
63.8 4: 17.9
95.6 1: 4.7
0.0 0.0
1.9 4.2
48.1 4 14.1
49.4 1: 30.0
Avg
59.9
59.2
84.5
92.7
4.6
15.3
82.4
76.5
38

---

Published as a conference paper at ICLR 2026
N ABLATION: ADDING SIMPLE GRU TOKEN
To examine whether the performance gain is merely from adding a recurrent GRU cell rather than our STAIRS design, we additionally evaluated baselines with simple GRU history token. Specifically, we appended an additional history token that passes through a GRU cell operating on 3-step temporal interval, which is identical to the interval used in our method.
The comparison results are summarized in Table 26 The table reports the average test win rate across all datasets. For example. for the Marine-Hard task, we average performance across Expert, Medium, Medium-Exper and Medium-Replay. As shown, incorporating GRU does not consistently improve either UpDeT-m UpDe'T-be, indicating that simply extending the temporal horizon is insufficient.
Table Average performance comparison of GRU addition.
Task Dataset
UpDeTTmm
UpDeTT-m GRU
UpDeTT-bc
UpDeT-bc GRU
Ours
Marine-Han
21.1
20.7
46.6
49.8
62.5
Stalker Stalker-Zealot
15.3
16.7
43.3
42.9
51.7
Marine-Easy 34.2
31.6
82.1
85.7
88.1
Moreover, we visualized the average attention maps over the entire trajectories. We observed that adding the GRU history token does not encourage the model to attend to either short local history or long-range GRU-based history. other words, the temporal cue introduced by the GRU largely ignored and fails to help the baselines integrate temporal structure effectively.
These results suggest that the performance gain of our method comes from the synergistic effect of its three components, Spatial Recursive Module, Temporal Module, and Token-Dropout mechanism, rather than the inclusion of recurrent GRU cell alone.
Seen (3m)
Unseen (4m)
UpDet-m
0.2
UpDet-m + GRU
0.1
Figure 19: Average attention maps of models trained the Marine-Harde Mard-Mediimm task, evaluated on the seen (3m) and unseen (4m) tasks. UpDeT-m (upper) vs UpDeT-m with an added GRU history token (lower).
39


Reviews:
[{"review_id": "TkilLhRTlv", "reviewer": "Reviewer_u6gB", "summary": "The main challenge in offline multi-agent reinforcement learning (MARL) with multi-task (MT) datasets is the varying number of agents across tasks, which changes the input structure. Prior transformer and hierarchical skill methods underutilized the transformer's attention mechanism, focusing instead on transferable skills, and crucially, they suffered from poor historical context: they compressed the entire history into a single token at each step, making them function like a basic recurrent neural network that largely ignores long-term historical information despite its criticality in partially observable MARL. The proposed STAIRS-Former addresses this by augmenting the transformer with spatial and temporal hierarchies to effectively leverage long history and properly attend to critical tokens, while a new token dropout technique is incorporated to improve generalization to diverse agent populations; experiments on the StarCraft Multi-Agent Challenge (SMAC) benchmark confirm that STAIRS-Former achieves new state-of-the-art performance.", "questions": "- What are the limitations of this work?\n- What is the memory footprint of the model and the introduced overhead compared to the other baselines?\n- In Figure 6, what is \"(\"wo RT\" excludes repeat & TSFFN)\"?\n- I did not understand well the analysis in Figure 5. Would you mind elaborating more and clarify the heatmaps in the figure?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- I believe this paper is well-structured and written.\n- The central problem of this work is well-motivated and does sound.\n- I think this work offers a good solution for the realization of MT-MARL problem with the corresponding challenges. \n- Another strength point is the ablation on the algorithmic decision by the authors and the informative discussion.", "weaknesses": "- A crucial weakness of this work is not stating the limitations.\n- I believe a limitation of this work could be the potential overhead and memory footprint due to the introduced components. Although, the overall training time is highlighted in the appendix, a deeper analysis would be appreciated where the training time or process time for each introduced component. This can be done by reporting the ablated training time if available or simply the overhead proccessing time compared to normal training step.\n- In the experimental section, there is no highlighting for the model sizes used across methods. This could result in an unfair comparison to the baselines.\n- Figure 5 is unclear. A more informative caption would be appreciated.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7EwT0cx7pb", "reviewer": "Reviewer_7K4n", "summary": "The paper targets offline multi-task multi-agent RL (MT-MARL) with varying agent counts, inputs and actions across tasks. It builds on UPDeT/transformer-style architectures and proposes (i) a “spatial recursive”/deeper transformer to get less uniform attention, (ii) a dual-timescale temporal/history module (short- and long-term), and (iii) token dropout to generalize across different token/entity counts. Experiments on SMAC multi-task offline datasets show improvements over UPDeT-m, ODIS, and HiSSD.", "questions": "1. Offline pretrained transformer-based MARL (MADT) show that one big sequence model can handle multiple SMAC tasks and benefit from offline pretraining. How does this method improve upon MADT and similar baselines?\n2. For offline MARL, optimizing just the TD3-loss with BC regulation has shown to yield poor results because of very weak regularizations on the exploding joint action spaces. How has that been tackled here? Why did the authors use this method of training over existing offline MARL framework? \n3. For the comparisons with other methods, how many layers were used for the baselines vs the STAIRS?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. Problem setting (offline + multi-task + variable agents) is relevant for MT-MARL and aligns with recent transformer-based MARL lines.\n2. Paper is clearly written and well-situated w.r.t. UPDeT/ODIS/HiSSD.", "weaknesses": "1. Overall, the paper integrates known ingredients rather than introducing a genuinely new architectural principle for offline MT-MARL.\n2. If the problem is “uniform attention” in UPDeT/HiSSD, there are other mechanisms: attention sharpening, entropy regularization, auxiliary supervision on heads, or stronger positional/task conditioning. The paper should justify why a relatively heavy spatial–temporal–recursive stack is preferable to these lighter alternatives.\n3. Baselines may be underpowered: The main UPDeT-style baselines in the paper use very shallow transformers (as the authors themselves note “one-layer transformer cannot capture diverse relations”). A fairer test is: what happens if we (i) increase the number of transformer layers, (ii) add a simple recurrent/history token with longer horizon. Right now, the improvement could just be due to “more depth + a GRUi.e., model capacity, not the specific STAIRS interleaving.\n4. Experiments on other benchmarks such as MAMuJoCo, WareHouse, etc. can also be presented to enhance the experimental evaluations to compare how this method compares against other offline MARL baselines.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oa6Tdz9Pc5", "reviewer": "Reviewer_7cme", "summary": "The paper introduces STAIRS-Former, a spatio-temporal transformer architecture for offline multi-task multi-agent reinforcement learning (MT-MARL). The proposed method includes (1) a novel spatial recursive model to extract correlations among local observations of different entities, and (2) a novel temporal module that helps mitigating partial-observability in MARL settings and allows for capturing long-term dependencies.\nThe authors carry out experiments on offline SMAC v1 datasets in multi-task fashion, showing improved performance over previous baselines. Particularly, STAIRS-Former displays impressive generalization over unseen tasks, as well as varying number of agents, likely due to the token dropout mechanism.", "questions": "- Why was SMACv1 chosen by the authors against the improved SMACv2 benchmark?\n- Please clarify how the tasks are divided into training and tests, with respect to number of agents and different goals/rewards. It appears to me that the authors do not rely on explicit task-conditioning information at training time, so I'm assuming agents must implicitly infer the task from observations. If so, how can they generalize to a task with a different objective? Or do tasks only differ by the amount of agents?\n- The authors claim in the abstract that a Transformer module is not able to capture long-range dependencies  because it compresses the entire history into a single token. However, this is in general a false claim, because that's exactly what transformers claim to do over RNNs. Could you please clarify what is the main drawback of previous methods and whether the limitation over long-horizons is a fundamental consequence of the architecture itself or rather a resulting effect?", "limitations": "", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The authors tackle a relevant and underexplored setting (offline MT-MARL) which is likely of interest to the community and opens important directions of application.\n- The paper shows clear improvements and strong empirical performance over relevant baselines in the field (UPDeT-m, ODIS, HiSSD), specifically over unseen tasks.", "weaknesses": "- Limited and outdated benchmark tasks: the authors solely present their comparison in the context of the SMAC benchmark. The authors neglected experimentation on more recent benchmarks such as the improved benchmark SMACv2 [1], as well as the MaMuJoCo benchmark [2]. Considering the nature of this work is mostly empirical, and that SOTA online MARL methods notoriously test on these benchmark, it is unclear why the authors only provide experimentations on SMAC v1. In turn, it's unclear how the architectural contributions in STAIRS-Former really compare against relevant benchmark tasks in the field.\n\n[1] Ellis, Benjamin, et al. \"Smacv2: An improved benchmark for cooperative multi-agent reinforcement learning.\" Advances in Neural Information Processing Systems 36 (2023): 37567-37593.\n\n[2] Peng, Bei, et al. \"Facmac: Factored multi-agent centralised policy gradients.\" Advances in Neural Information Processing Systems 34 (2021): 12208-12221.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SIhOLtHjqE", "reviewer": "Reviewer_PLFk", "summary": "The paper introduces STAIRS-Former, a transformer-based architecture for offline multi-task multi-agent reinforcement learning (MT-MARL). It addresses two main limitations in prior transformer-based MARL models (like UPDeT, ODIS, and HiSSD): poor handling of long-term temporal dependencies and limited relational reasoning among entities. STAIRS-Former introduces three modules — (1) a spatial recursive transformer for deeper inter-agent correlation modeling, (2) a dual-scale temporal module that maintains short- and long-term histories, and (3) a token-dropout mechanism to improve robustness across varying numbers of agents. Extensive experiments on SMAC benchmarks (Marine-Easy, Marine-Hard, and Stalker-Zealot) show consistent and significant improvements over state-of-the-art baselines, with ablation studies and interpretability analyses (e.g., attention and dormant neuron studies) supporting the architectural design", "questions": "Generality beyond SMAC: Have you tested or considered applying STAIRS-Former to environments with more complex or unstructured observations, such as multi-modal inputs (e.g., visual or continuous sensor data)? If not, how do you anticipate the model’s spatial recursion and tokenization scheme would adapt in such cases?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "This paper makes a clear and well-motivated contribution to the field of offline multi-task multi-agent reinforcement learning. The authors identify real and important limitations in existing transformer-based MARL models—specifically their difficulty in capturing long-term temporal relations and rich inter-agent dependencies—and address them with a framework that feels both technically sound and intuitively designed. The proposed STAIRS-Former architecture is an elegant combination of spatial recursion and dual-scale temporal modeling, allowing the system to reason jointly over history and agent relations in a way that existing models cannot.\n\nThe experimental section is particularly strong: the evaluation on multiple SMAC benchmarks is thorough, ablation studies are comprehensive, and the visualizations provide genuine insight into how the model learns. The paper is also very well written and organized, making it easy to follow both the motivation and the technical details. Finally, the work feels significant because it pushes transformer-based MARL toward more scalable and generalizable architectures, offering a practical foundation for future research in multi-agent decision-making systems.", "weaknesses": "While the framework is strong and the results are convincing, the paper would benefit from a clearer discussion of generalization beyond the SMAC environment. SMAC’s discrete and tokenizable observation space makes it naturally suited to transformer architectures, so it’s uncertain whether STAIRS-Former would maintain its advantages in less structured or multi-modal domains (e.g., visual-linguistic inputs or real-world sensor data). A brief evaluation or qualitative analysis in such settings would significantly strengthen the paper’s claim to generality.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)