TransfQMix release

[mttga]

No description provided.

[amacrutherford]

looks cool! couple of minor points, also would be interesting to see the SMAX results

[amacrutherford]

need to remove this no?

[mttga]

oh yes true, will do

[amacrutherford]

will each env need an if clause in __init__? maybe worth putting a note if so

[mttga]

yes every env that wants to use trasformers will need to wrap the environment observations in some way to be readable by the transformer. I could also reorganize it to have per-env specific wrappers if that would look better. I had put a note on the readme but now I see I deleted by mistake

[amacrutherford]

why not just import from flax?

[mttga]

the multihead attention from the fastattention script is significantly faster and more stable than the default flax one, which is beneficial in RL.

The script was taken in https://github.com/google-research/google-research/blob/master/performer/fast_attention/jax/fast_attention.py

Notice that the use of fast attention is optional:

https://github.com/FLAIROx/JaxMARL/blob/4d78674dc5195683f3cd6e7e9d6799ddd586a714/baselines/QLearning/transf_qmix.py#L51C9-L72C14

[amacrutherford]

typo?

[mttga]

yes 🙃 thanks

[mttga]

Here are the results average across 4 seeds:

• test won: https://api.wandb.ai/links/mttga/2beham1m
• test returns: https://api.wandb.ai/links/mttga/6yszv25o

Results are better than QMix for most of the maps with except of 3s5z and 3s5z_vs_3s_6z. But the main advantage of transformers is the potential transferability of the agent parameters and the learned qmix function between scenarios.

[amacrutherford]

why we removing the hanabi option?

[mttga]

That was a preliminary way to create a global state vector for Hanabi. I realized that including the players' hands wasn't adding any new information, since those hands are already represented within the concatenated agent observations.

[amacrutherford]

cools !

[amacrutherford]

Here are the results average across 4 seeds:

* test won: https://api.wandb.ai/links/mttga/2beham1m

* test returns: https://api.wandb.ai/links/mttga/6yszv25o

Results are better than QMix for most of the maps with except of 3s5z and 3s5z_vs_3s_6z. But the main advantage of transformers is the potential transferability of the agent parameters and the learned qmix function between scenarios.

awesome!