Is the creation of RL model Limited to JAX?

[coder-rakeshpaul]

Hi,

so I understand that the physics simulation in the mujoco 3.0.0.0 is done on the GPU and is more faster, does that use the mujoco-mjx for this simulation.

From what I understand the MJX uses JAX. Is this use of JAX only for the RL part or is the simulation itself been done on this?

I use Pytorch for the suppose of the creating RL models and it'll be some what time consuming to create the entire model using JAX since I'm not entirely familiar with it.

Or can I use the MJX just for simulation of the Physics and create the model in Pytorch??

[btaba]

Hi @coder-rakeshpaul

There are 2 modes: [1] MuJoCo, [2] MJX whose entry point is putting a MuJoCo MjModel and MjData on device

MJX is a jax clone of MuJoCo with a limited feature set, see https://mujoco.readthedocs.io/en/stable/mjx.html#feature-parity

You can create an environment in MJX and interface with pytorch, although the only example we have is doing this with Brax. Check out https://colab.sandbox.google.com/github/google/brax/blob/main/notebooks/training_torch.ipynb#scrollTo=D3y5o7-oSBm-

Hope that helps for now

[coder-rakeshpaul]

Hi @btaba Thank you for the response.

I have another question.

In case in the absence of a dedicated GPU will mujoco use the CPU for the calculations by default?? Or do we need to modify the code in order to make it compatible with it??

[btaba]

Hi @coder-rakeshpaul , MuJoCo runs on CPU, MJX runs on CPU/GPU/TPU depending on the device you have available. If you'd like to run on CPU, we strongly suggest you use MuJoCo

[coder-rakeshpaul]

Hi @btaba Thank you.