Converting mjx.Model into mujoco.MjModel

[nico-bohlinger]

Hi,
I was wondering if there is a quick way of converting a mjx.Model into a mujoco.MjModel?
There are the functions mjx.get_data and mjx.put_data for the data object but for the model only mjy.put_model exists.

Thank you for any advice in advance.

[WangHanSolo]

Something like this should do the trick for copying the jnp/np arrays.

for field in mjx.Model.fields():
    if field.type in [jax.Array, np.ndarray]:
        value = getattr(mjx_model, field.name)
        setattr(mj_model, field.name, value)

Hope this helps!

[btaba]

Keep in mind that certain fields might need to be reshaped, or may be MJX only fields.

Generally, you should keep around the mjModel used to create the mjx.Model. The correct way to create an mjModel is through the MuJoCo compiler. Feel free to copy fields back from mjx.Model though. We might consider implementing something like put_into_model(mj_model, mjx_model) that does what @WangHanSolo is doing at a high level.

What are your use-cases out of curiosity?