MuJoCo XLA (MJX) support in Unity

[Norite-dev]

The feature, motivation and pitch

Hi guys,
Is there any chance to have MuJoCo XLA (MJX) support in Unity for GPU-based physics calculation, assuming it will increase the training speed significantly. What do you think? What do you @Balint-H personally think?)

Thanks! Love your product

Alternatives

No response

Additional context

No response

[Balint-H]

Hello, this is already partially possible. The process involves:

1. Train with MJX outside Unity. This should incorporate at least moderate domain randomization for robustness.
2. Confirm the policy works to a satisfactory level in CPU MuJoCo outside Unity.
3. Converting the policy to an ONNX format.
4. Install the ML-Agents package to perform inference in Unity with the ONNX policy
5. Import the scene in Unity, confirming scene settings are well replicated from the MJX/MuJoCo scene (e.g. timestep, solver settings). Recreate the observation and action space in Unity as ML-Agents observations and action components.
6. Test performance of your policy trained with MJX in Unity!

@joanllobera and I have been experimenting with this process, with no full end-to-end solution demonstrated yet. You can check out our WIP brax-to-ML-Agents converter scripts here: https://github.com/Balint-H/modular-agents/tree/fix/deepmimic-training/mjx-to-unity , we welcome feedback and contributions. You can also find a lot of other MuJoCo in Unity with RL tools in that repo.

[Balint-H]

Or did you mean to perform the learning and inference directly in a Unity scene?

[Norite-dev]

Or did you mean to perform the learning and inference directly in a Unity scene?

Yes, exactly

[Norite-dev]

I asked because there is a new physics engine released recently https://genesis-embodied-ai.github.io/ with fantastic performance, party because it runs on a GPU

[Balint-H]

Could you elaborate on what role would Unity play in your expected use case?

You could use Unity to help build a scene (the editor serving as a GUI) then export the MuJoCo scene and use that in MJX, which runs parallelized on the GPU. Or run MJX inference and pipe its state to its Unity equivalent for rendering or interactivity.

There is no support currently for GPU parallelized physics directly in a Unity app. Unity has support for GPU computations via e.g., compute shaders, but I think there is no straightforward way to bridge those features to run JAX code (what MJX is written with). It would need a complete reimplementation of the engine the GPU code (which is strenuous, and the reason why MuJoCo didn't come to GPU prior to JAX).

[Norite-dev]

Got it. Thanks for the clarification