Multi-Agent Policy Transfer via Task Relationship Modeling
Abstract
Team adaptation to new cooperative tasks is a hallmark of human intelligence, which has yet to be fully realized in learning agents.
Previous works on multi-agent transfer learning accommodate teams of different sizes, but heavily rely on the generalization ability of neural networks for adapting to unseen tasks.
We posit that the relationship among tasks provides the key information for policy adaptation.
To utilize such relationship for efficient transfer, we try to discover and exploit the knowledge among tasks from different teams, propose to learn effect-based task representations as a common latent space among tasks, and use it to build an alternatively fixed training scheme.
We demonstrate that the task representation can capture the relationship among teams and generalize to unseen tasks.
As a result, the proposed method can help transfer learned cooperation knowledge to new tasks after training on a few source tasks, and the learned transferred policies can also help solve tasks that are hard to learn from scratch.
Introduction
Cooperation in human groups is characterized by resiliency to unexpected changes and purposeful adaptation to new tasks [40].
This flexibility and transferability of cooperation is a hallmark of human intelligence.
Computationally, multi-agent reinforcement learning (MARL) [61] provides an important means for machines to imitate human cooperation.
Although recent MARL research has made prominent progress in many aspects of cooperation, such as policy decentralization [27, 33, 43, 50, 5], communication [11, 21, 53, 14], and organization [20, 45, 46], how to realize the ability of group knowledge transfer is still an open question.
Compared to single-agent knowledge reuse [65], a unique challenge faced by multi-agent transfer learning is the varying size of agent groups.
The number of agents and the length of observation inputs in unseen tasks may differ from those in source tasks.
To solve this problem, existing multi-agent transfer learning approaches build population-invariant [26] and input-length-invariant [49] learning structures using graph neural networks [1] and attention mechanisms like transformers [17, 64].
[Table]
[Figure]
Although these methods handle varying populations and input lengths well, their knowledge transfer to unseen tasks mainly depends on the inherent generalization ability of neural networks.
The relationship among tasks in MARL is not fully exploited for more efficient transfer.
Towards making up for this shortage, we study the discovery and utilization of common structures in multi-agent tasks and propose Multi-Agent Transfer reinforcement learning via modeling TAsk
Relationship (MATTAR).
In this learning framework, we capture the common structure of tasks by modeling the similarity among transition and reward functions of different tasks.
Specifically, we train a forward model for all source tasks to predict the observation, state, and reward at the next timestep given the current observation, state, and actions.
The challenge is how to embody the similarity and the difference among tasks in this forward model, we specifically introduce difference by giving each source task a unique representation and model the similarity by generating the parameters of the forward model via a shared hypernetwork, which we call the representation explainer.
To learn a well-formed representation space that encodes task relationship, an alternative-fixed training method is proposed to learn the task representation and representation explainer.
During training, representations of source tasks are pre-defined and fixed as mutual orthogonal vectors, and the representation explainer is learned by optimizing the forward model prediction loss on all source tasks.
When facing an unseen task, we fix the representation explainer and backpropagate gradients through the fixed forward model to learn the representation of the new task by a few samples.
Furthermore, we design a population-invariant policy network conditioned on the learned task representation.
During policy training, the representations for all source tasks are fixed, and the policy is updated to maximize the expected return over all source tasks.
On an unseen task, we obtain the transferred policy by simply inserting the new task representation into the learned policy network.
On the SMAC [34] and MPE [27] benchmarks, we empirically show that the learned knowledge from source tasks can be transferred to a series of unseen tasks with great success rates.
We also pinpoint several other advantages brought by our method.
First, fine-tuning the transferred policy on unseen tasks achieves better performance than learning from scratch, indicating that the task representation and pre-trained policy network provide a good initialization point.
Second, training on multiple source tasks gets better performance compared to training on them individually and other multi-task learning methods, showing that MATTAR also provides a method for multi-agent multi-task learning.
Finally, although not designed for this goal, our structure enables comparable performance against single-task learning algorithms when trained on single tasks.
Method
In this paper, we focus on knowledge transfer among fully cooperative multi-agent tasks that can be modeled as a Dec-POMDP [28] consisting of a tuple G = ⟨
I, S, A, P, R, Ω , O, n, γ ⟩ , where I is the finite set of n agents, s ∈ S is the true state of the environment, and γ ∈ [0 , 1) is the discount factor.
At each timestep, each agent i receives an observation o
i ∈ Ω drawn according to the observation function
O ( s, i ) and selects an action a i ∈ A .
Individual actions form a joint action a ∈ A n , which leads to a next state s ′ according to the transition function P ( s ′ | s, a ) , and a reward r =
R
( s, a ) shared by all agents.
Each agent has local action-observation history τ i ∈
T ≡ (Ω × A ) ∗ × Ω .
Agents learn to collectively maximize the global action-value function
Q tot ( s, a ) =
E s 0: ∞ ,a
0: ∞ [ (cid:80) ∞ t =0 γ t R ( s t , a t ) | s
= s, a 0 =
a ] (a little notation abuse: the subscript here for s and a indicates the timestep while the subscript for observation and action elsewhere in this paper indicates the index of the agent).
Overall, our framework first trains on several source tasks {S i } and then transfers the learned cooperative knowledge to unseen tasks {T j } .
As shown in Fig.
1, our learning framework achieves this by designing modules for (1) task representation learning and (2) policy learning.
In the following sections, we first introduce how we design the representation learning module and its learning scheme in different phases.
Then, we describe the details of policy learning, including the population-invariant structure for dealing with inputs and outputs of varying sizes.
2.1 Task representation learning
Our main idea in achieving knowledge transfer among multi-agent tasks is to capture and exploit both the common structure and the unique characteristics of tasks by learning task representation.
A task distinguishes itself from other tasks by its transition and reward functions.
Therefore, we incorporate our task representation learning component into the learning of a forward model that predicts the next state, local observations, and reward given the current state, local observations, and actions.
We associate each task i with a representation z i ∈ R m and expect it to reflect the relationship of the transition dynamics of tasks.
For modeling task similarity, all source and unseen tasks share a representation explainer, which takes as input the task representations and outputs the parameters of the forward model.
We then train the representation explainer on all source tasks.
Concretely, for a source task S i , parameters of the forward model are generated as η i =
f θ ( z i ) , where θ denotes the parameters of the representation explainer.
The forward model contains three predictors: f sη i , f oη i , and
f rη i .
Given the current state s , agent j ’s observation o j , and action a j , these predictors estimate the next state s ′ , the next observation o ′ j , and the global reward r , respectively.
A possible method for training on source tasks is to backpropagate the forward model’s prediction error to update both the representation explainer and task representations.
However, in practice, this scheme leads to representations with very small norms, making it difficult to get a meaningful representation space.
To solve this problem, we propose pre-determining the task representation for each source task and learning the representation explainer by backpropagating the prediction error.
Such a method can help form an informative task representation space and build a mapping from task representation space to the space of forwarding model parameters.
In practice, we initialize source task representations as mutually orthogonal vectors.
Specifically, we first randomly generate vectors in R m for source tasks, and then use the Schimidt orthogonalization [4] on these vectors to obtain source task representations.
With the pre-defined task representations, the representation explainer is optimized to minimize the following loss function:
[Equation]
where N src is the number of source tasks, and
[Equation]
is the per-task prediction loss.
Here, D is the replay buffer, and λ 1 , λ 2 are scaling factors.
We fix the source task representations and learn the representation explainer during the training phase.
In the current implementation, the training data here is collected with the uniform random policy; the data is the same for task representation learning in the transfer phase.
We find it simple and effective in the experiment, while more efficient approaches may be adopted here.
When it comes to the transfer phase, we aim to find a good task representation that can reflect the similarity of the new task to source tasks.
To achieve this goal, we fix the trained representation explainer and learn the task representation by minimizing the prediction loss of the forward model on the new task.
Specifically, we randomly initialize a task representation z , keep θ fixed, and get the forward model parameterized by η = f θ ( z ) .
Then the task representation z is updated by backpropagating the prediction loss for transition and reward functions through the fixed f η .
To keep the new task representation in the well-formed space of source task representations, we learn the new task representation as a linear combination of source task representations:
[Equation]
[Figure]
InvariantMixing Layers
[Figure]
𝑄
[Figure]
ℎ
[Figure]
In this way, what we are learning is the weight vector µ .
To make the learning more stable, we additionally optimize an entropy regularization term H ( µ ) .
The final loss function for learning z is:
[Equation]
The detailed architectures for task representation learning are described in App.
E.
2.2 Task policy learning
After the task representation is learned by modeling the transition and reward functions, it will be used to learn and transfer the policy on the source and unseen tasks.
A difficulty faced by multi-agent transfer learning is that the dimensions of inputs and outputs vary across tasks.
We use a population-invariant network (PIN, Fig. 2) to solve this problem.
The idea is not novel, and our contribution here is a light-weight structure that achieves comparable or better performance to other complex PIN modules [19].
The single task experiments in App.
G demonstrate its advantage as the performance gains mainly come from our PIN design.
Our PIN uses the value decomposition framework and consists of two main components: an individual Q-network shared by the agents, and a monotonic mixing network [33] that learns the global Q-value as a combination of local Q-values.
For the individual Q-network, like in previous work [19], we decompose the observation o i into parts relating to the environment o env i , agent i itself o own i , and other entities o other i =
{ o other j ̸ = i i } .
We adopt attention mechanism to get a fixed-dimensional embedding h :
[Equation]
where [ · , · ] is the vector concatenation operation, d k is the dimension of the query vector, and bold symbols are matrices.
Embedding h is then fed into the subsequent network together with the task representation z for estimating action values.
For the mixing network, we decompose the state s into parts relating to different entities s entity and
the environment s env .
We apply a similar selfattention module to integrate information from these parts of the state and obtain a fixed-dimensional embedding vector h state .
h state , the task representation z , and s entity are used to generate the parameters of the mixing network.
In some multi-agent tasks, the number of actions also varies in different tasks.
We use a mechanism similar to other popular population-invariant networks [48, 17] to deal with this issue, which is discussed in detail in App.
F.
During the whole process of policy learning, we fix the task representation z .
Compared to policy learning, which typically lasts for 2 M timesteps, the training of task representation costs few samples.
In practice, we collect 50 K samples for learning task representations and the representation explainer.
When transferring to new tasks , we use the individual Q-network and the representation explainer trained on source tasks.
We learn the task representation for 50 K timesteps and insert it into the individual Q-network.
Agents execute in a decentralized manner according to this Q-network.
[Table]
[Table] [Table]
Experiments
In this section, we design experiments to evaluate the following properties of the proposed method
: (1) Generalizability to unseen tasks.
Our learning framework can extract knowledge from multiple source tasks and transfer the cooperation knowledge to unseen tasks, and task representations play an indispensable role in transfer (Sec. 3.1).
(2) A good initialization for policy fine-tuning.
Fine-tuning the transferred policy can succeed in super hard tasks, which can not be solved when learning from scratch (Sec. 3.2).
(3) Benefits of multi-task training.
Our multi-task learning paradigm helps the model better leverage knowledge of different source tasks to boost the learning performance compared to training on source tasks individually (Sec. 3.3).
(4) Performance advantages on single tasks.
Although
we did not design our framework for single-task learning, our network performs better against the underlying algorithm (App. G).
We evaluate MATTAR on the SMAC [34] and MPE [27] benchmarks.
To better fit the multi-task training setting, we extend the original SMAC maps and sort out three task series, which involve various numbers of Marines, Stalkers/Zealots, and Marines/Maneuvers/Medivacs, respectively.
For the MPE tasks, we test on a discrete version of Spread and Gather , with different numbers of agents trying to cover all or one landmarks.
The detailed environmental settings are described in App.
B.
To ensure fair evaluation, we carry out all the experiments with five random seeds, and the results are shown with a 95% confidence interval in all figures.
For a more comprehensive description of experimental details, please refer to App.
D.
[Table] [Table]
3.1 Generalizability to unseen tasks
As the major desiderata of our method
[Table]
we expect that MATTAR has better transfer performance on unseen tasks.
We note that there are few related works applicable to tasks of varying sizes
[Table]
and we compare our method against the state-of-the-art multi-agent transfer method UPDeT [17] and
REFIL [19].
UPDeT transfers knowledge from a single source task.
For a fair comparison, we transfer from each source task to every unseen task and calculate the best (UPDeT-b) and mean (UPDeT-m) performance on each unseen task.
For the test phase, we conduct transfer experiments on both source tasks and unseen tasks.
Results on the SMAC benchmark are shown in Tab. 1 ∼ 3, and those on the MPE are shown in Tab. 4 and 5.
We find that MATTAR shows superior transfer performance on unseen tasks, significantly outperforming UPDeT-b and REFIL, especially in complex settings requiring sophisticated coordination like the MMM series.
MATTAR is effective even for unseen tasks that are very different from source tasks.
For example, the MATTAR policy learned on 2s3z , 3s5z , and 3s5z _ vs _ 3s6z can win 79% of games on 1s8z .
On MPE tasks, MATTAR also shows good generalization performance on unseen tasks with different numbers of agents and overall outperforms the baselines and ablations.
For example, On Spread tasks, MATTAR exhibits significant advantage over UPDeT on all unseen tasks and outperforms REFIL with a margin of 13% on the unseen task with 6 agents.
To explain the performance of our method, we first check the learned representations on unseen tasks and then carry out ablation studies.
What representations are learned for unseen tasks?
When encountering an unseen task, we first learn its representation as a linear combination of the representations for source tasks.
Specifically, we directly update the coefficients of this linear combination by backpropagating the prediction error of the forward model.
For a deeper understanding of how our method transfers the learned knowledge, here we investigate the learned coefficients of the linear combination.
In Tab. 6, we show the coefficients of two unseen tasks for each series of maps.
We observe that the largest coefficient typically corresponds to the source task that is the most similar to the unseen task.
For example, 5m is the closest source task to the unseen task 4m , and the coefficient of 5m is 61% .
There are also some exceptions.
For example, for the unseen task 10m _ vs _ 12m , the coefficients of two source tasks, 5m and 10m _ vs _ 11m , are equal, and they together take up 86% of all the coefficients.
While 10m _ vs _ 11m is very similar to 10m _ vs _ 12m as the team composition is similar, the policy for solving 5m is different from that for 10m _ vs _ 12m , we find that agents first form some groups to set up an attack curve quickly in 10m _ vs _ 12m .
Therefore, on a local battlefield, there are around five
[Figure]
fighting against a similar number of enemies.
In this case, the policy learned from 5m can be used locally to promote coordination in 10m _ vs _ 12m .
We conclude that, in MATTAR, task representation learning captures the similarity of task dynamics and thereby discovers opportunities of policy reuse.
These representations help unseen tasks effectively leverage the knowledge from the most similar source tasks or reuse the knowledge from a mixing of source tasks.
Ablations .
To further investigate the role of task representations in our method, we design two ablations.
(1) 0 task rep . uses the trained MATTAR models but feed a zero-valued task representation into the policy network when transferring.
(2) w / o task rep . denotes completely remove components related to task representation, including the forward model, the representation learning process, and the task representation in the policy network.
Train and transfer the population-invariant policy.
This ablation can reveal the generalization ability of the policy network itself.
Results are shown in Tab. 1 ∼ 3.
We can see that these ablations bring about a large drop in the performance of MATTAR on nearly all unseen tasks.
For example, after trained on 2s3z , 3s5z , and 3s5z _ vs _ 3s6z , MATTAR achieves a win rate of 0 .
79 on 1s8z , while w / o task rep . obtains a win rate of 0 .
12 and 0 task rep . does not win any games.
We can thus conclude that task representations play an indispensable role in policy transfer .
3.2 A good initialization for policy fine-tuning
When evaluating the performance of MATTAR on unseen tasks, we only train the task representations and reuse other parts of the policy.
In this section, we investigate the performance of MATTAR after fine-tuning the policy.
Specifically, on an unseen task, we first fix the representation explainer and train the task representation z for 50 K timesteps.
Then we fix the learned task representation and finetune the policy for 2 M timesteps.
We compare against learning from scratch in Fig.
Learning from scratch means the policy is randomly initialized and trained for 2 M timesteps.
We consider two versions for learning from scratch.
The first version ( Learn from scratch w / o repr . ) totally removes task representations from the policy, while the second version ( Learn from scratch w / repr . )
[Figure]
We observe that the task representation and the reused policy provide a good initialization.
For example, on 10m _ vs _ 12m , after fine-tuning for 2 M timesteps, MATTAR fine - tuning converges to a win rate of around 0 . 86 , while both versions of learning from scratch can only achieve a win rate of about 0 . 4 .
Furthermore, empirically only MATTAR fine - tuning achieves non-zero winning rates on 3s5z _ vs _ 3s7z , a task harder than the super hard map 3s5z _ vs _ 3s6z .
3.3 Benefits of multi-task learning
MATTAR adopts a scheme where multiple sources are learned simultaneously.
Our aim is to leverage knowledge from more tasks and to be able to generalize the learned knowledge to a larger set of unseen tasks.
Empirically, we find that this multi-task training setting helps not only unseen tasks but also the source tasks themselves.
In Fig.
4, we present the performance of MATTAR on source tasks when training with multiple and a single source task.
The experiments are carried out on three tasks from three different series.
We can see that training on multiple tasks significantly boosts learning performance.
For example, on 3s5z _ vs _ 3s6z , after 2 M training samples, MATTAR with multiple tasks converges to the win rate of around 0 . 6 , while training solely on this task can only achieve a win rate of about 0 . 05 .
Furthermore, we can also observe that MATTAR significantly outperforms the state-of-the-art deep multi-agent multi-task method (REFIL [19]).
These results reveal that MATTAR also provides a good learning framework for multi-agent multi-task learning.
It can leverage experience on other tasks to improve performance on a similar task.
Related Work
Multi-task learning with metadata [56, 63] approaches have been explored to exploit richer context information in the form of learned task embeddings, with a focus on task relation discovery.
[36] use context information in high-level, under-specified, and unstructured natural language for multi-task learning.
Compared to the settings studied in this paper, multi-agent problems pose additional challenges.
Transition dynamics and reward functions in multi-agent tasks are influenced by the interaction among all agents so the relation between tasks has to be conditioned on the interaction pattern of multiple agents.
Different from previous contextual multi-task RL methods, we include agent interaction modeling in the task representation learning module to deal with this problem.
Nevertheless, we think incorporating natural language as prior high-level metadata of tasks can further improve the performance of our method in further work.
Related to our work, context-based meta-RL also reuses experience from multiple tasks and infers how to solve a new task from collected small amounts of experience [66].
For example,
PEARL [32] performs online task inference by probabilistic filtering of latent task variables.
Other meta-RL methods exploit the dynamics of recurrent networks for implicit task inference and fast adaptation [42, 10].
These methods tend to form dynamic variables or embeddings, while the context information is utilized to obtain fixed task representations in our approach.
Besides, although their ability of task adaptation has been proved, context-based meta-RL has not been extended to multiagent cases and cannot deal with varying numbers of actions and different lengths of observation inputs.
For future work, it is important to discuss whether the inspiration of context-based meta-RL can further improve the performance of our multi-agent multi-task learning framework.
Learning task or skill embeddings for multi-task transfer reinforcement learning has also been extensively explored.
When testing, [15] learn a new skill embedding by interpolating learned skill embeddings on source tasks by approximate variational inference.
[2] also focus on learning skill embeddings but is conditioned on optimal Q-functions for different skills.
[6] adopt a hierarchical framework and learns a high-level policy to control a low-level skill latent space.
Similar to ours, [6] learn the low-level skill space by encoding experience trajectories and decoding states and actions.
Another work that is similar to ours is [60], which conditions the policy on latent representations learned by dynamics and reward module.
[24] learn task embeddings, fixes the policy, and updates the encoder at the test time.
Compared to these methods, our task representation learning method differs by (1) it is specially designed for multi-agent settings and considers varying input sizes and intra-agent interaction; and (2) to the best of our knowledge, the alternatively fixed learning scheme is proposed for the first time.
More contents include multi-task reinforcement learning, multi-agent transfer learning, multi-agent representation learning, and modular RL are shown in App. A.
Conclusion
In this paper, we study cooperative multi-agent transfer reinforcement learning by learning task representations that model and exploit task relationship.
Previous work on multi-agent transfer learning mainly deals with the varying population and input lengths, relying on the generalization ability of neural networks for cooperation knowledge transfer, ignoring the task relationship among tasks.
Our method improves the transfer performance by learning task representations that capture the difference and similarities among tasks.
When facing a new task, our approach only needs to obtain a new representation before transferring the learned knowledge.
Taking advantage of task relationship mining, our proposed method MATTAR achieves the best transfer performance and other bonuses compared with multiple baselines.
An important direction in the future is the transfer among tasks from different task distributions, and whether a linear combination of source tasks’ representations can fully represent unseen tasks is also a valuable topic.  Acknowledgements
This work is supported by the National Key Research and Development Program of China (2020AAA0107200), the National Science Foundation of China (61876119, 61921006), and the
Natural Science Foundation of Jiangsu (BK20221442).