Brax's Simulator Engine Swap

[scott-yj-yang]

We've been impressed with Brax's support for the MJX engine and its efficiency in simulation and computation. However, due to the MJX engine's current development status and some feature limitations, we'd like to temporarily switch back to the CPU-based MuJoCo engine to take advantage of its full feature set.

Our computational pipeline is currently built around Brax and Jax's NN Linen. We understand that Brax is designed as a differentiable physics engine, and we're hoping there's a straightforward way to swap out the underlying simulator from MJX to CPU MuJoCo without significant rework.

Any guidance or suggestions on how to achieve this transition smoothly would be greatly appreciated.

Scott

[btaba]

Hi @scott-yj-yang , we have an internal version where we swap MJ CPU with the MJX backend using JAX's FFI. I'm hoping that can be released in the coming months, but if you'd like to keep using a brax/JAX API with MJ, then FFI is the way to go.