Dramatic performance difference between JaxMARL SMAX vs. Mava SMAX?

[Chulabhaya]

Hi there! I was wondering if you guys had any ideas as to why you're seeing such a big performance difference on SMAX with your implementations vs. the JaxMARL implementations? https://github.com/instadeepai/Mava?tab=readme-ov-file#smax

[RuanJohn]

Hi @Chulabhaya. Thank you for the question.

We are in the process of reworking Mava's Readme, so the results should be updated soon as well.

There are a few potential reasons for the performance differences:

1. There are two model architecture differences between Mava and JaxMARL's implementations.
   a. In Mava we do not use a shared actor-critic network with two separate heads that is optimised by a single optimiser with a shared loss. Instead, we have completely independent actor and critic networks which get optimised with their own losses and optimisers. We did this based on the recommendations from The Surpirising Effectiveness of PPO in Cooperative Multi-Agent Games. As well as the implementation details from Is Independent Learning All You Need in the StarCraft Multi-Agent Challenge?, the results from What Matters In On-Policy Reinforcement Learning? A Large-Scale Empirical Study and the suggestion from The 37 Implementation Details of Proximal Policy Optimization. Although the two latter references are from single-agent literature we still thought the suggestions were sound and well tested.

   b. In JaxMARL it seems like the environment and agent dimensions are flattened into one and then used as the input to the RNN in the case of recurrent algorithms. This implies that the hidden states are shared among agents. In Mava we keep agent observations and hidden states completely separate between agents.

2. Mava is also written in a way to leverages JAX's vmap and pmap function transformations by default. This implies that even though we train for the same number of environment interactions as JaxMARL (10 Million Timesteps in this case) Mava will vmap the entire update (including batched experience generation in the environment and gradient updates) twice over batches of 64 parallel environments each. This leads to smaller effective batch sizes that sync gradients via JAX's pmean operation. This could have an impact on system performance as well.

3. The differences could also stem from hyperparameter optmisation. We have optimised the hyperparameters of Mava's systems as best we could within reasonable computational constraints. We have taken the JaxMARL baselines results directly from a version that was on arXiv at the time when we had just rewritten Mava in JAX and assumed that the authors of JaxMARL would have also tuned their systems. For the performance of our most recent version of Mava we invite you to have a look at Figure 11 in Appendix C in our most recent paper. Here all systems were tuned with a tuning budget of 40 trials per task. All hyperparameters can also be found in the source code uploaded to this website.

Please let us know if you have any further questions, and please keep checking back since we will have a newly overhauled Readme live soon highlighting all our most recent developments.