Sable: a Performant, Efficient and Scalable Sequence Model for MARL

Thank you for the feedback. Your comments on the retention equations and other aspects of our work have helped us identify areas for improvement. We address your questions and comments below.

wondering if agents share parameteres, espicially on Q, V, K matrices

Sable uses a single network for all agents. The QVK matrices are not unique per agent, but shared in the sense that each agent is treated as another element in the sequence. While longer sequences incur higher memory usage, RetNets address this by constraining computational memory to a fixed chunk size.

the effect when the chunk size and the number of agents increase jointly

The chunk size can be increased as the number of agents increases, but this will affect memory requirements. The optimal setting is to process the entire training sequence at once, but if this isn't feasible, Sable allows the chunk size to be tuned to maximize hardware utilization and enable training on arbitrarily large sequences.

which decomposition theorem are you referring to

Please see our reply to D22g marked (**).

agents' actions are dependent in a order. How would \hat{h}_0 be?

As shown in Line 177 Column 2, $\hat{h}\_{0}$ is $h\_{t-1}^{dec}$, where $h_{t-1}^{dec}$ is the decayed hidden state from the previous timestep. $\hat{h}$ is an intermediary variable that accumulates the hidden state over agents within a single timestep. This is decayed once per timestep to produce $h_{t}^{dec}$. Regarding the agents’ actions order dependency, we mitigate any potential bias by shuffling agent order during training.

it is unclear about the \tau and \tau_prev using \tau to refer to a chunk can be confused

(***) Thank you for going through our paper in such detail, we will make the following changes to the paper to help with the clarity of our method.

For the chunkwise representation, we split a trajectory $\tau$ consisting of $L$ timesteps, with $N$ agents gathered during inference into smaller chunks each of length $C$ such that the retention for chunk $i$ can be given as:

$$ \begin{aligned}Q_{[ \tau_{i} ]} &= Q_{C(i-1):Ci}, K_{[\tau_i]} = K_{C(i-1):Ci}, V_{[\tau_i]} = V_{C(i-1):Ci} \\h_i &= K^T_{[\tau_i]} \left( V_{[\tau_i]} \odot \zeta \right) + \delta \kappa^{ \lfloor L/C \rfloor } h_{i-1}, \zeta = D_{N \cdot \lfloor L/C \rfloor, 1:N \cdot \lfloor L/C \rfloor} \\\text{Ret}(\boldsymbol x_{[\tau_i]}) &= \left( Q_{[\tau_i]} K^T_{[\tau_i]} \odot D \right) V_{[\tau_i]} + \left( Q_{[\tau_i]} h_i \right) \odot \xi \\\text{where } \ & \xi_{j} = \begin{cases} \kappa^{\left\lfloor j / N \right\rfloor + 1}, & \text{if } j\leq Nt_{d_0} \\ 0, & \text{if } j > Nt_{d_0} \\\end{cases}.\end{aligned} $$

Here $h$ is an intermediary variable carrying information from one chunk to the next, $h_0$ is the hidden state at the beginning of $\tau$ that will be used for training, $\zeta$ is the last row of the decay matrix that is created from the data of the chunked trajectory for chunk $i$ and $\xi$ is a duplicated column vector. Please see our answer to (C1) of reviewer qBu1.

This removes the confusing $\tau$ and $\tau_{prev}$ notation from the text and should link Equations 3 and 6 more clearly. It also removes $i$ from the definition of $\xi$ giving clarity around $Nt_{d_0}$, as $j$ is now an index in $\xi$.

it is unclear whether you refer to a chunk as a set of agents or a set of observations over time step

Sable's flexibility allows for treating either entire rollouts with multiple agents at each timestep, or just the number of agents, as the training sequence length. With E environments, L timesteps, N agents, and C chunks, the default training batch shape is $(E, NL)$, divisible into $(E, N [L / C] )$ size chunks. When using only the number of agents as the sequence length, the shape is $(EL, N)$, divisible into $(EL, N/C)$ size chunks.

in section 2 you are using continuing tasks while in equation 4 you consider epsiodes

All tasks we consider have fixed length time horizons and termination conditions and we allow for environments to automatically reset once an episode terminates. Thus, for a fixed rollout length it is possible for there to be multiple terminal timesteps.

