Trajectory-Class-Aware Multi-Agent Reinforcement Learning
Authors: Hyungho Na
Kwanghyeon Lee
Sumin Lee
Il-chul Moon
Venue: ICLR 2025 Poster (2025)
Keywords: trajectory clustering
multi-agent reinforcement learning
trajectory-class-aware policy
multi-task
Primary area: reinforcement learning
TL;DR: TRAMA enables agents to recognize task types by identifying the class of trajectories and to use this information for action policy.

Abstract: In the context of multi-agent reinforcement learning, *generalization* is a challenge to solve various tasks that may require different joint policies or coordination without relying on policies specialized for each task. We refer to this type of problem as a *multi-task*, and we train agents to be versatile in this multi-task setting through a single training process. To address this challenge, we introduce TRajectory-class-Aware Multi-Agent reinforcement learning (TRAMA). In TRAMA, agents recognize a task type by identifying the class of trajectories they are experiencing through partial observations, and the agents use this trajectory awareness or prediction as additional information for action policy. To this end, we introduce three primary objectives in TRAMA: (a) constructing a quantized latent space to generate trajectory embeddings that reflect key similarities among them; (b) conducting trajectory clustering using these trajectory embeddings; and (c) building a trajectory-class-aware policy. Specifically for (c), we introduce a trajectory-class predictor that performs agent-wise predictions on the trajectory class; and we design a trajectory-class representation model for each trajectory class. Each agent takes actions based on this trajectory-class representation along with its partial observation for task-aware execution. The proposed method is evaluated on various tasks, including multi-task problems built upon StarCraft II. Empirical results show further performance improvements over state-of-the-art baselines.

Published as a conference paper at ICLR 2025
TRAJECTORY-CIASS-AAAR MULTI-AGENT REIN- FORCEMENT LEARNING
Hyungho Na', Kwanghyeon Lee', Sumin Lee ' I1-Chul Moon
"Korea Advanced Institute of Science and Technology (KAIST), Ssummary.ai (oudgh723]] gma com, (rhkogus041 sumlee, (aaoooon@kaiit..e ac kr
ABSTRACT
In the context of multi-agent reinforcement learning, generalization a challenge to solve various tasks that may require different joint policies coordination with- out relying on policies specialized for each task. We refer to this type of problem multi-task, and we train agents to be versatile in this multi-task setting through a single training process. To address this challenge, we introduce TRajectory- class-Aware Multi-Agent reinforcement learning (TRAMA). In TRAMA, agents recognize a task type by identifying the class of trajectories they are experienc- ing through partial observations, and the agents use this trajectory awareness or prediction as additional information for action policy. To this end, we introduce three primary objectives TRAMA: (a) constructing a quantized latent space generate trajectory embeddings that reflect key similarities among them; (b) con- ducting trajectory clustering using these trajectory embeddings; and (c) building a trajectory-clase -llectory-clsssaaaree policy. Specifically for (c), we introduce trajectory-class predictor that performs agent-wise predictions on the trajectory class; and we de- sign a trajectory- -laass representation model for each trajectory class. Each agent takes actions based on this trajectory-class representation along with its partial observation for task-aware execution. The proposed method is evaluated vari- ous tasks, including multi-task problems built upon StarCraft II. Empirical results show further performance improvements over state-of-the-ee tfrttheart baselines.
INTRODUCTION
The value factorization framework (Sunehag et al. 2017 Rashid et al. 2018 Wang et al. 2020a) under Centralized Training with Decentralized Execution (CTDE) paradigm (Olichoek et al. 2008 Gupta et al., 2017) has demonstrated its effectiveness across range of cooperative multi-agent tasks (Lowe et al. 2017]; Samvelyan et al.) 2019). However, learning optimal policy often takes a long complex tasks, and the trained model often falls into suboptimal policies. These suboptimal results are often observed in complex tasks, which require agents to search large joint octionnobservation spaces.
training time in
Researchers have introduced task division methods in diverse frameworks to overcome this limita- tion. Although previous works have used different terminologies, such skills (Yang et al., 2019 Liu et al. 2022), subtasks (Yang et al. 2022), and roles (Wang et al. 2020b 2021 |); they have shared the major objective, such as reducing the search space of each agent during training or encouraging committed behavior for coordination among agents. For this purpose, agents first determine its role, skill, or subtask often by upper-tier policies (Wang et 202T Liu al.. 2022) Yang et al.) 2022); and the agents determine actions by this additional condition along with their partial observations. Compared to common MARI. approaches, these task division methods show strong performance in some complex tasks.
Recently, multi-agent multi-tasks have become new challenges in generalizing the learned policy to be effective in diverse settings. These tasks require agents to learn versatile policies for solving distinct problems, which may demand different joint policies or coordination among agents through single MARL training process. For example, in SMACv2 (Ellis et al. 2024),, agents need to learn policies neutralizing enemies in different initial positions and even with different unit combinations unlike the original SMAC (Samvelyan al. 2019). In this new challenge, the previous task division

---

