Kaleidoscope: Learnable Masks for Heterogeneous Multi-agent Reinforcement Learning

Thank you for your insightful feedback. We appreciate your questions and suggestions and have provided clarifications below. Please let us know if you have any follow-up questions or comments; we would be happy to discuss further.

Q1: The authors do not include FuPS+ID ... a baseline in prior work [1,2] and is the most common ...

A1: We apologize for the misunderstanding. The FuPS in the original manuscript in fact refers to FuPS+ID, as the reviewer mentioned. We agree that adding agent IDs is the default PS implementation in MARL, and therefore we did not explicitly mention it. We will rename it as FuPS+ID for clarity.

Q2: .. missed a relevant related work [2] that also uses pruning masks for adaptive parameter sharing in MARL...

A2: We appreciate the suggestion and will incorperate discussion regarding this concurrent work into the related work section: AdaPS [2] combines SNP and SePS by proposing a cluster-based partial parameter sharing scheme.
It is also worthnoting that in our manuscript, "adaptive" refers to whether the sharing scheme evolves during training (see Table 1 in our manuscript). Although more flexible than SePS, AdaPS is not adaptive in this sense because its masks are initialized at the beginning of training and kept fixed throughout the training. Without the learnable masks, AdaPS cannot adaptively adjust the parameter sharing scheme during the training to boost the performance. Its performance highly relies on good initialization. A fixed partial parameter sharing scheme in general can not exhance the joint network representational capacity as our proposed Kaleidoscope does (see analysis in line 136-143 in our manuscript).

Q3: The motivation for using ... different weights for each agent, versus ... different activations ... is not clear to me. I understand that STR learns the pruning schedule itself and this is useful, but it is not clear that learning weights is the best way...

A3: We appreciate this insightful question. In fact, we also considered this when designing our method. During the rebuttal period, we conducted an experiment to justify this design choice, with results provided in Fig. 9 in the response PDF. As the results indicate, masking weights is more effective than masking activations. The rationale is that masking weights offers better flexibility. In a fully-connected layer, masking an activation is equivalent to masking all the weights connected to that activation. In Kaleidoscope, masks are learnable. In mask learning, better flexibility corresponds to more parameters to optimize, increasing the chance of finding optimal masks.
We agree this is an important design choice to justify and will incorporate these new results and discussions into our manuscript.

To sum up, comparing with baselines, our Kaleidoscope enjoys better flexibility at two levels:

• The learnable nature of masks provides adaptive PS during training (as you already mentioned).
• The weight-level sharing allows for partial sharing that is more precise (as verified by this additional experiment).

Q4: ... not clear that the proposed method is significantly better than the baselines... I recommend using IQM ... to help clarify if the proposed method is better in a statically significant way ... details of the evaluation procedure ...

A4: Following your suggestion, we now report IQM with 80% CI in the response PDF. For more seeds in SMAC v2, we are currently running those, but the results are not fully ready due to the limited time of the rebuttal period. Regarding the evaluation procedure, we follow the standard configurations provided by the corresponding codebase, listed below for your reference:

Q5: An analysis of the computation cost (speed and memory) during training time might be relevant...

A5:

• Cost results: We follow your suggestion and summarize the training and execution costs of all methods in Table 2,9,10 in global rebuttal text. Overall, Kaleidoscope enjoys a higher performance with lower execution complexity but does introduce some extra training overhead. The training overhead is comparable with some baselines such as SePS.

• Discussion: Kaleidoscope sacrifices training computational cost for the following benefits:

  • During training, Kaleidoscope achieves better policy diversity while maintaining high sample efficiency.
  • During execution, Kaleidoscope enjoys better performance with lower overhead.

  We believe this tradeoff is worthwhile because:

  • In a learning task, training is conducted once. After deployment, execution performance and efficiency are what matter.
  • In an RL setting, sample efficiency is more crucial than computational efficiency because online data in the real world is usually difficult to collect.
  • Training computational cost can potentially be reduced by employing strategies such as mixed precision training.

We will incorporate these results and discussions into our manuscript.

Thank you for your positive review. Regarding your questions and suggestions, we would like to provide clarifications below. If you have any follow-up questions or comments, please let us know, and we will be happy to discuss further.

Q1: The current limitations of the work are presented in an Appendix. It would be good to ... as part of the main text. For example, it would be interesting if the authors provide a guideline for how their method would scale with the number of agents.

A1: We appreciate the suggestion and will elaborate on how to protentially scale Kaleidoscope to more agents. Instead of assigning one unique mask to each agent, we can consider clustering $N$ agents into $K$ ($K \lt N$) groups and train $K$ masks with Kaleidoscope. This would reduce computational costs and achieve a better trade-off between sample efficiency and diversity. Within the same group, agents share all parameters, while agents from different groups share only partial parameters. Techniques for clustering agents based on experience, as proposed in [1], could be useful.

