Multi-agent Environments with MJX

[willthibault]

Hello,

I have been enjoying using Brax with the MJX backend for my single agent RL problems and am interested in investigating multi-agent problems. I understand some previous work (Braxlines) did a bit of work in this direction, but Brax has grown since that work.

I would appreciate any suggestions on how to implement this and am happy to contribute any developments in this direction.

Thanks!