Published as a conference paper at ICLR 2025
Class A!
22
Clustering
K3
Trajectory- class
R
Class Al
82
:
Partial observation
Predicted trajectory class
Class
Trajectory- class labels
#1
trajectories
i-th agent
(a) Trajectory Clustering
(b) Trajectory- -lassss Prediction
(c) Trajectory-Class-Avei Policy
Figure 1: Illustration of the overall procedure for trajectory-class claas
aware policy learning: (a) through trajectory clustering, each trajectory is labeled. (b) Each agent predicts which trajectory class it
experiencing based on its partial observation. (c) After identifying the trajectory class, agents perform
decision-making In (c), each agent succeeds in identifying the same trajectory class based on its partial observations denoted with different colors.
approaches show unsatisfactory performance because specialized policy for a single task will not work under different settings.
Motivated by this new challenge and the limitations of state-of-the-art MARI algorithms, we de- velop framework that enables agents to recognize task types during task execution. The agents then use these task-type predictions conditions in decision-making for versatile policy learning.
Contribution. This paper presents -Rajectory-class-A Awrre Multi-Agent reinforcement learning (TRAMA) to perform the newly suggested functionality described below.
Constructing quantized latent space for trajectory embedding:
To generate trajectory
embeddings that reflect key similarities among them, we adopt the quantized latent space via Vector Quantized- Variationnl Autoencoder (VQ-VAE) (Van Den Oord et al., (017)- However. the naive adaptation of VQ- VAE for state embedding in MARL results in sparse usage of quantized vectors highlighted by (Na & Moon 2024). To address this problem in multi-task settings, we introduce modified coverage loss considering trajectory class in VQ- VAE training to spread quantized vectors evenly throughout the embedding space of feasible states.
Trajectory clustering:
In TRAMA, we conduct trajectory clustering based on trajectory embeddings the quantized latent space to identify trajectories that share key similarities. Since we cannot conduct trajectory clustering at every MARL training step, we introduce classifier to determine the class of the newly obtained trajectories from the environment.
Trajectory-class-avees policy: After identifying the trajectory class, we train the agent- wise trajectory-class predictor, which predicts a trajectory class using agents" partial ob- servations. Using this prediction, the trajectory-class representation model generates representation. This task type trajectory representation is then provided to an action policy along with local observations. In this way, agents can learn ceeettoy-class-eeeendeee policy.
jjjjectory-clss-eeeeee
Figure iillustrates the conceptual process of TRAMA by enumerating key functionalities on inte- grating the trajectory-clas. information in decision decision-making.
PRELIMINARIES
2.1 DECENTRALIZED POMDP
Decentralized Partially Observable Markov Decision Process (Dec-POMDP) (Oliehoek & Amato 2016) is a widely adopted formalism for general cooperative multi-agent reinforcement learning (MARL) tasks. In Dec-POMDP, we define the tuple G (1 A, P, R, Ω, 0, n, n,,,,, where / is the finite set of n agents; E is the true state of the global state space S: A X is the joint action space and the joint action is formed by each agent's action a' E A'; P(s'|s, a) is the state transition function to new state E S given and a; a reward function R provides a scalar reward R(s, a, s') E IR given transition (8, s'; O is the observation function generating each
2

---

Published as a conference paper at ICLR 2025
agent's observation o' E S2' from the joint observation space = x,S';; and finally, is a discount factor. At each timestep, an agent receives a local observation o² and takes an action a' E A' given o'. Given the global state s and joint action a, state transition function P(s'|s, , a) determines the next state s'. Then, Rt provides a common reward r = R(s, a, 8') to all agents. For MARI training, follow the conventional value factorization approaches under the CTDE framework. Please refer to Appendix A.3 for details.
2.2 MULNTI-AGENT MULTI-TASK
This section introduces formal definition of multi-agent multi-task T. under dec-POMDP set- tings. In this paper, we omit the term multi-agent and denote multi-task for conciseness.
Definition 2.1 (Multi-agent multi-task T) A partially observable multi-agent multi-task T is de- fined by a tuple (I,SAA,P, R, S2,O, O,n.n. n,,,, Y, K)), where K. is a of tasks, S= =U and 5 = S2 for task-specific state space Sy and joint observation space Su. Then, a partially-observable single- task Tk for k K. is defined by a tuple of (I, A, P, R. Sh O,n, n, n,,)), such that Vki, k2 K. Sky O SS% + #0 and S2k, n 12 + 0.
Support Difference Although each task Tx shares the governing transition P, reward R, and ob- servation O functions, Definition 2.1 implies that vk1, h2 K, dom(P);; n dom (omm()), * 0, (om(R)ss &, n dom(R) * and dom O)2, n dom(O){ &f 0, where subscript kı and 12 repre- sent task specific values. For example, two tasks with different unit combinations SMACv2 (Ellis) ett al., 2024) satisfy this condition. Figure iillstratts the state diagram of multi-task settings. Thus, agents need to learn generalizable policies to maximize the expected return obtained from T.
Unsupervised Multi-Task Importantly, in multi-task setting in this pa- per, task ID & is unknown in both training and execution. This differs from general multi-task learning, where task ID is generally given during training (Omidshafici et al. 2017| Hansen et al. [224]] Yu et al., 2020 Tassa 2018). This unsupervised multi-task setting practical for multi-agent tasks. For example, in a football game, allied teammates predict opponents' strategies based on their observations during compe- tition to respond appropriately. In such cases, task label indicating the opponents' strategy (or task type) is not provided to the allied team or cooperating agents.
= S₁ U2
S
S2
Figure 2: State Diagram of multi-task setting
This setting can also be viewed as a formal definition of SMACv2 task (Ellis et al. 2024). To address this challenging multi-task problem. TRAMA begins by identifying which task given trajectory belongs to through clustering, assuming that trajectories from the same task are more similar than those from different tasks.
2.3 QUANTIZED LATENT SPACE GENERATION WITH VQ-VAE
In this paper, we utilize VQ-VAE (Van Den Oord et al. 2017) to generate trajectory emheddings in a discretized latent space. We follow VQ-VAE adoption in MARL introduced by LAGMA (Na & Moon 2024). The VQ-VAE for state embedding in MARL contains an encoder network fī :
Rd decoder network fg Rd S, and trainable embedding vectors used codebook with size n. denoted by e = f21,
e2, .n.,n) where E IR for all j= 21,2.... ne). An encoder output x = (s) E Rd replaced to discretized latent xg by quantization process [lg. which maps æ to the nearest embedding vector in codebook follows.
[m]g = Cz, where z = argmin, |2 ey11a
Then, decoder få reconstructs the original state s given quantized vector input z.. We follow the objective presented by LAGMA to train an encoder Jã. a decoder fg. and codebook
Ctr (0, e) = Lve e) + Aevi [J(()
Σ (EF(()
[|sg[f[(s)] - e,1
(2)
Lvglé,e) =11(((((()))) -s113 Avy [gg[{[(s]]
x. 1 Haammitll/fg6 -g(
(3)
3

---

Published as a conference paper at ICLR 2025
Here, sg[-] represents stop gradient. Avq. Acommit, and Åeve are scale factors for corresponding terms. We follow a straight- estimator to approximate the gradient signal for an encoder (Bengio et al. 2013). The last term in Eq. (2) is a coverage loss to spread the quantized vectors throughout the embedding space. (t) is a timestep-dependent index, which designates some por- tion of quantized vectors
given timestep. Although the previous coverage loss works well in general MARL tasks, it is observed that such (t) has limitations in multiple tasks. We modify this (t) by identifying types of trajectories during the training. As result, we expand the index to include timestep and task type, and thus the indexing function is now [t, k),, incorporating both timestep and trajectory class k.
METHODOLOGY
This section presents TRajectory-clas--Aaass Multi-Agent reinforcement learning (TRAMA). To learn
policy, we first generate trajectory embeddings before performing trajectory clustering. To this end, we construct (1) quantized latent space using modified VQ- VAE. With the trajectory embeddings in quantized latent space, we then describe the process for (2) trajectory clustering and trajectory classifier learning. Finally, we present the (3) trajectory- class-aware policy, which consists of trajectory-class predictor and a trajectory-class representa- tion model, in addition to the action policy network.
(+) Quantized Latent Space Generation
[b) Trajectory Clustering and Classifier Learning
Replay Buffer
K-meams
R
f6
fz!
[g
Loss on
Sampled Quantized Sequence
Xxi
Coverage Loss
labeled?
No
Codebook
Standard CTDE framework
Mixing Network
04
05 r(Q (.'
gre-hot encoding
Environment
Trajectory-Class Representation Model
1/2
of
1
(ECIE)
He(: lof, (.
T (of.
[c) Trajectory- Saarrr-CCa---eeeee DeperdeDttddntt Policy
Figure 3: Overview of TRAMA framework. The purple dashed line represents a gradient flow.
3.1 QUANTIZED LATENT SPACE GENERATION WITH MODIFIED VQ-VAE
This paper adopts VQ-VAF (Van Den Oord et al. 2017) for trajectory embedding in quantized latent space. In this way. trajectory embedding can be represented by the sequence of quantized vectors. As illustrated in Section 2.3 LAGMA (Na & Moon 2024) presents the coverage loss utilizing timestep dependent indexing (tt) to distribute quantized vector evenly throughout the embedding space of states stored in the current replay buffer D, denoted as x = {x E Rd x = f=(s), E D). However, in multi-task settings, state distributions are different according to tasks, so Leer with J(1) does not guarantee quantized vectors evenly distributed over X- To resolve this, we additionally consider the trajectory class in the coverage loss through a modified indexing function, [tt, k)), which designates specific quantized vectors in the codebook, according to (t,k) pair. When the class of given trajectory T = {st=o, St=1, ... )t=r) } is identified as k. we consider state s, E T,, to belong to the k-th class and denote this state as st Then, the modified coverage loss, considering both timestep t and the trajectory class k, is expressed as follows.
(ev(
[[sg]5((())] eg113 jEJ(C.K)
Σ
k(1
(4)
Eq. (4) adjusts quantized vectors assigned the k-th class towards emhedding of * The details
4

---

Published as a conference paper at ICLR 2025
EAY
(a) Training with [t) (b) Training with [[t, k)
(c) Clustering of (a)
(d) Clustering of (b)
Figure 4: PCA of sampled embedding I E D. Colors from red to purple (rainbow) represent early late timestep (a) and (b). (a) and (b) are the results of sample multi-tasks with three different unit combinations with various initial positions. (c) and (d) are the clustering results of (a) and (b), respectively, and each color (red, green, and blue) represents each class.
Here, the class number Ret = 3 is assumed.
of (t, k) are presented in Appendix Then, with a given sf, the final loss function LVQ to train VQ- VQ-VAE becomes
- Lve(d, =vv((e))=(v(((ee) + Aeve [T,
Σ
[sg[f[(})]]
e,112
(5)
JEJ((LA)
Figure Dillustrates the embedding results with the proposed coverage loss compared the original [t).. As in Figure 4 (b), the proposed coverage loss [t, k) distributes quantized vectors more evenly throughout X multiple tasks compared to the cases with J(t) (a). Notably, we assumed the number of classes as net 6 in (b). Still, our model successfully captured three distinct initial unit combinations in the task, as illustrated by the three red branches in Figure (b). However, to adopt this modified coverage loss, we need to determine which trajectory class given state belongs to. Thus, we conduct clustering to annotate the trajectory class. Then, we use pseudo-class k obtained from clustering for Eq. (5).
3.2 TRAJECTORY CLUSTERING AND CLASSIFIER LEARNING
Trajectory Clustering With a given trajectory Termo we get quantized latent sequence with VQ- VAE as Txi_o = [ ]l = ( Fq.t=1," ... fq Here, only the indices of quantized vectors, i.c., TZ--o = [===0+ St-l, [T=T], J. are required to express Tx--o Thus, we can efficiently store the quantized sequence =2,-o to D along with the given trajectory, Ts- In addition to MARL training. we sample M trajectory sequences [T 1M
from D and conduct K-means clustering (Lloyd 1982 Arthur & Vassilvitskii 2006 periodically.
With the m-th index sequence T2. we compute a trajectory em- bedding 2m using quantized vectors in the codebook.
80-
(6)
Presuvved
Ceravoo
Cartrod bet
7
20-
Then, with trajectory embeddings [EM @TM clustering with the predetermined number of class Net- In this pa- per, the class labels are denoted as K 1fm==1,km=2)
we conduct K-means
0.0
0.5
1.0
1.5
2.0
2.5
3.0
T [mi]]
km=
Figure 5: Preserved Labels
Figures (c) and (d) illustrate the clustering results with trajectory embedding constructed by Eq. 6 based on quantized vectors of (a) and (b), respectively. The visual results and their silhouette score emphasize the importance of the distribution of quantized vectors. Appendix D.4 provides further analysis. However, the problem is that the class labels may change whenever the clustering is updated. Consistent class labels are important because we update the agent-wise predictor based these labels. To resolve this, we conduct centroid inirialization with the previous centroid results. Figure Sssows without considering centroid initialization.
shows the ratio of preserved labels after clustering with and
Trajectory Classifier Training Even though we determine trajectory-claat labels K stored D at specific training time, do not have labels for new trajectories obtained by interacting with the
5

---

Published as a conference paper at ICLR 2025
environment. To determine the labels for such trajectories before additional clustering update, develop classifier (+(())"") predict trajectory-class kn based eth We train fo whenever clustering is updated in parallel to MARL training through cross-entropy loss, Lyp with M samples.
1 M
((v)
Σ k =mm log(J (
Here, is an indicator function. Figure illlustrates the loss classifier training proceeds. With centroid initialization, the classifier learns trajectory embedding patterns more coherently than the case without it.
3.3 TRAECTORY-CLASS-AWW POLICY
Trajectory-Class Predictor After obtaining the class labels K for sampled trajectories either determined by clustering or classi- fier fw. we can train trajectory-class predictor πेς shared by all agents. Unlike the trajectory classifier fø. a trajectory-class pre- dictor मę only utilizes partial observation given to each agent. In other words, each agent makes a prediction on which trajectory type or class it experiencing based on its partial observation, such as kj ~ TC( (-Ih,.) Here, hi represents the observation history computed by GRUs in æT. With predetermined trajectory labels for sampled batches size of B, we train og with the following loss.
1.50
No Cempoi #1 (0060)
Gane WO0TT spc
0.75
0.00
Figure 6: Classifier Loss
((C)
1 B
k- log (Kglm
6=
t=0
i=1
Erraeectory-llas Representation Learning Base on TT. each agent predicts kj at each timestep Then, the i-th agent utilizes one-hot encoding of kj as an additional condition or prior when determining its action. Instead of directly utilizing this one-hot vector, we train trajectory-class representation model f
(-1kt) ( to generate more informative representation, gt- The purpose of gi to generate coherent information for decision-making and to enable agents to learn a trajectory- class-dep policy. As illustrated in Figure 3 we directly train fä through MARL training. In addition, we use a separate network for action policy mg and utilize the additional class representa- tion of along with partial observation of for decision-making
3.4 OVERALL LEARNING OBJECTIVE
This paper adopts value function factorization methods (Rashid et al. 2018 Wang al. 2020a Rashid et al. 2020 Zheng et al. 2021 presented in Section to train individual Q3 via Q For mixer structure, Aadividaal-Global Q6 and f can be expressed
mamly adopt QPLEX (Wang et al. 2020a), which guarantees the complete (IGM) condition (Son et al. 2019) The loss function for the action policy
(())- E E E
ymax. "pa g a') Q0 9, a)) max Qtaa (o', g a') Qt (o.g,
= E
(9) Here, p(k) is the portion of samples (o, a, o²)x generated by Ta within D. However, since we randomly sample tuple from D. the expectation over & can be omitted. In addition, g represents the joint trajectory-class representation. With this loss function, we train Qú fg and क. together with the following learning objective:
(10) Here, Ag is scale factor. Note that 0 denotes neural network parameters contained in both Qo and fg. Algorithm in Appendix C specifies the learning procedure with loss functions specified by Eqs. (5) and (10).
((0,C)= (()) + ACC())

---

Published as a conference paper at ICLR 2025
4 RELATED WORKS
4.1 TASK DIVISION METHODS IN MARL
In the field of MARL, task division methods are introduced in diverse frameworks. Although pre- vious works use different terminology. such as a subtask (Yang et al.] (022), role (Wang et al.) 2020b 202T) or skill (Yang et al. 2019; Liu et al. 2022), they share the primary objective, such as reducing search space during training encouraging committed behaviors among agents utilizing conditioned policies. HSD Yang et al. 2019). RODE (Wang et al., 2020b). LDSA (Yang et al, 2022) and HSL (Liu et al.) (022) adopt hierarchical structure where upper-tier policy network first determines agents' roles, skills, or subtasks, and then agents determine actions based on these additional conditions along with their partial observations. These approaches share commonality with goal-conditioned RL in single-agent tasks; however, the major difference is that the goal is not explicitly defined in MARL. MASER (Jeon et al. 2022) adopts subgoal generation scheme from goal-conditioned RL when it generates an intrinsic reward. On the other hand, TRAMA first clusters trajectories considering their commonality among multi-tasks. Then, agents predict task types by identifying the trajectory class and use this prediction as additional information for action policy. In this way, agents utilize wwapectorr-clss-depppeeee
or task-specific policies.
Appendix A presents additional related works regarding state space abstraction and some prediction methods developed for MARL..
5 EXPERIMENTS
In this section, we evaluate TRAMA through multi-task problems built upon SMACv2 (Eliss et al. 2024) and conventional MARL benchmark problems (Samvelyan et al., 2019 Ellis al. 2024). We have designed the experiments to observe the following aspects.
Q1. The performance of TRAMA in multi-task problems and conventional benchmark problems compared to state-of-the-art MARL frameworks
Q2. The impact of the major components of TRAMA on aginttwise trajectory-class pre- diction and overall performance
Q3. Trajectory-claat distribution in the embedding space
To compare the performance of TRAMA. consider various baseline methods: popular bascline methods such as QMIX (Rashid et al., 2018) and QPLEX (Wang et al. 2020a); subtask-based meth- ods such as RODE (Wang et al.) 2021), LDSA YYang et al. 2022), and MASER (Jeon et al.] (022); memory-based approach such as EMC (Zheng et al. 2021 and LAGMA (Na & Moon 2024). For
the baseline methods, we follow the hyperparameter settings presented in their original paper and implementation. For TRAMA, the details of hyperparameter settings are presented in Appendix B
5.1 COMPARATIVE EVALUATION ON BENCHMARK PROBLEMS
We first evaluate TRAMA on the conventional SMACv2 tasks (multi- tasks) such as 15_155vv... vs_5 5 and t5_vs_5 5 in (Ellis et al.) 2024). In Figure 7 TRAMA shows better learning et- ficiency and performance compared to other baseline methods, including a memory-based approach such as EMC @Zhegg et al., (021), utilizing an additional memory buffer. Besides multi-task problems, we conduct ad- ditional experiments on the original SMAC (Samvelyan et al.] 2019) tasks to see how TRAMA works in single-task settings. Figure g illustrates the results, and TRAMA shows comparable or better performance compared other methods. Notably, TRAMA succeeds
TRAMA(DUR))
20.0
25.w w.W.W.
200
11
17.5
12.5
10.0
12.0
255
25
6.0
1.0
2.0
3.0
0.5
2.0
30
((all
Tom0
Figure 7: Performance comparison of TRAMA against baseline algorithms on p5 vs.5 and t5.vs.5 in SMACv2 tmulti-taskss) Here, == is assumed.
7

---

Published as a conference paper at ICLR 2025
in learning the best policy the end in super hard tasks such MM2 and 6h.vs-8z. In the fol- lowing section, we present a parametric study on nett and explain how we determine the appropriate value based on the tasks.
TRAMA (ours)
QMIX
RODE
LDSA
MASER
EMC
LAGMA
Test Maen Return (-)
20
1c3s52
20
5m, vs 8m
20
MMM2
20
8h ve 82
15-
15
15
15-
10
10
10
10
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1,0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
0.0
1.0
2.0
3.0
4.0
T (mil
T (mll)
T
(mil)
T
[mil)
Figure 8: Performance comparison of TRAMA compared to baseline algorithms SMAC task (single-task).
5.2 COMPARATIVE EVALUATION ON VARIOUS MULTI-TASK PROBLEMS
To test MARL algorithms additional multi-task problems, introduce the four modified tasks built upon SMACv2 as presented in Table In these new tasks, two types of initial position distribu- tions are considered, and initial unit combinations are randomly selected from the designated sets. In addition, agents get rewards only when each enemy unit is fully neutralized, similar to sparse reward settings in eJeon et al. 2022 Na Moon 2024 Appendix B provides further details of multi-task problems and SMACv2.
Table 1: Task configuration of multi-task problems
Name
Initial Position Type
Unit Combinations (Ncomi)
SurComb3
Surrounded
(3s22, 2c32, 2c3s) } {3s22, 2c32, 2c3s)
reSurComb 3
Surrounded and Reflected
SurComb4
Surrounded
{1c2s22, 3s22, 2c32, 2c3s) c1c2s2z, 3s2z, 2c3z, 2c3s)
reSurComb4
Surrounded and Reflected
To evaluate performance, we consider the overall return value instead of the win-rate, as the learned policy may be specialized for specific tasks while being less effective for others among multiple tasks. Figure Dillustrates the overall return values for multi-task problems. As illustrated in Figure TRAMA consistently demonstrates better performance compared to other baseline methods.
9
TRAMA (ours)
QMIX
RODE
LDSA
MASER
QPLEX
EMC
LAGMA
Test Mesin Retur [-]
SurComb3
reSurComb3
SurComb4
reSurComb4
3
3-
2
2
1
1
0.0
1.0
2.0
3.0
4.0
5.0
0.0
1.0
2.0
3.0 (mil)
4.0
5.0
0.0
1.0
2.0
3.0 (mil)
4.0
5.0
0.0
1.0
2.0
3.0
4.0
5.0
(mll)
(mil)
Figure 9: The mean return of TRAMA compared to baseline algorithms on four multi-task problems presented in Table !
To see how well each agent predicts the trajectory class, we present the learning loss of LG and the overall accuracy of the prediction on reSurComb4 task in Figure 10 Appendix [D.7 presents an additional analysis for a trajectory-class prediction made by agents. Notably, the agents accurately identify which types of trajectory classes they are experiencing based on their partial information throughout the episodes. In this case, agents can coherently generate trajectory-class representation through f and condition on this additional prior information for decision maaiiig...maing.... With this extra information, agents can leam distinct policies based on trajectory class and execute different joint policies specialized for each task, resulting in improved performance. In Section we will discuss how trajectories are divided into different classes and represented in the embedding space.
8

---

Published as a conference paper at ICLR 2025
5.3 PARAMETRIC STUDY
In this study, we check the performance variation ac- cording key parameter Nht evaluate the impact of the number of trajectory classes on the general perfor- mance. We considered ndt {2, 4, 8, 16} for multi- task problems SurComb3 and SurComb4, and Ret (4, 6, 8, 16} for the original SMACv2 task p5.vs vs..... Fig- ure [1 presents the overall return according to differ- ent Ret. To evaluate the efficiency of training and per- formance together, we compare cumulative return, PR. which measures the area below the mean return curve. The high value of fæ represents better performance. In Figure (d), P is normalized by its possible maximum value. From Figure [1 (d), we see that peaks of jin occur around Rat = 4 for SurComb. and SurComb4, and Rd = 8 for p5_vs_5, respectively. Interestingly, these numbers seem highly related to variations in unit combinations. In SurComb3
1-30
(a) Le
(b) Accuracy [%]
Figure 10: With n learning loss тÇ and the mean accuracy of the trajectory-class prediction made by agents in (reSurComb4).
15
TOMAA -g
19006 (y/1)
20.11
00
25
(a) SurComb3
(b) SurComb4
(c) p5_vs_5
(d) Normalized ps
Figure 11: Parametric study of ne on surComb3, surComb 4 and p5_vs_5.
and SurComb4, 3 and 4 unit combinations are possible, and thus jectory diversity. Therefore, combinations of multi-tasks. When larger Re! is selected, the agent-wise prediction accuracy may degrade because the number of options increases. However, some classes share key similarities even though they are labeled different classes. Trajectory embedding in quantized embedding space can efficiently capture these similarities. Thus, agents can still learn coherent policies specialized on each task in multi-task settings even with different trajectory class labels. We will further elaborate on this in Section 5.5
4 well captures the tra- recommended to determine Nel considering the diversity of unit
5.4 ABLATION STUDY
In this subsection, conduct the ablation study see the effect of major components of TRAMA. First, to see the importance of coherent label generation in class clustering, we consider No-Init representing clustering without centroid initialization presented in Section 3.2 In addition, ab- late the proposed coverage loss [(t, k) and consider (tt) instead when constructing quantized embedding space. Figure 12 illustraaees the corresponding results.
HAMO RAMINTUALEOO
[44)
(a) Preserved Label Ratio
(b) Ce
t=
(d) Mean Return
Figure 12: Ablation study on SurComb4.
In Figure 12 (a), without centroid initialization, the trajectory class labels change frequently, and the loss for Te fluetuates as illustrated in (b). As a result, कę in both TRAMA (No-Init) and TRAMA (No-Init & (t)) predicts trajectory class labels almost randomly as illustrated in Figure (c). lead- ing to degraded performance shown in Figure 12 (d). In the of TRAMA (J(I)) coherent

---

Published as a conference paper at ICLR 2025
0.4
0.3
0.2
0.0
-0.1
Early
-0.3
-022
-0.2
Timestep
0.0
0.2
0.4
+
0.2
Axic
1.0 -0.2
2c32
1c252z
Visualization of embedding results and test episodes for SurComb4. Here, nw==8 is assumed, and gray dots represent quantized vectors in the VQ codebook. Solid lines represent each test episode, while colored dots represent the majority opinion on the trajectory class predictions made by agents. Each color denotes different class.
Figure 13:
labels are generated due to centroid initialization. However, prediction accuracy is lower and fluc- tuates compared to the full version of TRAMA, the quantized vectors not evenly distributed over x as described in Figure Therefore, TRAMA without [(,, k) cannot sufficiently capture the key differences in trajectories. In Appendix D class representation model fg.
present additional ablation studies regarding the
5.5 QUALITATIVE ANALYSIS
In this section, we evaluate how the trajectory classes are identified in the quantized embedding space. We consider SurComb4 task as test case and assume = 8 for this test. Figure T3 illus- trates the visualization of embedding results and test episodes. Notably, four branches are developed in quantized embedding space after training, and each branch is highly related the initial unit com- binations. The result implies that the initial unit combinations significantly influence the trajectories of agents. Although the larger number of nat is chosen compared to the number of possible unit combinations (Ncomt = 4) in SurComb4, 4, two classes are assigned to each branch in the quantized embedding space, making the model less sensitive
misclassification within each pair. In Figure 13 classes 4 and 8 are assigned to 3s2z; classes 2 and 3 to cc2s2z; classes 5 and to 2c3z; and classes 1 and to 2c3s. By identifying the trajectory class, agents can generate additional prior information and utilize sijectory-class-eepee can see, this trajectory-class identification is important in solving multi-task problems T. Thus, it would be interesting to see how TRAMA predicts trajectory classes in -istii-disttibuutin tasks. As TRAMA learns identify trajectory class an unsupervised manner, without task ID, agents can identify similar tasks among in-distribution tasks. Then, agents rely on these predictions during decision-makinn making, thereby promoting a joint policy that benefits OOD tasks. Appendix D.6 presents OOD experiments and their corresponding qualitative analysis, demonstrating TRAMA's generaliz- ability across various OOD tasks.
policies conditioned on this prediction. As
CONCLUSION
This paper presents TRAMA, a new framework that enables agents to recognize task types by iden- tifying the class of trajectories and to use this information for action policy. TRAMA introduces three major components: 1) construction of quantized latent space for trajectory embedding, 2) policy. The constructed quantized latent space allows trajectory embeddings to share the key commonality between trajectories. With these tra- jectory embeddings, TRAMA successfully divides trajectories into clusters with similar task types in multi-tasks Then, with a trajectory-ellan predictor, each agent predicts which trajectory types agents experiencing and uses this prediction to generate trajectory-class representation. Finally, agents learn wlwttooryy--sssaaeeee policy with this additional information. Experiments validate the effectiveness of TRAMA in identifying task types in multi-tasks and in overall performance.
trajectory clustering, and 3) trajectory-class-ave
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
This work was supported by the IITP(Intiiute of Information & Communications Technology Plan- ning & Vvabuation)IITRCC(fformallt Technology Research Center) grant funded by the Korea Movisty of Science and
REFERENCES
David Arthur and Sergei Vassilvitskii. k-means++: The advantages of careful seeding. Technical report, Stanford, 2006.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv::IXi..... 2013.
Nicolas Carion, Nicolas Usunier, Gabriel Synnaeve, and Alessandro Lazaric. A structured pre- diction approach for generalization in cooperative multi-agent reinforcement learning. In H. Wallach, Larochelle, Beygelzimer, F. d'Alché-Buc, Fox, and Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., wiwess/aper//011111
2019. URL ftttss///poceedinng.neeeiis..... Papeerappppppaa.
File/ File/ http://cc11.hh000001111/10111110
Filippos Christianos, Georgios Papoudakis, Muhammad A Rahman, and Stefano Albrecht, Scal- ing multi-agent reinforcement learning with selective parameter sharing. In International Confer- ence on Machine Learning, pp. 1989-1998. PMLR, 2021.
Sanjoy Dasgupta. Experiments with random projection. arXiv preprint arXiv: 2013.
Benjamin Ellis, Jonathan Cook, Skander Moalla, Mikayel Samvelyan, Mingfei Sun, Anuj Mahajan, Jakob Foerster, and Shimon Whiteson. Smaev2: An improved benchmark for cooperative multi- agent reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Aditya Grover, Maruan Al-Shedivat, Jayesh Gupta, Yuri Burda, and Harrison Edwards. Learning policy representations in multiagent systems. In International conference on machine learning, pp. 1802-1811. PMLR, 2018.
Marek Grzeš and Daniel Kudenko. Multigrid reinforcement learning with reward shaping. In Inter- national Conference on Artificial Neural Networks, pp. 357-366. Springer, 2008
Jayesh K Gupta, Maxim Egorov, and Mykel Kochenderfer. Cooperative multi-agent control using deep reinforcement learning. In International conference on autonomous agents and multiagent systems, pp. 66-83 83. Springer, 2017.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with dis- crete world models. arXiv preprint arXiv:2010.022 023, 2020.
Nicklas Hansen, Hao Su, and Xiaolong Wang. TD-MPC2: Scalable, robust world models for contin- uous control. In The Twelfth International Conference on Learning Representatiinss 2024. URL
net forum? (irrrmmmiid=OOxhtttttt..
Hado Hasselt. Double q-learning. Advances in neural information processing systems, 23, 2010.
He He, Jordan Boyd-Graber, Kevin Kwok, and Hal Daumé III. Opponent modeling in deep rein- forcement learning. In International conference machine learning, pp. 1804-1813, PMLR, 2016.
Siyi Hu, Fengda Zhu, Xiaojun Chang, and Xiaodan Liang. Updet: Universal multi-agent reinforce- ment learning via policy decoupling with transformers, arXiv preprint arXiv:2101.0800, 2021.
Jeewon Jeon, Woojun Kim. Whiyoung Jung, and Youngchul Sung. Maser: Multi-agen reinforce- ment learning with subgoals generated from experience replay buffer. In International Conference Machine Learning, pp. 10041-10052. PMLR, 2022.

---

Published as a conference paper at ICLR 2025
Nan Jiang, Alex Kulesza, and Satinder Singh. Abstraction selection in model-based reinforcement learning. International Conference on Machine Learning, pp. 179-188. PMLR, 2015.
Guangyu Li, Bo Jiang, Hao Zhu, Zhengping Che, and Yan Liu. Generative attention networks for multi-agent behavioral modeling. In Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 34, pp. 7195-7202, 2020.
Zhuo Li, Derui Zhu, Yujing Hu, Xiaofei Xie, Lei Ma, Yan Zheng, Yan Song, Yingfeng Chen, and Jianjun Zhao. Neural episodic control with state abstraction. arXiv preprint arXiv:2301.11490 2023.
Yuntao Liu, Yuan Li, Xinhai Xu, Yong Dou, and Donghong Liu. Heterogeneous skill learning for multi-agent tasks. Advances in Neural Information Processing Systems, 35:3701 2022.
Stuart Lloyd. Least squares quantization pem. IEEE transactions on information theory, 28(2): 129-137, 1982.
Ryan Lowe, Yi Wu, Aviv Tamar, Jean Harb, Pieter Abbeel, and Igor Mordatch. Multi-agent actor- critic for mixed cooperative-competiiis environments. Neural Information Processing Systems (NIPS), 2017.
Hyungho Na and Il-chul Moon. Lagma: Latent goal-guided multi-agent reinforcement learning.
arXiv preprint arXiv:2405. 19998, 2024.
Hyungho Na, Yunkyeong Seo, and Il-chul Moon. Efficient episodic memory utilization of coopera-
tive multi-agent reinforcement learning. arXiv preprint arXiv:2403.01112, 2024.
Frans A Olichock and Christopher Amato. A concise introduction to decentralized POMDPs.
Springer, 2016.
Frans A Oliehoek, Matthijs TJ Spaan, and Nikos Vlassis. Optimal and approximate q-value func- tions for decentralized pomdps. Journal of Artificial Intelligence Research, 32:289-353, 2008.
Shayegan Omidshafiei, Jason Pazis, Christopher Amato, Jonathan P How, and John Vian. Deep decentralized multi-task multi-agent reinforcement learning under partial observability. In Inter- national Conference Machine Learning, pp. 2681-266999 PMLR, 2017.
Georgios Papoudakis and Stefano Albrecht. Variational autoencoders for opponent modeling in multi-agent systems. arXiv preprint arXiv:2001.10829 2020.
Georgios Papoudakis, Filippos Christianos, and Stefano Albrecht. Agent modelling under partial ob- servability for deep reinforcement learning. Advances in Neural Information Processing Systems, 34:19210-19222. 2021.
Roberta Raileanu, Emily Denton, Arthur Szlam, and Rob Fergus. Modeling others using oneself in multi-agent reinforcement learning. In International conference machine learning, pp. 4257- 4266. PM 2018.
Tabish Rashid, Mikayel Samvelyan, Christian Schroeder, Gregory Farquhar, Jakob Foerster, and Shimon Whiteson. Qmix: Monotonic value function factorisation for deep multi-agent reinforce- ment learning. In International conference machine learning, pp. 4295-4304. PMLR, 2018.
Tabish Rashid, Gregory Farquhar, Bei Peng. and Shimon Whiteson. Weighted qmix: Expanding monotonic value function factorisation for deep multi-agen reinforcement learning. Advances in neural information processing systems, 33:10199-10210, 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 starcraft multi-agent challenge. arXiv preprint
04043, 2019.
Siqi Shen, Chennan Ma, Chao Li, Weiquan Liu, Yongquan Fu, Songzhu Mei, Xinwang Liu, and Cheng Wang. Riskq: risk-sensitive multi aalti- reinforcement learning value factorization. Ad- vances in Neural Information P Processing Systems, 36:34791-34425, 2023.
12

---

Published as a conference paper at ICLR 2025
Kyunghwan Son, Daewoo Kim, Wan Ju Kang, David Earl Hostallero, and Yung Yi. Qtran: Learning to factorize with transformation for cooperative multi-agent reinforcement learning. In Interna- tional conference machine learning, pp. 5887-5896. PMLR, 2019.
Peter Sunehag, Guy Lever, Audrunas Gruslys, Wojciech Marian Czarnecki, Vinieius Zambaldi, Max Jaderberg, Mare Lanctot, Nicolas Sonnerat, Joel Z Leibo, Karl Tuyls, et al. Value-decomposition networks for cooperative multi-agent learning. arXiv preprint arXiv:1706.005: 2017.
Yunhao Tang and Shipra Agrawal. Discretizing continuous action space for on-policy optimization. In Proceedings of the auai conference artificial intelligence, volume 34, pp. 5981-5988, 2020.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv: 1801.000 2018.
Aaron Van Den Oord, Oriol Vinyals, al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q- learning. Proceedings of the AAAI conference on artificial intelligence, volume 30, 2016.
Jianhao Wang, Zhizhou Ren, Terry Liu, Yang Yu, and Chongjie Zhang. Qplex: Duplex dueling
multi-agent q-learning. arXiv preprint arXiv:2008. 01062, 2020a.
Tonghan Wang, Heng Dong, Victor Lesser, and Chongjie Zhang. Roma: Multi-agent reinforcement learning with emergent roles. arXiv preprint arXiv:2003.08053 2020b.
Tonghan Wang, Tarun Gupta, Anuj Mahajan, Bei Peng. Shimon Whiteson, and Chongjie Zhang. Rode: Learning roles to decompose multi-agent tasks. In Proceedings of the International Con- ference Learning Representations (ICLR), 2021.
Jiachen Yang, Igor Borovikov, and Hongyuan Zha. Hierarchical cooperative multi-agent reinforce- ment learning with skill discovery. arXiv preprint arXiv: 2019.
Mingyu Yang, Jian Zhao, Xunhan Hu, Wengang Zhou, Jiangcheng Zhu, and Houqiang Li. Ldsa: Learning dynamic subtask assignment in cooperative multi-agent reinforcement learning. Ad- vances in Neural Information Processing Systems, 15:1698-17101 2022.
Tianhe Yu, Deirdre Quillen, Zhampeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: benchmark and evaluation for multi-task and meta reinforcement learning. In Conference robot learning, pp. 1094-1100, PMLR, 2020.
Xiaopeng Yu, Jiechuan Jiang, Wanpeng Zhang, Haobin Jiang, and Zongqing Lu. Model-based op- ponent modeling. Advances in Neural Information Processing Systems, 35:28208-28221, 2022.
Chongjie Zhang and Victor Lesser. Multi-agent learning with policy prediction. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 24, pp. 927-934, 2010.
Lulu Zheng, Jiarui Chen, Jianhao Wang, Jiamin He, Yujing Hu, Yingfeng Chen, Changjie Fan, Yang Gao, and Chongjie Zhang, Episodic multi-agent reinforcement learning with curiosity-driven exploration. Advances in Neural Information Processing Systems, 34:3757-3769, 2021.
Derui Zhu, Jinfu Chen, Weiyi Shang, Xuebing Zhou, Jens Grossklags, and Ahmed E Hassan. Deep- memory: model-based memorization analysis of deep neural language models. In 202/ 36th JEEE/ACM International Conference on Automated Software Engineering (ASE), pp. 1003-1015. IEEE, 2021.
13

---

Published as a conference paper at ICLR 2025
A
ADDITIONAL RELATED WORKS AND PRELIMINARIES
A.1 1 STATE SPACE ABSTRACTION
has been effective in grouping similar characteristics into a single cluster, which is called state space abstraction, in various settings such model-basee RL (Jiang et al. 2015 Zhu et al. 2021
Hafner et al. 2020) and model-free settings (Grzes & Kudenko 2008 Tang & Agrawal 2020). NECSA, introduced by (Li et al.) 2023), is the model that facilitates the abstraction of grid based state-action pairs for episodic control, achieving state-of-the-art (SOTA) performance in a general single-reinforcement learning task. The approach could alleviate the usage of inefficient memory in conventional episodic control. However, an additional dimensionality reduction process is in- evitable, such as random projection in tigh-dimensional tasks in (Dasgupta) 2013). In (Na et al.) 2024), EMU utilizes a state- semantic embedding for efficient memory utilization. LAGMA (Na Moon 2024) employs VQ-VAE for state embedding and estimates the overall value of ab- stracted states to generate incentive structure encouraging transitions toward goal-reaching trajec- tories. Unlike previous works, we use state-space abstraction to generate trajectory embedding in quantized latent space, ensuring that these embeddings share key similarities. We then cluster the trajectories into several classes based on task commonality. In this way, TRAMA can learn task- aware policy by identifying trajectory classes in multi-tasks
A.2 PREDICTIONS IN MARI
Predictions in MARL are generally used to model the actions of agents (Zhang & Lesser 2010 He
et al., 2016 Grover et al. 2018 Raileanu et al., 2018 Papoudakis et al. [021 Yu et al.] 2022).
In (Carion al., 2019: Li et al. 2020; Christianos et al. (202]], models also include groups or
agents' tasks for their prediction. The authors in (He et al. 2016) utilize the Q-value to predict
opponents' actions, assuming varying opponents' policies. On the other hand, Raileanu et al. (2018)
presents the model to predict other agents' actions by updating hidden states. [Papoudakis & Al- brecht (020) introduce the opponent modeling adopting variational autoencoder (VAE) and A2C without accessibility to opponents' information. LIAM (Papoudakis et al.] (202T) learns the trajec- tories of the modeled agent using those of the controlled agent. the other hand, our approach, agents predict which trajectory class they are experiencing to generate additional inductive bias for decision-making based on this prediction.
A.3 CENTRALIZED TRAINING WITH DECENTRALIZED EXECUTION (CTDE)
Under Centralized Training with Decentralized Exccution (CTDE) framework, value factorization approaches have been introduced by Suunhha et al. 2017: Rashid et al. 2018 Son et al.) 2019 Rashid et al. 2020 Wang et al. 2020a) to solve fully cooperative multi-agent reinforcement learning (MARL) tasks, and these approaches achieved of-eeoof-the-rr performance in challenging bench- mark problems such as SMAC (Samvelyan et al.] 2019). Value factorization approaches utilize the
joint action-value function Qaat with learnable parameter 0. Then, the training objective L(0) can be expressed as
(0) = E..a.... + ymax Q a') - t a))
(11)
where D is a replay buffer; is the joint observation; Q is target network with online parameter 0- for double Q-learning (Hasselt 2010 Van Hasselt et al. 2016;; and Q8o and Q8aa include both mixer and individual policy network.
B EXPERIMENT DETAILS
B. EXPERIMENT DESCRIPTION
In this section, we present details of SMAC (Samvelyan et al., 2019). SMACv2 Ellis et al. 2024) and multi-task problems presented in Table built upon StarCraft II. To test the generalization of policy, SMACv2 contains highly varying initial positions and different unit combinations within one map, unlike the original SMAC tasks. In new tasks, agents may require different strategies against
14

---

Published as a conference paper at ICLR 2025
enemies with different unit combinations and initial positions. TRAMA makes agents recognize which task they are solving and then utilize these predictions as additional conditions for action policy. Figure compares the different characteristics between single-task and multi-tasks.
Varying Unit Combination
Init. distribution
Unit Combination
Init. distribution
Reflected
Surrounded
Reflected
Single-Task
Multi-Task
Figure 14: Comparison between single-task and multi-tasks.
In SMACv2, unit combinations for agents randomly selected from set of given units based the predefined selection probability for each unit. For example, in t5_va_5, five units are drawn from three possible units, such Marine, Marauder and Medivac, according to predetermined probabilities. Initial unit positions are randomly selected between Surrounded and Re f eceee..
On the other hand, in our multi-tasks presented in Table initial unit combinations are selected from the predefined sets. For example, three unit combinations are possible among s3s2z, 2c3z, 2 2c3s} in SurComb3. 3. Here, s, 2, and с represent Stalker, Zealot, and Colo colossus, respectively. In addition, in our multi-tasks we adopt sparse reward settings similar to (Jeon et al.) [202]] Na & Moon 2024) unlike SMACv2. 2 the reward structure of the multi-tasks we presented.
Table Reward settings for multi-tasks.
Condition
Sparse reward
All enemies die (Win) Each enemy dies
+200
+10
Each ally dies
Note that both SMAC (Samvelyan et al. 2019) and SMACv2 (Ellis et al. 2024) normalize the reward
output by the maximum reward agents can get so that the maximum return (without considering discount factor) becomes 20. Due this setting, the maximum
in multi-tasks becomes around 3.5. 5. Table 3 presents details of tasks evaluated in the experiment section.
Table 3: Task Specification
Task
Nagent
Dim. of state space
Dim. of action space
Episodic length
SurComb3
130
200
reSurComb3
130
200
Multi-task
SurComb4
5
130
200
reSurComb4
130
200
p5_vs_p5
11
200
SMACv2
15_vs.15
120
11
200
1c3s5z
270
15
180
5m.vs.6
120
11
70
SMAC
MMM2
322
18
180
6h_vs_8z _z
140
14
150
15

---

Published as a conference paper at ICLR 2025
B.2 EXPERIMENT SETTINGS
For the performance evaluation, measure the mean return computed with 128 samples: 32 episodes for four different random seeds. For baseline methods, we follow the settings presented in the original papers their original codes. We use almost the same hyperparameters throughout the various tasks except for Ret- For VQ-VAE training, we use the fixed hyperparameters for all Pq tasks, such =025 Acommi=0.125. Aev=0.122 in Eq. (5) ny=500, and fireq =10. Here, no the update interval for clustering and classifier learning, and "freq represents the update interval of VA-VAE. Algorithm
2 presents details of TRAMA training and the parameters used in overall summarizes the task-dependent hyperparameter settings for TRAMA. Here, the dimension of latent space is denoted as d.
training. Table 4
Table 4: Hyperparameter settings for TRAMA experiments.
Task
nee
d
CT
SurComb3
4
reSurComb3
6
Multi-task
256
4
50K
SurComb4
reSurComb4
8
p5 p5 vS L5
8
SMACv2
256
4
50K
8
1c3s5z 5m 6m MMM2
3
50K
SMAC
256
4
3
6h_vs_82 8z
200K
B.3 INFRASTRUCTURE AND CODE IMPLEMENTATION
For experiments, we mainly use GeForce RTX 3090 and GeForee RTX 4090 GPUs. Our code is built on PyMARL (Samvelyan et al. 2019) and the open-sourced code from LAGMA (Na & Moon 2024). Our official code available at: https: / halpp://gittiie //gitubb... com aai waw.lllaaaaittaaaa
B.4 TRAINING AND COMPUTATIONAL TIME ANALYSIS
This section presents the total training time required for each model to learn each task. Before that, present the computational costs of newly introduced modules in TRAMA. Table presents the results. In Table 5 the MARL training module includes training for prediction and VQ-VAE mod- ules. The computational cost is about 20% increased compared to the model without prediction and VQ- VAE modules, overall training time does not increase much compared to other complex baseline methods illustrated in Table 6 This is because most computational load in MARL of- ten comes from rolling out sample episodes by interacting with the environment. In addition, as VQ-VAR module is periodically updated, we measure the updates. Clustering and classifier training are called sparsely compared to MARL training; over- all, additional computational costs
computational time with VQ-VAE
burdensome. In addition, we can expedite computational costs for classifier training, as classifier already converges to optimal one after sufficient training, as illustrated in Figure 6
Table Computational costs of TRAMA modules.
Module
Computing time per call [s]
MARL training module
0.83
Clustering
0.13-0.3
Classifier training
2~15
16

---

Published as a conference paper at ICLR 2025
Training times of all models measured in GeForce RTX 3090 or RTX 4090. In Table 6 marker (*) represents the training time measured by GeForce RTX 4090. Others are measured by RTX 3090. As in T TRAMA does not take much training time compared to other baseline methods, even with a periodic update for trajectory clustering, classifier learning, and VQ-VAR training. Again, as classifier already converges to optimal after sufficient training as illustrated in Figure [6] can further expedite training speed by reducing the frequency of updating the classifier, fv.
Table 6: Training time for each model in various tasks (in hours).
Model
5m.vs 6m (2M)
1c3s5z (2M)
p5.vs .9.55 (5M)
SurComb3 (5M)
EMC
8.6
23.1
23,2
21.6*
MASER
12.7
12.9
21.8
23.5
RODE
6.0
10.5
15.0
20.6
TRAMA
9.1
10.5
12,8*
15.1
C IMPLEMENTATION DETAILS
In this section, present the implementation details of TRAMA. In TRAMA, we additionally consider the trajectory class k as an additional condition along with timestep for selecting indices of quantized vectors during VQ-VAE training. We denote this indexing function T (t. k). Algorithm presents details of J[I. k).
Algorithm Compute kt,
1: Input: For given the number of codebook ne, the maximum batch time T, the current timestep
1, the number trajectory class nett and the index of trajectory class k
2: if = 0 then
3:
7kK = [m./na]
4: d ddnı//T nT/T
5: Keep the values of ng, d until the end of the episode
6: end if
7: i, = пи x (I: 1)
8: if > then
9:
[(1,k) = i, + [d x 1] 1: + +|dx (1 (1+1))
else
11:
J(t,k)= = I
end if
13:
Return Jt,
The computed [t., k)) for a given (t. k) pair is then used for coverage loss in Eq. 5 to spread the quantized vectors throughout the embedding space of feasible states, X. Since Eq. (5) is expressed for given s/ we further elaborate on the expression considering batch samples. Modified VQ-VAE loss objective for given state 8t, nearest vector Ft.g [[[()))lg and given class & is expressed as follows.
cot cVolo,e, = Lvold,e,st) + Aovi
1
Σ JEJ((,M
[|sg[f](s)) e
k(t
(12)
Lvolo)e,s))
(13)
(1 - sel1 + Av {vg||sg][g(s.))
x Acommit (lf sgl (
For batch-wise training for VQ-VAF, we train VQ-VAE with the following learning objective:
B T-1
C (0,e)
ΣΣ
tot C
(ò, e, $t,b,
(14)
B
17

---

Published as a conference paper at ICLR 2025
Algorithm 2 presents the overall training algorithm for all learning components of TRAMA, in- cluding f.. o e VQ-VAE; trajectory classifier fø: trajectory-class predictor Tç' trajectory-class representation model fg: action policy Q3 for i E I; and Q for mixer network.
Algorithm 2 Training algorithm for TRAMA
1: Parameter: Batch size B for MARL training, batch size M for classifier training, classifier update interval Rig. VQ-VAE, update interval ng. and the maximum training time T Tene Individual Q-network Q3 for n agents, trajectory-class representation model fg. VQ- VAE encoder fz. VQ-VAE decoder fz. VQ-VAE codebook e, trajectory-class predictor nç. replay buffer D, trajectory classifier f
2: Input:
3: Initialize network parameter 0, Φ, V
4: lenv =0
5: Mepisode = 0
6: while tenu S T Tene do
7: Interact with the environment via e-greedy policy with [Q)] and get trajectory Ts
8: lenv Lene + lepisode
9:
Repisode Repiaode +
10:
Encode Tec_o by f and get quantized latent sequence Txe=o = Lf by Eq. 1
Get indices sequence T2.=n from Txr=o Append - to D Get sample trajectories [(T T2 1 ~
12:
14:
for b < B do
15:
if kg is None then
16:
Get trajectory-class label ky via fs
17:
end end for if mod(nepisode: To. then
18:
19:
20:
Compute Loss Civeeh (o) e) by Eq. 14 with [[ Tz k11b11
21:
Update Φ, e
22:
end if
23:
Compute Loss L(A, by Eq. 10 with [{T Ter k
24:
Update 0.
25:
if mod(Repisode; Rey.) then
26:
Get M sample trajectories [T,
M
m Compute trajectory embedding (Em
27:
m==
28:
Get class labels K by K-means clustering of [em mm=1
M
Compute Loss C(b) by Eq. with [ Update s/
Jm=1 and K
30:
end
32:
end while
18

---

Published as a conference paper at ICLR 2025
D ADDITIONAL EXPERIMENTS
D.1 OMITTED EXPERIMENT RESULTS
In Section [ we evaluate the various methods based on their mean return values instead of the mean win-rate. In single-task settings, both win-rate and return values show similar trends since agents learn policy to defeat enemies with a fixed unit combination under marginally perturbed initial positions. On the other hand, in multi-task settings, agents' policies may converge suboptimality specialized on specific tasks, preventing them from gaining rewards in other tasks. In such a case, high win-rate derived by specializing in certain tasks does not guarantee the generalization of policies multi-task problems. Thus, we measure the mean return values instead. In the following, present the omitted win-rate performance of experiments presented in Section [5]
In Figures 15 and T6 TRAMA shows the best or comparable performance compared to other base- line methods. However, the performance gap between TRAMA and other methods not distinc- tively observed in these win-rate curves due to the aforementioned reason. Thus, measuring the mean return value more suitable for multi-task problems.
TRAMA (ours)
QMIX
RODE
LDSA
MASER
QPLEX
EMC
LAGMA
SurComb3
reSurComb3
SurComb4
reSurComb4
100
100
100
100
Test Win (%)
50
50
50
50
0.0
1.0
2.0
3.0 (mil)
4,0
5.0
0.0
1.0
2.0
3.0 (mil)
4.0
5.0
0.0
1.0
20
3.0 (mil)
4.0
5.0
0.0
1.0
20
3.0
4.0
5.0
(mill
Figure 15: The mean win-rate of TRAMA compared to baseline algorithms on four multi-task problems presented in Table [
TRAMA (ours)
QMIX
RODE
LDSA
MASER
GPLEX
EMC
p5 vs_6
15 v5_
100
00
Test Win 1h)
60
40-
20-
20-
0.0
1.0
2.0
3.0
0.0
1.0
2.0
3.0
T(ml))
T(mml
Performance comparison of TRAMA with win-rate against baseline algorithms on p5_vs_5 and 55_vs_5 in SMACv2. Here, n is assumed.
Figure 16:
On the other hand, in Figure 17 similar performance trends are observed, illustrating the better performance of TRAMA in some tasks.
D.2 PERFORMANCE COMPARISON WITH ADDITIONAL. BASELINE METHODS
In this section, we present additional performance comparison with some omitted baseline methods, such as Updet (Hu et al.) (02T) and RiskQ (Shen et al. (023). Updet utilizes transformer architecture for agent policy based on the entity-wise input structure. Please note that this input structure is different from the conventional input settings of a single feature vector, which are widely adopted in MARL baseline methods (Rashid et al. 2018 Wang et al. 2020a| Zheng et al. 2021 Na & Moon, 2024). Therefore, we modified the input structure provided by the environment to evaluate Updet in SMACv2 tasks. RiskQ introduces the Risk sensitive Individual-Glodde
(RIGM) to consider the common risk metrics such the Value at Risk (VaR) metric distorted risk measurements.
19

---

Published as a conference paper at ICLR 2025
(ours)
MASER
LAGMA
1c3352
5m _______mm 6m
MMM2
6h vs 8z
100
100
100
Test Vin 1%6
50
30
50
0.0
0.5
0.0
0.5
1.0 (mil)
1.5
2.0
0.5
1.5
2.0
1.0
2.0
4.0
T (mil)
(mil)
T (mil)
Figure 17: The mean win-rate of TRAMA compared baseline algorithms on SMAC task.
For evaluation, we consider four multi-task problems presented in Table | and the conventional SMACv2 tasks, such as p5_vs_5 and 5_vvs_5. use the default settings presented their codes for evaluation.
TRAMA
Updet
Mest Mean Relurn H
SurComb3
reSurCombb3
SurComb4
reSurCamb4
3-
3-
2
0
0.0
1.0
2.0
3.0
4.0
5.0
0.0
1.0
2.0
3.0
4.0
5.0
0.0
1.0
2.0
3.0
4.0
5.0
0.0
1.0
2.0
3.0
4.0
5.0
T
(mil)
(mil)
T
(mil)
Figure 18: The mean return comparison of various models multi-task problems.
TRAMA (ours)
RiskO
Updet
SurComb3
reSurComb3
SurComb4
reSurComb4
Test Win 1%6)
50
50
50
50
0.0
2.0
1.0
2.0
3.0
4.0
5.0
((mil)
T (mil)
(mi)
T (mil)
Figure 19: The mean win-rate comparison of various models on multi-task problems.
TRAMMA
200
15.20 1e.
20.0
ve.
17.5
17.5
Toshhon Ritmnl
15:0
15.0
12.5
7.8
1.0
28
2.5
0.0
1.0
2.0
3.0
010
0.0
1.0
2.0
3.
Timm:
T(mi)
Figure 20: The mean return comparison of various models on SMACv2 p5_vs_5 and t5_vs_5.
In Figures 18 21
TRAMA shows the best performance compared additional baseline methods, in terms of both return and win-rate in all multi-task problems.
20

---

Published as a conference paper at ICLR 2025
TRAMA-(OUn))
R36Q
Updrr
60
03
Tool Tosk/Wne05
80
40
40
Figure 21: The mean win-rate comparison various models on SMACv2 5 and t5_vs_5.
D.3 ADDITIONAL ABLATION STUDY
In this subsection, we present additional ablation studies on multi-task problems, such as SurComb3 and reSurComb4 evaluate the impact of trajectory-class representation 9 gener- ated by fg. To this end, we ablate f in TRAMA and consider the one-hot vector, instead of g', additional condition individual policies. We denote this model TRAMA (one-hot), In addition, we also ablate [t, k) and consider Jt) to understand further the role of [tt k). Figure 22 illusttaates the results.
30
300
25-
2.5
1.0-
100
0.5
06-
00
05
1.5
2.0
25
30
00
0.
1.0
1.
240
25
30
(a) SurComb3
(b) reSurComb4
Figure 22: Ablation studies on SurComb3 3 and reSurComb4.
Similar to the results in Figure 12 J(t, improves the performance as the quantized vectors are evenly distributed throughout X. yielding the clusters of trajectories with task similarities. the other hand, the one-hot vector also gives additional information to the policy network as agents predict the trajectory class labels accurately, generating consistent signals the policy for given trajectory class. However, trajectory-class representation signifies this impact and further improves performance of TRAMA.
We also conduct an additional ablation study the conventional SMACv2 tasks (Ellis et al. 2024) to see the effectiveness of the trajectory-class representation. Figure 23 illustrates the result.
2000
17.5
15.0
15.0
17.5
17.5
13.0
10.0
75
TRAMA
75
TRAMA bonal
TRAM (eee-oo
50-
TRAN iee
DFLEX
25
aPUEX
EVC
60
U1
1.
220
20
00
05
1/0
1.0
200
20
30
(a) p5_vs.5 task.
(b) 15 15_v8.5 task.
Figure 23: Additional ablation tests trajectory-class representation, g.
21

---

Published as a conference paper at ICLR 2025
that one-hot vector additional conditional information for decision- making also benefits the general performance. However, the trajectory- claass representation signifies this gain, illustrated by the comparison between TRAMA (our) and TRAMA (one-hot).
In Figure 25
D.4 ADDITIONAL ABLATION STUDY ON CLUSTERING MODULE
In this section, we study the effect of the distribution of quantized vector throughout the embedding space, X = {ir Rdd x = [ 8 E D). We utilize VQ-VAE embeddings to generate trajectory embeddings via Eq. 6 so that it can capture the commonality among trajectories. Although we do not have prior knowledge of [K], we can identify some tasks sharing similarity based on trajectory embeddings by assuming net- In addition, we can utilize some adaptive algorithm to determine optimized Ne.. Please see Appendix D.5 for an adaptive clustering method.
In addition, even though k and l22 are actually different tasks if their differences are marginal, then they can be clustered into the same trajectory class. This mechanism is important since similar tasks may require similar joint policy. Thus, having the exact trajectory class representation as an additional condition can be more beneficial in decision-makin than having vastly different trajectory-class representation, which could encourage different policies.
If quantized embedding vectors are not evenly distributed through the embedding space, the seman- tically dissimilar trajectories may share the same quantized vectors, which is unwanted. Without well-distributed quantized vectors, it becomes hard to construct distinct and meaningful clustering results. To see the importance of the distribution of embedding vectors in VQ-VAF, we ablate the coverage loss: (1) training with Aeer considering J(t) only, and (2) TRAMA model, i.e., training with Aoer considering [t, k). In addition to Figure in Section 3.1 Figure 24 presents additional ablation results on SMACv2 task p5_vs_5.
of
-2
(a) Training with ((t) (b) Training with [(t, k)
(c) Clustering of (a)
(d) Clustering of (b)
Figure 24: PCA of sampled trajectory embedding and VAE embedding vectors (gray circles). Colors from red to purple (rainbow) represent early to late timestep in (a) and (b). In (c) and (d), five clusters assumed (net = 5), and each trajectory embedding is colored with designated class (red, green, blue, purple and yellow). p5 vs 5 task is used for testing, and we ablate components related to coverage loss.
From Figures 4 and 24 we can see that having evenly distributed VQ- VAE embedding vectors is critical in clustering, which is highlighted by Silhouette score and the result of visualization. Since tasks in p5 vs 5 can have marginal differences, the clustering results are unclear compared to the results from the customized multi-task problems in Figure 4
D.5 IMPLEMENTATION DETAILS OF CLUSTERING AND ADAPTIVE CLUSTERING METHOD
This section presents some details of the K-means++ clustering we adopt and possible adaptive clustering method based on the Silhouette score. For centroid initialization, centroids are initially selected randomly from the data points. Then next centroid is selected probabilistically, where the probability of selecting a point is proportional to the square of its distance from the nearest existing centroid. At first, iterate K-means with 10 different initial centroids. However, in Section 3.2, generating coherent labels is also important. Thus, once we get the previous centroid value, this prior value as initial guess for the centroid and run K-means just once with it.
discussed
As discussed in Section 5.3 we may select n. large enough so that TRAMA can capture the possible diversity of unit combinations in multi-task problems. To consider some automatic update for nat,
22

---

Published as a conference paper at ICLR 2025
implement possible adaptive clustering algorithm by adjusting nat based the Silhouette scores candidate Rett values. We test this on multi-task problems and Figure 23 presents the result. We maintain other components in TRAMA the same. We initially assumee nad and ned 6 for adaptíve methods. In Figure [5] the adaptive clustering method succeeded in finding optimal Ad = nud in terms of Silhouette score. When initial value nett 2 is close the optimal value, the overall performance in terms of return shows better performance as it quickly converges to the value ne and generates coherent label information for trajectory-class predictor.
THAMM (mbapten ng 2
6.-
TRAMA (ndosfiee
30-
50-
20-
2.0
35
1.3
TRAAA (t-
25-
TRAMA Indaptive n5 21
(150
(a) Nal variation
(b) Return value
Figure 25: Test of adaptive clustering method on SurComb4 4 multi-task problem.
Interestingly, the presented adaptive method converges net 3 instead of net 4. Figure 13 the clusters for c and 3s2z are close and share common units, such as 2s2z. Thus, in terms of Silhouette score, Nat has a marginally higher score and thus adaptive method, which converged to Net = 3, yields similar performance in the case of fixed nit = 4.
23

---

Published as a conference paper at ICLR 2025
D.6 EXPERIMENTS ON OUT-OF- OF-DISTRIBUTIO (OOD) TASKS
In this section, we present additional experiments on Out-Of-Distribution (OOD) tasks. Here, for to Kood- define OOD tasks as Tk. such that Th. & T vain according Definition 2.1 In other words, vk. E Kood and Vk; E Kuain. Sk. n.sk * @ and S2 n * @ should be satishedd Thus, we can view OOD tasks as unseen tasks.
For this test, we use four models trained under different seeds for each method, and the correspond- ing win-rate and return curves are presented in Figure 26
85-
2-5
60-
1.5
CKX
TRANA tours)
100
QUIE
TRAMA [Dum)
OPLEX
ORLEX
EMC
EMC
00
T ((n]
(a) Winrate
(b) Return uEn
Figure 26: Performance comparison of various models SurComb4.
In this test, we construct OOD tasks by differentiating either the unit combination or initial position distribution so that agents experience different observation distributions throughout the episode. Table presents details of OOD tasks. We set two different types of unit combinations in OOD tasks, in addition to in-distribution (ID) unit combinations, (1) one similar to ID unit distribution and (2) the other largely different unit combinations. For example, unit combinations in OOD (#1) or OOD (#4) share some units with ID unit distribution. Specifically, 1c2s1z 1z are common units in both the ID task with 1 1c2s22 z and the OOD task with 1c3s12. 2. On the other hand, OOD (#2) or OOD (#5) accompanies largely different unit combinations.
In addition, we also construct OOD via largely different initial positions in OOD (#3) ~ OOD (#5). Thus, expect that OOD (#5) is the most -itto--distributed task among various task settings in 7
Table 7: Task configuration of OOD tests
Name
Initial Position Type
Unit Combinations (Ncomb)
ID (SurComb 4)
Surrounded
{1c2s2z, 3s2z, 2c3z, 2c3s)
OOD (#1)
Surrounded
{1c4s, 1c3s1z, 2s3z]
OOD (#2)
Surrounded
{5s, 52, 5c]
OOD (#3)
Surrounded and Reflected
f1c2s22, 3s22, 2c32, 2c3s} (1c4s, 1c3s1z, 2s3z)
OOD (#4)
Surrounded and Reflected Surrounded and Reflected
0OD (#5)
{5s, 5z, 5c}
Based on four models of each method trained with different seeds, we evaluate them across six different task settings, as shown in Table 7 For the evaluation, we run 50 test episodes per each trained model, resulting in a total of 4 x 50 200 runs per method. Tables and present test results. In the tables, the star marker (*) represents the best performance for given task among various methods.
In Tables Jand 8 TRAMA shows the better or comparable performance, in terms of both win-rate and return, in all cases including OOD tasks. In the OOD (#2) task, the differences in the win-rate are not evident compared to other baseline methods. This is reasonable because TRAMA cannot gain significant benefits from identifying similar task classes and encouraging similar policies through trajectory-claat classs representations when there is clear task similarity. In addition, TRAMA also shows the best performance in OOD tasks with highly perturbed initial positions, as in OOD (#3) ~ OOD (#5).
24

---

Published as a conference paper at ICLR 2025
Table Return of OOD tests
TRAMA
QMIX
QPLEX
EMC
(SurComb4) 4)
5.063"#0.48%
2..99
2.323+0.119
2.34840.406
0OD (#1)
2.706"
2.077=0:422
2.248+0.471
1.661#0.220
0OD (#2)
1.259"=0.148
0.939=0.463
1.04940.240
0.774+0.326
0OD (43)
2.307*=0.485
1.899=0.299
2.025#0.221
2.153#0.557
OOD (#4)
#0.270
1.884=0.219
1.993+0.199
1.457#0.22 127
0OD (#5)
0.921'#0.679
0.532#0.200
0. 0.792+0.282
0.605#0.295
Table 9: Win-rate of OOD tests
TRAMA
QMIX
QPLEX
EMC
ID (SurComb4)
0.707*#0.025
0.540=0.059
0.667:0.034
0.613#0.025
OOD (#1)
0..05
0.475=0.062
0.525#0.073
0.430=0.057
OOD (#2)
0.280*#0.00
0.270=0.057
0.265+0.017
0.275#0.057
0OD (#3)
4.620"=0.111
0.465=0.050
0.545+0.050
0.525#0.084
0OD (#4)
0.545' #0.052
0.355=0.071
0.475+0.055
0.365+0.033
0OD (#5)
0.160=0.037
0.215+0.050
0,205#0.062
Discrepancy between return and win-rate in multi-task problems In Tables 8 and the win-rate difference (in ratio) between TRAMA and QPLEX about 6% while the return difference (in ratio) is about 32%.
This discrepancy can be observed multi-task problems since a trained model can specialized on some tasks yet less effective in other tasks, making not much reward. For example, both models A and B succeed in solving task while failing on task 2 at the same frequency. However, model A nearly succeeds in solving task 2 but ultimately fails, while model B completely fails in solving task 2, this scenario results in a similar win-rate for and B but a different return.
In another case, both models A and B become specialized in of the tasks, but the total return from tasks 1 and 2 can differ. If task yields a larger return compared to task 2 for success, and model overfits to task while model B overfits task this scenario results in similar win rate for A and B but a different return, i.e., a larger return for model A.
In both cases, model is preferred over model B. This highlights why it is important to focus on the return difference when evaluating model performance in multi-task problems.
Qualitative analysis on trajectory class prediction in OOD tasks To understand how TRAMA predicts trajectory classes in OOD tasks, we first evaluate how accurately agents in TRAMA predict trajectory classes in-distribution (ID) tasks.
64
82
00
3622
.22
FD Le242Z
-04
059
023
100
15
22
269t
12
-04
(a) Clustering Results (b) Epi sode instances (ID)
Figure 27: Qualitative analysis on in-distribution task (SurComb4).
Figure a7)a presents the clustering results of the ID task to determine which cluster corresponds to which type of task. From this result, we identify types of task denoted by each trajectory class
25

---

Published as a conference paper at ICLR 2025
in Figure b)) and present test episodes #A #D. We also present the overall prediction made by agents across all timesteps in each test episode by the conventional box plot. The box plot denotes 1st (Q1) and 3rd (Q3) quartiles with color box and median value with yellow line within the color box.
(a) Case #A
(b) Case
4B
(c) Case
#C
(d) Case ID
Figure 28: Overall prediction on trajectory class made by agents (ID task). Please refer to Figure 27 for each episode case.
In Figure 28 agents are confident in predicting trajectory class. In Case #A, agents predict a possi- bility that the given task belongs to class 4 instead of class 3, their unit combinations can be quite similar, such as 1 lc2 z, when some units are lost.
2036
cb 2::0
9622
1065 DL
02
012
Wit Sc
-3.
240
9.
5.00
Sz
2,
$..00
de 20262
1.00
d S62522
2.50
-0
1.50
02
(a) 8pi sode instances (COD#1)
(b) Epi sode instances (000#2)
Figure 29: Qualitative analysis out-of-distribution tasks (OOD #1 and OOD #2).
Figures 29 present qualitative analysis results of out-of-distribution tasks, OOD#1 and OOD#2 and
Figures 30 and 31 illustrate the predictions made by agents for each test case. In Figure 30 agents predict the class of given disttiibbiiiino task as the closest class experienced during training. Some tasks in OOD#1 share some portion of unit combinations, yielding strong predictions on Case #B and Case #C. This may yield OOD task adaptation, as a predicted trajectory class representation can encourage joint policy that is effective in tasks sharing some units with II) tasks.
On the other hand, when there no evident similarity between OOD tasks and ID tasks, agents make weak predictions on trajectory class as presented in Figure 31 In this case, can detect highly OOD tasks by setting certain threshold of confidence level, such as 50%.
(a) Case #A
(b) Case #B
(c) Case #C
Figure 30: Overall prediction on trajectory class made by agents (OOD#1 task). Please refer to
Figure ( for each episode case.
26

---

Published as a conference paper at ICLR 2025
(a) Case #A
(b) Case #B
(c) Case #C
Figure 31: Overall prediction on trajectory class made by agents (OOD#2 task). Please refer to Figure b) for each episode case.
D.7 TRAJECTORY CLASS PREDICTION
In this section, we present the omitted results of trajectory class predictions made by agents. We present the accuracy of predictions in ratio. Tables I0 and 12 demonstrate the prediction accu- racy of each training time (tene). Columns of each Table represent each timestep (tepisode) within episodes. Throughout various multi-task problems, agents accurately predict the trajectory class. In the results, agents predict more accurately as the timestep proceeds since they get more information for prediction through their observations.
Table 10: The accuracy of trajectory class prediction (IM)
11
tepisode=
tepisode =10
tepisode=22
=
SurComb3
0.841#0.062
0.855+0.044
0.86140.044
0.886+0.043
SurComb4
0.766+0.085
0.791 104
0.848
0.876#0.088
reSurComb3
0.980+0.014
0.991#0.011
0.997+0.004
0.997+0.004
reSurComb4
0.838+0.092
0.900=0.044
0.920#0.055
0.934+0.046
Table 11: The accuracy of trajectory class prediction (3M)
1em0=3M
lepisode=0
(episode=11 =10
lepisode =20
Lepisode=
SurComb3
0.880+0.038
0.909=0.062
0.925#0.041
0.930#0.00
SurComb4
0.908+0.063
0.939+0.022
0.94110.0233
0.947+0.022
reSurComb3 0.947+0.046
0.961+0.029
0.975+0.017
0.968+0.023
reSurComb4
0.889#0.062
0.927=0.037
0.933+0.045
0.931
Table 12: The accuracy of trajectory class prediction (5M)
5e
tepssode=0
tepisode= =10
tepieode =20
tepisode=30
SurComb3
0.853+0.034
0.902+0.046
0.933+0.052
0.945+0.044
SurComb4
0.925+0.084
0.939+0.070
0.963+0.049
0.957+0.045
reSurComb3
0.963+0.07.
0.953+0.094
0.967+0.066
0.972+0.056
reSurComb4
0.88610.024
0.913+0.017
0.91740.035
0.930+0.027
27


Reviews:
[{"review_id": "a1Y4fVSwOn", "reviewer": "Reviewer_8EYy", "summary": "This paper focuses on utilizing multi-agent reinforcement learning (RL) for solving the multi-goal task. An approach of TRajectory-class-Aware Multi-Agent reinforcement learning (TRAMA) is proposed by generating the trajectory embeddings, making trajectory clustering, learning trajectory-class-aware policy with trajectory-class representation. Finally, experiments were conducted on SMACv2 domain, including four modified tasks (multi-goal tasks) and two conventional SMACv2 tasks, and compared TRAMA with QMIX, RODE, LDSA,  MASER, QPLEX, EMC, and LAGMA.", "questions": "- see my comments in the weaknesses.\n- Did the replay buffer store the trajectories only, instead of transitions?\n- How many agents were used? what are the number of trajectory classes and centroid? How did you select them?\n- How did you collect the trajectories? What's size of the replay buffer?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "- The multi-agent RL has the potential to solve more complex tasks. This paper aims to use multi-agent RL to solve the multi-goal task. The main contribution of this paper is by proposing a framework with doing clustering and prediction on the latent space and learning the policies based on the generated trajectory-class representation. Although the main idea is straightforward, there is some novelty in the algorithm design.", "weaknesses": "- The presentation should be improved. What's the formal definition of the multi-goal task? Why multi-agent RL can solve this better than the single-agent RL? What's the relationship between the multi-goal and trajectories? In Figure 2, what's the difference between \"$s_t$\" and \"$o_t^i$\"? what are the goals? How different are they? How to formulate them into the objective? How agents interact with each other and the environment, and learn various goals? Is the whole training process end-to-end? or are there some pre-training process? I do believe there needs some clarifications on the problem background, assumptions, and details of the methodology, currently it's not clear to me.\n- Empirical evaluation: what's the difference between the modified tasks and conventional tasks on SMACv2? Are they all multi-goal settings? What are these goals? In the results shown in Figure 6, it looks the improvement of TRAMA over others is marginal. For Figures 7 and 8, it's better to report the numbers instead of showing the curves. In conventional tasks on SMACv2, it looks TRAMA barely didn't have the advantage over others. If this is the case, how could TRAMA work better on multi-goal tasks than other methods? It feels like the methodology is quite complex, but the benefit is small. Overall, the current conclusion is still vague to me.", "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": "EA3YGGaqet", "reviewer": "Reviewer_9FqX", "summary": "The paper introduces a new multi-agent reinforcement learning framework called TRAMA which aims to address the challenge of multi-goal tasks. The multi-goal tasks are especially challenging since they require agents to adapt to goals without relying on task-specific policies. TRAMA incorporates the vector quantize variational autoencoder (VQ-VAE) approach to create a discrete space for trajectory embeddings. With clustering, the embeddings provide information on which trajectory class agents are experiencing and then can be used to improve the decision making process. Experiments are conducted over a suite of SMACv2 tasks and the proposed method achieves superior performance.", "questions": "1. Since the clustering process is trained on a dataset, would the model be able to detect or adapt to out-of-distribution trajectories which better simulate real-life applications? \n2. Given the unsupervised nature of VQ-VAE learning and trajectory clustering, how can the model ensure that the resulting clusters are distinct and meaningful with respect to specific tasks?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper is well written and easy to follow. The approach to incorporate VQ-VAE for identifying goals is a novel branch of study in MARL. The authors conducted extensive experiments demonstrating the capability of the model. Notably the paper provides insightful visualizations, such as the VQ codebook visualization, which makes the result more convincing. Overall the proposed method achieves a significant improvement in performance in many different settings in SMACv2.", "weaknesses": "1. One major concern I have is regarding the novel contribution of TRAMA as compared to the LAGMA framework. Since LAGMA proposed adapting the VQ-VAE framework to multi-agent RL, the paper may benefit from a more detailed discussion on the unique contributions brought by TRAMA. \n2. The motivation behind applying k-means clustering to the VQ-VAE embeddings requires further clarification. Since the embeddings are already quantized, it would be helpful to further explain why the k-means clustering is still necessary \n3. Following the previous point, the authors may consider providing additional insights into the stability of the k-means clustering process. For me while the result in figure 13 is indeed encouraging, it is in the meantime counterintuitive to see why VQ-VAE embeddings consistently form clear clusters. Additional explanation would strengthen the arguments.", "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": "XqwRNC3IAQ", "reviewer": "Reviewer_q68e", "summary": "This paper proposes a novel multi-agent reinforcement learning framework called TRAMA (TRajectory-class-Aware Multi-Agent reinforcement learning) to address the issue of policy generalization in multi-goal tasks. TRAMA introduces the idea of trajectory clustering, where trajectories are embedded and clustered to enable agents to recognize the type of task they are experiencing based on their trajectories, and use this prediction as additional information for decision-making, thereby learning policies that can adapt to different goal tasks. The main components of TRAMA include: 1) constructing a quantized latent space using a modified VQ-VAE to generate trajectory embeddings; 2) performing trajectory clustering based on the embeddings; 3) training a trajectory-class predictor and class representation model to generate trajectory-class-aware policies. TRAMA achieves better performance than multiple baseline methods on various multi-goal tasks in StarCraft II.", "questions": "1. The paper does not theoretically analyze the superiority of TRAMA. The authors only demonstrate the advantages of TRAMA over other methods from experimental results, lacking rigorous theoretical derivation and complexity analysis. For example, does the introduction of trajectory clustering and class prediction modules significantly increase the computational overhead? No quantitative analysis is provided.\n2. The experimental evaluation is not comprehensive enough. The paper only constructs and tests 4 multi-goal tasks in the StarCraft II environment. However, there are many ways to compose multi-goal tasks. Can these 4 tasks represent the main scenarios in the real world? Moreover, the authors only compare TRAMA with a few representative MARL algorithms. Are there other multi-goal generalization methods that are not included in the comparison? More benchmarks and baselines would make the experimental results more convincing.\n3. The number of trajectory classes is predetermined, and the paper does not provide a principled method for setting the classes. It only discusses the impact of different class numbers on performance in the ablation study. So how to adaptively determine the optimal number of classes based on task characteristics? What are the criteria for class division? These questions require further exploration.\n4. This paper only focuses on tasks with discrete action spaces. Is it equally applicable to continuous action spaces? Additionally, in other types of multi-agent tasks, such as mixed cooperative-competitive tasks, is trajectory clustering still effective? These issues need to be analyzed and validated in future work.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. TRAMA effectively addresses the policy generalization problem in multi-goal tasks. Traditional MARL methods often learn policies that are only effective for specific tasks and lack generalizability. By enabling agents to predict trajectory classes and learn class-related policy representations, TRAMA significantly improves the adaptability of policies, which is validated by the experimental results showing that TRAMA achieves higher average returns than other methods on multiple multi-goal tasks (Figure 6).\n2. The modified VQ-VAE quantization method better learns trajectory embeddings. By considering both time steps and trajectory classes in the coverage loss, the quantized vectors are more evenly distributed in the embedding space of different trajectory classes (Figure 3), providing a good foundation for subsequent clustering.\n3. Trajectory clustering and the class predictor allow agents to accurately infer the trajectory class from local observations. Experiments show that after training, the agents' prediction accuracy for trajectory classes can stabilize at a high level (Figures 7, 8), enabling the policy to make class-related decisions based on the class representation.\n4. Ablation studies and parameter analysis investigate the impact of main modules and hyperparameters (Figures 11, 12, 18), enhancing the interpretability and reproducibility of the method.", "weaknesses": "1. The paper does not theoretically analyze the superiority of TRAMA. The authors only demonstrate the advantages of TRAMA over other methods from experimental results, lacking rigorous theoretical derivation and complexity analysis. For example, does the introduction of trajectory clustering and class prediction modules significantly increase the computational overhead? No quantitative analysis is provided.\n2. The experimental evaluation is not comprehensive enough. The paper only constructs and tests 4 multi-goal tasks in the StarCraft II environment. However, there are many ways to compose multi-goal tasks. Can these 4 tasks represent the main scenarios in the real world? Moreover, the authors only compare TRAMA with a few representative MARL algorithms. Are there other multi-goal generalization methods that are not included in the comparison? More benchmarks and baselines would make the experimental results more convincing.\n3. The number of trajectory classes is predetermined, and the paper does not provide a principled method for setting the classes. It only discusses the impact of different class numbers on performance in the ablation study. So how to adaptively determine the optimal number of classes based on task characteristics? What are the criteria for class division? These questions require further exploration.\n4. This paper only focuses on tasks with discrete action spaces. Is it equally applicable to continuous action spaces? Additionally, in other types of multi-agent tasks, such as mixed cooperative-competitive tasks, is trajectory clustering still effective? These issues need to be analyzed and validated in future work.", "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": "cGQ68bm6Bw", "reviewer": "Reviewer_YpdJ", "summary": "This paper studies multi-task multi-agent reinforcement learning. It proposes a trajectory clustering-based approach, where trajectories from different environments or tasks are clustered based on similarities in the latent space. The clustering label is then used as an additional input for the policy。", "questions": "[1] It seems that the policy module does not propagate gradients back to the clustering module. I wonder why these two modules are not coupled, allowing the clustering module to have some trainable components that could help extract useful information for the policy?\n\n[2] The paper emphasizes the coverage loss significantly. However, why is distributing throughout the embedding space considered beneficial?\n\n[3] In equation (7), why do you directly sum the embedding vectors? Does this make sense, or could there be better ways to handle this?\n\n[4] How do you manage different observation and input dimensions at the same time for different tasks? \n\n[5] \"To evaluate performance, we consider the overall return value instead of the win-rate, as the learned policy may be specialized for specific goal tasks while being less effective for others in multi-goal tasks.\" Could you elaborate further on this point?\n\n[6] For comparison with QMIX, do you use PyMARL2 (see 1') or PyMARL? PyMARL2 fine-tunes the parameters for QMIX and shows better performance. For example, in Figure 10, in the 6h_vs_8z scenario, based on my experience, QMIX could achieve a better return at 4M environment steps.\n\n\n1', Hu, J., Jiang, S., Harding, S. A., Wu, H., & Liao, S. W. (2021). Rethinking the implementation tricks and monotonicity constraint in cooperative multi-agent reinforcement learning. arXiv preprint arXiv:2102.03479", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Multi-task multi-agent reinforcement learning is an important challenge for MARL. The proposed trajectory clustering approach is innovative and offers a novel solution.", "weaknesses": "[1] The primary concern is whether the clustering approach is truly necessary. From my understanding, the purpose of clustering is merely to provide a label as additional information for the policy. It seems that the proposed approach only evaluates tasks seen during training, without considering unseen tasks during test time. Does this approach generalize to unseen tasks? If so, could the authors provide experimental results to support this? If not, why is the clustering necessary to provide a label, since we could easily have labels for the tasks during training and use it to train the trajectory class predictor? Perhaps simply using the task ID as a one-hot vector could replace the entire clustering module, as this has been effective in some single-agent RL work, see 1'.\n\n[2] The notations in Section 3 are unclear and could lead to confusion. Please consider improving them.\n\n[3] I think some discussions and comparisons with Updet ,see 2', is necessary, since Updet could be thought of a multi-task marl algorithom\n\n\n\n1' Hansen, N., Su, H., & Wang, X., 2023. TD-MPC2: Scalable, robust world models for continuous control. arXiv preprint arXiv:2310.16828\n\n2' Hu, S., Zhu, F., Chang, X., & Liang, X. (2021). Updet: Universal multi-agent reinforcement learning via policy decoupling with transformers. arXiv preprint arXiv:2101.08001.", "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)