not clear how's the decay matrix updated according to your equation

In our case, the decay matrix is blockwise lower diagonal of size $(LN, LN)$ with block size $(N, N)$ where $N$ is the number of agents. Each element is exponentially decayed given its position in time for a given trajectory which follows Equation 2. We discuss our adaptations to the decay matrix after Equation 6 in Lines 213-240 and give an example in Appendix D.

Would it be possible to use a smaller chunk size, e..g, 8, for all tasks?

This is possible, but not practical. Smaller chunks during training use less memory at the cost of wall clock time. In practice, we train with a chunk size that is as large as our computational memory permits for the fastest training wall clock time.

Thank you for your clarifying questions and feedback. We provide detailed responses below.

The claim on the scalability may be a bit untenable

RetNets' chunkwise formulation allows us to process long sequences in small chunks, scaling to arbitrarily long sequences regardless of whether they consist of many agents at a single timestep or over multiple timesteps. This scalability difference is evident in Figure 4, where MAT scales much worse than Sable.

Sable’s details are vague, it is essentially the same as MAT, both are centralized methods, the claim of "a new sequence model for MARL" is debatable

We emphasize that Sable introduces a fundamentally different sequence modeling approach compared to MAT. While both Sable and MAT are CL methods, Sable can reason temporally, which allows it to model sequences of agents over time and capture long-term dependencies, unlike MAT, which is limited to reasoning within a single timestep.

Due to the page limit, we did not have enough space to add all the method details in our initial submission but they can all be found in the appendix. We have an extra page for the camera-ready version and will move these method details into the main text to improve clarity.

What does b mean

The b subscript denotes a batch of trajectories from the buffer. We acknowledge that this notation could be made more clear, and we will update Algorithm 1 and the notation accordingly to avoid confusion.

the decoder only takes the actions as input

We are only referring to the first block of the decoder, the second block performs cross retention with the output of the first block and the encoded observations as can be seen in Figure 13.

Sable does not take the trajectories as input

This statement referred to is specifically to the scaling strategy described in Section 3 (Scaling the number of agents) for handling thousands of agents. It represents a variant of Sable optimized for extremely large agent counts, not the default implementation used in most experiments. The algorithm in Appendix D shows the full version of Sable that conditions on trajectories.

Sable does not rely on the retention

Sable still relies on retention. This statement referred to a limitation when using the agent-chunking scaling strategy described in Section 3. When chunking across a large number of agents, self-retention can only be applied per chunk, which means that agents in different chunks do not “retend” to each other. When the number of agents all fit into a single chunk, it is possible to perform full self-retention across all agents as they all fit within the same chunk.

This provides a design trade-off to the user when considering the scale of the problem at hand and the memory available. Furthermore, we would like to point out that in Figure 4, we show that even though Sable cannot perform full self-retention across all agents (only per chunk) at large scales, it still achieves the best performance, whereas MAT is not even able to fit into memory at the extreme end.

unsure what advantage decomposition theorem it refers to

(**) We are referring to the advantage decomposition theorem (ADT) originally derived in [Kuba (2022)] (there called Lemma 1). This theorem underpins the Fundamental Theorem of Heterogeneous-Agent Mirror Learning (HAML) as mentioned on Line 58 Column 2 in the introduction. We will amend the text to make the link between HAML and the ADT clear.

degradation of IPPO performance looks a bit suspicious

Please refer to our answer to a similar question from reviewer pmav marked (*).

memory usage, is it in the training or in the inference?

The memory usage reported in Figure 5 is measured during training as this is when the bulk of the memory is used.

On the usage of memory efficient transformers

Memory efficient transformers without a dual form do not maintain a hidden state, thus would need to keep a cache of per-agent/timestep observations during inference. The two most important aspects of retention in Sable for scaling is low memory requirements and the dual form for efficient inference. For a further discussion on the differences between RetNets and Transformers, we refer the reviewer to Section 2.4 in [Sun (2023)]. We are not against adding some of these related works into the text if the reviewer thinks it would be useful.

should be gradient descent for ϕ

Thank you for pointing this out, we will update it.

final configs used to produce the reported results?