Q2: While the presentation of the material is clear, it is repetitive in places. It would be nice to make the text more concise and provide some more interpretation of the results.

A2: We appreciate the suggestion and provide further visualization of agent trajectories in Fig. 7 in the response PDF. We visualize the pairwise mask differences among agents and the agent trajectories at different training stages. As training progresses, the test return increases and diversity loss decreases, indicating better performance and greater diversity among agent policies. Correspondingly, mask differences among agents increase, and the agent trajectory distribution becomes more diverse. This reflects:

• The masks become more diverse as the training proceeds.
• The more diverse masks results in more diverse policies (reflected by the trajectories).
• The diverse policies lead to better performance.

Q3: More self-contained Figure captions and typos.

A3: Thank you for catching those, and we will revise them accordingly.

Thank you for your constructive feedback. Here are our clarifications. If you have any follow-up questions or comments, please let us know and we will be happy to have further discussions.

Q1: The proposed method is similar to SNP...

A1: As listed in Table 1, our proposed method differs from SNP in the following key aspects:

• Kaleidoscope learns masks that indicate which parameters to share during training, whereas SNP relies on fixed mask initialization. We then further design diversity regularization and resetting mechanism to facilitate the mask learning with the objective to boost the task performance.
• Kaleidoscope applies masks to weights, while SNP applies masks to neurons (activation outputs).

These design choices give Kaleidoscope greater flexibility (elaborated below) in parameter sharing and result in better performance (supported by Fig. 4-5).

• The learnable nature of masks provides adaptive PS during training.
• The weight-level sharing allows for partial sharing that is more precise (supported by Fig. 9 in the response PDF).

Q2: The learned masks can not be explained. The proposed method can be regarded as enlarging the network into a Mix-of-Expert network with a dense router, ... does not leverage the relationships among agents.

A2: In Fig. 7 of the original manuscript, we visualize mask differences among agents throughout the training process, reflecting the pairwise similarity between agent policies. Initially, agents share all parameters and gradually learn diverse policies through different masks. During the rebuttal period, we further visualized agent trajectories as training proceeds, with results provided in Fig. 7 of the response PDF.

Q3: Pseudocode is lacking...

A3: We appreciate the suggestion and will revise the manuscript accordingly. Kaleidoscope generally requires an extra resetting step, as shown in Fig. 3 of the original manuscript. It follows the same end-to-end training procedure as the base algorithms (MATD3 or QMIX), with modified network structures and loss functions.

Q4: The proposed method maintains N masks whose sizes are similar to those of the actor and critic networks, which may result in low training efficiency.

A4: Please see A6.

Q5: ... implementing ... a shared trunk with multiple action heads.

A5: We follow the suggestion and implement this baseline as MultiH during the rebuttal period. The results are reported as purple curves in Fig. 3-4 in the response PDF. Overall, the performance of MultiH lies between FuPS and NoPS. Like other non-adaptive partial parameter sharing algorithms, MultiH finds a fixed middle point between full parameter sharing and no parameter sharing. While this particular middle point may be favored in certain scenarios, its performance across different scenarios is not stable due to the lack of flexibility. This observation motivates us to develop an adaptive partial parameter sharing algorithm such as Kaleidoscope.

Q6: Apart from Table 2, the authors should compare the training costs among those algorithms.

A6:

• Cost results: We summarize the training and execution costs of all methods in Table 2,9,10 in global rebuttal text. Overall, Kaleidoscope enjoys a higher performance with lower execution complexity but does introduce some extra training overhead. The training overhead is comparable with some baselines such as SePS.

• Discussion: Kaleidoscope sacrifices training computational cost for the following benefits:

  • During training, Kaleidoscope achieves better policy diversity while maintaining high sample efficiency.
  • During execution, Kaleidoscope enjoys better performance with lower overhead.

  We believe this tradeoff is worthwhile because:

  • In a learning task, training is conducted once. After deployment, execution performance and efficiency are what matter.
  • In an RL setting, sample efficiency is more crucial than computational efficiency because online data in the real world is usually difficult to collect.
  • Training computational cost can potentially be reduced by employing strategies such as mixed precision training.

Q7: The ablation study is incomplete. The main new idea ... share parameters using learnable masks...

A7: We appreciate the suggestion and have added an ablation with fixed masks rather than learnable masks in Fig. 6 of the response PDF. This new ablation underperforms Kaleidoscope with the same sparsity level. It showcases the need to use learnable masks to dynamically adjust the parameter sharing throughout the training (as supported by Fig. 7 in the original manuscript).

Q8: What is the purpose of citing LLM works in lines 35 and 36?

A8: To support the current trend of scaling up model sizes and motivate the need for an effective partial parameter-sharing technique. Although our method focuses on MARL, we believe parameter sharing is potentially a topic of interest in other multi-agent systems. We will polish the wording to make this clearer.

