search

google / brax

Massively parallel rigidbody physics simulation on accelerator hardware.
Apache License 2.0
2.38k stars 257 forks source link

Brax's Simulator Engine Swap #505

Closed scott-yj-yang closed 2 months ago

scott-yj-yang commented 4 months ago

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 commented 2 months ago

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.