[Feature Request] Mujoco Playground Integration (with multiagent support)

[itwasabhi]

Motivation

Mujoco Playground [1] is the latest comprehensive, mujoco-based RL control suite offered by deepmind.

Currently torchrl offers the BraxEnv, which offers the standard single-agent version of the original deepmind control suite. Brax authors have reqeusted users to transition to mujoco_playground for environments and only use brax for training algorithms.

Separately, JaxMARL [2] offers a multi-agent version decomoposition of the brax environments as proposed in the FACMAC paper [3]. This continues to be a very useful application domain for cooperative MARL.

[1] https://playground.mujoco.org/
[2] https://github.com/FLAIROx/JaxMARL
[3] https://arxiv.org/abs/2003.06709

Solution

Create a single Mujoco-Playground environment wrapper that would additionally offer multi-agent decomposition.

Upon completing this,

• We can deprecate the BraxEnv
• We can use the multi-agent decomposition in marl examples / benchmarks

Checklist

• I have checked that there is no similar issue in the repo (required)

[itwasabhi]

@vmoens I'm hoping to work on this this week, please tell me if you have any thoughts or advice!