[`MJX`] `mujoco.mj_resetData` equivalent for `MuJuCo-MJX`

[Kallinteris-Andreas]

Hi,

I'm a maintainer of Gymnasium & Gymnasium-Robotics, and I'm trying to use MuJoCo-MJX for "prototyping MJX-based RL environments in Gymnasium".

Here is the model tested: Gymnasium/Ant (though it should be not relevant for this question)

I want to use mj_resetData with MJX data & model.

The tutorial (which is based on the old Brax tutorial from what I can tell) replaces the data with initial values and ctrl to 0 and then steps once (as shown below):

    data = data.replace(qpos=qpos, qvel=qvel, ctrl=jp.zeros(self.sys.nu))
    data = mjx.forward(self.sys, data)

Which is a way to get behavior similar to mujoco.mj_resetData (but with stepping once), is there a way to get the same behavior without stepping (which is what mujoco.mj_resetData does from what I can tell).

Thanks!

[yuvaltassa]

mjx.forward and mujoco.mj_forward do not step the dynamics.

[btaba]

Hi @Kallinteris-Andreas

From looking at mj_resetData, you can call mjx.make_data(model) (sets a bunch of zeros), then you can populate the default values (qpos) as in your example. In any case, all the zeroed out fields will get populated on the next forward call

[Kallinteris-Andreas]

@yuvaltassa is mj_resetData the same as setting data->qpos = model->qpos0 and data->qvel = model->qvel0 and data->ctrl = ZEROS and the running mj_forward(qpos, qvel) once (which is exactly what the code above does),

if so what is the purpose of mj_resetData, if not so how do does it differ with code above?

@btaba I would rather not run mjx.make_data() since I want to avoid unnecessary mallocs, plus it might break the pointers of the renderers routines

Thanks!