marl-jax

JAX library for MARL research

Demo Video Paper Link

Implemented Algorithms

• Independent-IMPALA for multi-agent environments
• OPRE

Environments supported

• Meltingpot
• Overcooked
• Sequential Social Dilemma

Other Features

• Distributed training (IMPALA style architecture)
  • Dynamically distribute load of multiple agents across available GPUs
  • Run multiple environment instances, one per CPU core for experience collection
• Wandb and Tensorboard logging
• PopArt normalization

Help

• Installation Instructions
• Environment Details

Implementation References

• Deepmind's Acme

Citation

If you use this code in your project, please cite the following paper:

@article{mehta2023marljax,
      title={marl-jax: Multi-agent Reinforcement Leaning framework for Social Generalization}, 
      author={Kinal Mehta and Anuj Mahajan and Pawan Kumar},
      year={2023},
      journal={arXiv preprint arXiv:2303.13808},
      url={https://arxiv.org/abs/2303.13808},
}