Migrate MABrax environments from Brax to MJX

[rechtevan]

Summary

Brax is no longer actively maintained and the project recommends migrating to MJX (MuJoCo XLA) for JAX-based physics simulation. This issue tracks the migration of MABrax environments from Brax to MJX.

Background

When running the test suite, we see deprecation warnings from Brax:

Brax System, pipelines and environments are not actively being maintained.
Please see MJX for a well maintained JAX-based physics engine.

These warnings originate from brax.io.mjcf and affect all MABrax environments (Ant, HalfCheetah, Hopper, Humanoid, Walker2d).

Affected Files

• jaxmarl/environments/mabrax/__init__.py
• jaxmarl/environments/mabrax/mabrax_env.py
• jaxmarl/environments/mabrax/mappings.py

Migration Path

MJX (MuJoCo XLA) is the official successor to Brax:

• Actively maintained by Google DeepMind
• Native JAX/XLA implementation
• Better performance and compatibility
• Official MuJoCo integration

Resources

• MJX Documentation
• MJX GitHub
• MuJoCo Playground (MJX Examples)

Tasks

• Research MJX API differences from Brax
• Prototype one MABrax environment (e.g., Ant) on MJX
• Validate agent observation/action mappings
• Update all MABrax environments
• Update MABrax baseline algorithms
• Add migration documentation
• Consider deprecating Brax-based MABrax

Complexity

Medium-High - Requires rewriting MABrax environment wrappers and testing all baselines.

Related

• Issue td_lambda loss in q_learning algs #34: Deprecation warnings analysis
• Analysis document: .local/deprecation-warnings-analysis.md