Q9: Are there any other works discussing parameter sharing apart from references [12,13,14] in this paper?

A9: The references in our manuscript are good representatives of recent progress on parameter sharing in MARL. However, as reviewer BN4p suggested, we will add another concurrent work to the related work section: AdaPS [2], which combines SNP and SePS by proposing a cluster-based partial parameter-sharing scheme.

Q10: How to distinguish Neurons and Weights in Table 1?

A10: Neurons refer to the output of the activation function, while weights refer to the network parameters. In a fully-connected layer, removing a neuron is equivalent to removing all the weights associated with that neuron.

Q11: ... are algorithms using similar sizes of networks?...

A11: Yes, we use the same network structure across different algorithms, with details listed in Appendix A.1.3.

Thank you for your positive review. Regarding you questions and suggestions, we would like to provide clarifications and additional results below. If you have any follow-up questions or comments, please let us know and we will be happy to have further discussions.

Q1: The paper primarily focuses on environments with a relatively small number of agents. The scalability of Kaleidoscope to environments with tens or hundreds of agents. In that case it is not clear if a large number of masks could have a negative performance compared with the parameter sharing baseline.

A1:

• The benchmark environments we use (MPE, MaMuJoCo, SMACv2) are widely adopted and challenging in MARL. They cover various characteristics (as specified in Table 7 in Appendix A.2): discrete/continuous action spaces, heterogeneous/homogeneous agent types, and varying numbers of agents. This diversity makes them suitable for testing MARL algorithms.
• We agree that scalability is a long-standing challenge in MARL. We acknowledge that scalability to hundreds of agents is a potential limitation of our approach, which we have discussed in Appendix B.3.
  • Protential solution: To scale Kaleidoscope to hundreds of agents, a possible approach is to cluster $N$ agents into $K$ ($K \lt N$) groups and train $K$ masks with Kaleidoscope. This would reduce computational costs and achieve a better trade-off between sample efficiency and diversity. Within the same group, agents share all parameters, while agents from different groups share only partial parameters. Techniques for clustering agents based on experience, as proposed in [1], could be useful. However, to truly scale current MARL algorithms to hundreds of agents, besides better parameter sharing schemes, we also need consider problems such as severe partiall observation, effective credit assignment and complex coordination among agents, which we will leave for future work.

Q2: How does the computational overhead of Kaleidoscope scale with the number of agents?

A2: In theory, the computational overhead of Kaleidoscope and all baselines scales linearly with the number of agents. However, thanks to PyTorch's parallel implementation framework, the actual training and execution time scales sublinearly, unless memory becomes a bottleneck.

Q3: Impact of Mask Sparsity: Does the sparsity level of the masks induced by STR vary across different tasks and environments? Are there any insights into how the sparsity level affects the policy diversity and overall performance?

A3: Generally, the final sparsity levels of the masks are similar across different tasks. Below are the results, followed by a discussion. Since overall sparsity comparison can be difficult to interpret across environments (affected by network structure, RNN layers, etc.), we report the average sparsity of the fully connected layers where Kaleidoscope is applied

Kaleidoscope tends to induce higher sparsity levels with a larger number of agents. However, overall sparsity across different environments does not vary significantly, generally converging between 0.2 and 0.3. This is because an overly sparse network may lack sufficient representational capacity for complex policies, especially if the original network is small, while networks that are not sparse enough may fail to induce diversity.

Q4: Integration with other MARL Algorithms: The paper demonstrates Kaleidoscope with MATD3 and QMIX. How straightforward would it be to adapt the method to other MARL algorithms such MAPPO? Are there any specific challenges or considerations?

A4: It should be fairly straightforward to adapt Kaleidoscope to other MARL algorithms because it mainly operates at a network level rather than an algorithm level. We used MATD3 and QMIX to illustrate Kaleidoscope because they represent the actor-critic and value-based families, respectively. While incorporating Kaleidoscope into different families of algorithms may vary slightly, implementation within the same family should be similar. To incorporate Kaleidoscope into MAPPO (another actor-critic algorithm), one should follow the Kaleidoscope-MATD3 implementation.

Q5: Real-world Applications: Can you envision any specific real-world applications where Kaleidoscope's ability to promote agent heterogeneity would be particularly beneficial?

A5: Kaleidoscope would be most useful in scenarios where agents need to share some aspects of their policies but not end up with identical policies. For example, two dexterous hands cutting a steak: one hand uses a fork to hold the steak, while the other uses a knife to cut it. Here, the agents should share the grasping skill but not compete for the same utensil and cause clashes. Since it's challenging to know the exact level of heterogeneity needed for a specific multi-agent task beforehand, a flexible framework like Kaleidoscope can be beneficial.