The final configs, source code, and raw experimental data are available at the link at the end of Section 3. You can download the data by pressing "Experiment Data" at the top of the page and find optimal hyperparameters at All experiments data/Benchmark/optimal-hyperparams/… to reproduce our results. We chose not to add the configs directly to the appendix as it would add a large number of extra pages.

Thank you for your feedback, especially your comments on our positioning within the context of different MARL algorithms, and questions on implementation and experimental details, which helped us improve the paper. We provide detailed responses below.

Confusion from the introduction's contrast between Sable and prior MARL approaches.

We acknowledge that the introduction to Sable can be improved and will make it clear that Sable is a CL method. Our narrative was that Sable breaks the typical CL mold by being both performant, memory efficient and scalable, unlike other CL methods. We will update the introduction to convey this more clearly.

Include the training objective in Section 3 and add Figure 13 to Section 3 to visualize the network architecture.

Due to the page limit constraint, we moved some additional details and the architecture figure to the appendix. We will add these back into the main text of the updated version.

How is the Sable network and policy adjusted for continuous vs discrete action spaces, and does the optimization objective differ?

Sable's policy network uses different output heads for discrete and continuous actions, but the architecture and PPO optimisation objective remain the same. For discrete actions, the policy head outputs action logits per agent, which are used to sample actions and train the policy. For continuous actions, the policy outputs mean values and a shared log standard deviation parameter, to sample actions from a Gaussian distribution. We will add this and more details to the appendix.

in what way the retentive network architecture applied in Sable differs from Structured state space models for in-context reinforcement learning?

While both approaches share the core idea of replacing attention with a more memory-efficient mechanism to enable scalable sequence processing, they differ in their underlying architecture, and research goals. Sable relies on the cross-retention mechanism which is an extension we added to RetNets. There is no obvious analogue for this in S5. Additionally, the focus of S5 is on single-agent RL, in-context learning and meta-learning, while our work focusses on computationally efficient long context memory in MARL. If interested, we refer the reviewer to Section 2.4 of the RetNet paper [Sun (2023)] where the difference between S4 and RetNets are discussed in detail.

Why can't the multi-agent transformer architecture use observation history as inputs?

MAT’s architecture lacks a recurrent formulation and handling temporal memory with transformers is challenging [Parisotto (2019), Meng (2022)]. Although it is possible to maintain a cache at inference time for memory over the sequence, it is less scalable due to high memory requirements from maintaining a cache. Our RetNet for RL is advantageous in that it only requires a hidden state and constant memory at inference time.

Why does IPPO's performance degrade during training in the 128-agent LBF task in Figure 4?

(*) Sharing parameters makes distinguishing agents harder, and partial observability leads to non-stationarity. Sable and MAT overcome this with auto-regressive action selection, which aids coordination and non-stationarity. Hyperparameters shouldn't be the issue, as they were tuned. For additional information, we refer the reviewer to Appendix A3, Lines 808-824.

Would the authors be able to provide information on the training and inference cost of Sable in comparison to MAT, IPPO and MAPPO?

Below we show training and inference SPS in Neom.

Table 1: Training SPS

Table 2: Inference SPS

Sable is significantly faster than MAT but slower than IPPO. This is expected as Sable and MAT both use larger transformer style networks, while IPPO uses a smaller MLP. Additionally, we believe this is a more fair comparison as MAT is also a centralised learning method and the previous SOTA in MARL, while IPPO is independent and has significantly worse performance than Sable.

Would I be correct in understanding that a smaller chunk size comes at a cost of reduced training speed due to lower degrees of parallelisation?

Yes, exactly. Decreasing chunk size reduces training speed. Larger chunk sizes allow more parallel computation and faster training.

Thank you for your feedback, especially the close attention paid to our equations/notation. We provide detailed responses below.

(Q1) Do all baseline methods are trained in via centralized training?

In addition to answering the above question, we also wish to clarify a misunderstanding evidenced by the following comment:

“However, to my understanding, some of the baselines utilize partial information (in the default setting), unlike the proposed method or MAT. This information gap makes direct comparisons unfair.”

While MAT and Sable process information from all agents using a single network, they only use local observations during training and not the global state, unlike CTDE methods, e.g. MAPPO, QMIX. All methods use the same observations at inference time. We will clarify the definition of CL in the introduction. Not all baselines belong to the CL paradigm; we include baselines from IL, CTDE and CL.

