  Recent advances in Reinforcement Learning (RL) have led to many exciting
applications. These advancements have been driven by improvements in both
algorithms and engineering, which have resulted in faster training of RL
agents. We present marl-jax, a multi-agent reinforcement learning software
package for training and evaluating social generalization of the agents. The
package is designed for training a population of agents in multi-agent
environments and evaluating their ability to generalize to diverse background
agents. It is built on top of DeepMind's JAX ecosystem~\cite{deepmind2020jax}
and leverages the RL ecosystem developed by DeepMind. Our framework marl-jax is
capable of working in cooperative and competitive, simultaneous-acting
environments with multiple agents. The package offers an intuitive and
user-friendly command-line interface for training a population and evaluating
its generalization capabilities. In conclusion, marl-jax provides a valuable
resource for researchers interested in exploring social generalization in the
context of MARL. The open-source code for marl-jax is available at:
\href{https://github.com/kinalmehta/marl-jax}{https://github.com/kinalmehta/marl-jax}