(Q2) Formal problem setting

In Section 2, we define the problem setting as a decentralised-POMDP.

(Q3) Random permutation of agents and performance

We didn't investigate this, believing it's more principled to randomly permute agent order each timestep. This prevents the model from relying on specific orderings, avoiding bias [Kuba (2022)].

(Q4) Do QMIX and MAPPO still leverage partial information during decision-making, while Sable and MAT utilize global information?

Please refer to our answer in Q1.

(Q5) Dimensions of matrices

Please refer to our answer to (C1) below and to how we intend to rewrite Equation 6 in our answer to reviewer wrvH marked (***).

outdated literature/baselines on value-based methods.

A well-tuned QMIX has been shown to outperform various extensions [Hu (2023)]. For this reason, we feel that QMIX represents a sufficiently strong value-based baseline. We will also clarify this in the experiments section.

(C1) In Eq. (3), $\nu_{ij}$ and $\zeta_{ij}$ contain index $j$ but $j$ does not appear in their expressions. Perhaps, it could be expressed differently to avoid any confusion.

Since Equation 3 does not have $\xi$ or $\nu$, we assume the reviewer is referring to $\xi$ and $\zeta$ in Equation 6. We acknowledge that our presentation of Equation 6 was imprecise. This misrepresentation was inadvertently transferred from the original RetNet Paper. Both $\zeta$ and $\xi$ are matrices of the same shape as the decay matrix $D$, with dimensions $C \times C$, where $C$ is the chunk size. These matrices contain values that are constant across columns but vary across rows. We will revise Equation 6 to eliminate the ambiguity caused by the overloaded use of the index $i$, and we will explicitly define the role of $j$ to avoid confusion. We refer the reviewer to our response to reviewer wrvH marked (***) for an overview of how we intend to update Equation 6.

(C2) In training part, the corresponding loss functions and algorithm presented in Appendix should be mentioned for readers to refer to them.

Due to the page limit, we did not have enough space to add this in our initial submission (moving it to the appendix), but since we have an extra page for the camera-ready version, we will add the loss function back to the main text.

(C3) Some pictorial illustration would be helpful for readers to understand Neom, a newly introduced task, if possible.

Thank you for the suggestion, we will add a render of a step in Neom.

major contribution is replacing the attention mechanism in MAT with the retention.

Indeed, the reviewer is correct that this is our main contribution. However, we do not see it as a weakness of our work. The use of retention in Sable goes beyond a straightforward replacement of attention in MAT. To get it to work, we had to change several aspects of the original retention mechanism including:

• Introducing a reset mechanism within the decay matrix to ensure that memory is retained within episodes and not across their termination boundaries.
• Carefully control the decay over timesteps, which, unlike the original RetNet’s decay that operates only over token positions, has to handle multiple tokens/observations in each timestep.
• Developing a cross-retention mechanism, a retentive encoder and an encoder-decoder RetNet, none of which are part of the original RetNet design and are also not straightforward implementations.

Therefore, Sable as a working retention-based sequence model for RL, is a highly non-trivial algorithmic implementation. This should also be clear when comparing our code with the original implementation of RetNets and/or MAT. Additionally, retention enables Sable to attend over entire trajectories, which is impossible in MAT and is the main reason for Sable's impressive performance. We are excited about what Sable is capable of, with extensive empirical evidence giving such a strong signal for its potential use in applications.

|-----|-----|-----|-----|-----|-----|-----|
|  #  |  #  |  #  |  #  |  #  |  #  |  #  |
|  #  |  .  | A2  |  .  |  .  |  G  |  #  |
|  #  |  .  |  #  |  .  |  .  |  .  |  #  |
|  #  |  .  |  .  |  .  |  .  |  .  |  #  |
|  #  | A1  |  .  |  .  |  #  |  .  |  #  |
|  #  |  #  |  #  |  #  |  #  |  #  |  #  |
|-----|-----|-----|-----|-----|-----|-----|

  #  |  .  |  .  
  #  | A1  |  . 
  #  |  #  |  # 

 #  |  # | #
 .  | A2 | .
 .  |  # | .