From d59978066416298d519dc77b572f56a287f06bec Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Fri, 3 Jul 2026 17:51:10 +0200 Subject: [PATCH 1/5] Add reset_data_keyframe function and corresponding tests - passing mujoco.MjModel as an argument --- mujoco_warp/__init__.py | 1 + mujoco_warp/_src/io.py | 184 ++++++++++++++++++++++++++++++++++++ mujoco_warp/_src/io_test.py | 90 ++++++++++++++++++ 3 files changed, 275 insertions(+) diff --git a/mujoco_warp/__init__.py b/mujoco_warp/__init__.py index 81f5970ab..729ccb689 100644 --- a/mujoco_warp/__init__.py +++ b/mujoco_warp/__init__.py @@ -58,6 +58,7 @@ from mujoco_warp._src.io import put_data as put_data from mujoco_warp._src.io import put_model as put_model from mujoco_warp._src.io import reset_data as reset_data +from mujoco_warp._src.io import reset_data_keyframe as reset_data_keyframe from mujoco_warp._src.io import set_const as set_const from mujoco_warp._src.io import set_const_0 as set_const_0 from mujoco_warp._src.io import set_const_fixed as set_const_fixed diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index 466ac26b9..02fa245ec 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -2576,6 +2576,190 @@ def reset_sleep( sleep.update_sleep(m, d) +def reset_data_keyframe( + mjm: mujoco.MjModel, + m: types.Model, + d: types.Data, + key: int, + reset: Optional[wp.array] = None, +): + """Reset data, set fields from specified keyframe; optionally by world. + + Unlike `reset_data`, this function requires the host-side MuJoCo model as an argument + becasue the keyframe data are not tracked by `types.Model`. + + Args: + mjm: The host-side MuJoCo model. + m: The model containing kinematic and dynamic information (device). + d: The data object containing the current state and output arrays (device). + key: The keyframe index to initialize the data with. An IndexError is raised if + key<0 or key>=mjm.nkey. + reset: Per-world bitmask. Reset if True. + """ + reset_data(m, d, reset) + + @wp.kernel(module="unique", enable_backward=False) + def reset_time( + # In: + target_time: float, + reset_in: wp.array[bool], + # Data out: + time_out: wp.array[float], + ): + worldid = wp.tid() + + if wp.static(reset is not None): + if not reset_in[worldid]: + return + + time_out[worldid] = target_time + + @wp.kernel(module="unique", enable_backward=False) + def reset_qpos( + # In: + target_qpos: wp.array[float], + reset_in: wp.array[bool], + # Data out: + qpos_out: wp.array2d[float], + ): + worldid, qid = wp.tid() + + if wp.static(reset is not None): + if not reset_in[worldid]: + return + + qpos_out[worldid, qid] = target_qpos[qid] + + @wp.kernel(module="unique", enable_backward=False) + def reset_qvel( + # In: + target_qvel: wp.array[float], + reset_in: wp.array[bool], + # Data out: + qvel_out: wp.array2d[float], + ): + worldid, vid = wp.tid() + + if wp.static(reset is not None): + if not reset_in[worldid]: + return + + qvel_out[worldid, vid] = target_qvel[vid] + + @wp.kernel(module="unique", enable_backward=False) + def reset_activation( + # In: + target_act: wp.array[float], + reset_in: wp.array[bool], + # Data out: + act_out: wp.array2d[float], + ): + worldid, aid = wp.tid() + + if wp.static(reset is not None): + if not reset_in[worldid]: + return + + act_out[worldid, aid] = target_act[aid] + + @wp.kernel(module="unique", enable_backward=False) + def reset_mocap( + # Model: + body_mocapid: wp.array[int], + # In: + target_mpos: wp.array[wp.vec3], + target_mquat: wp.array[wp.quat], + # In: + reset_in: wp.array[bool], + # Data out: + mocap_pos_out: wp.array2d[wp.vec3], + mocap_quat_out: wp.array2d[wp.quat], + ): + worldid, bodyid = wp.tid() + + if wp.static(reset is not None): + if not reset_in[worldid]: + return + + mocapid = body_mocapid[bodyid] + + if mocapid >= 0: + mocap_pos_out[worldid, mocapid] = target_mpos[mocapid] + mocap_quat_out[worldid, mocapid] = target_mquat[mocapid] + + @wp.kernel(module="unique", enable_backward=False) + def reset_control( + # In: + target_ctrl: wp.array[float], + reset_in: wp.array[bool], + # Data out: + ctrl_out: wp.array2d[float], + ): + worldid, cid = wp.tid() + + if wp.static(reset is not None): + if not reset_in[worldid]: + return + + ctrl_out[worldid, cid] = target_ctrl[cid] + + reset_input = reset or wp.ones(d.nworld, dtype=bool) + + target_time = mjm.key_time[key] + wp.launch( + reset_time, + dim=d.nworld, + inputs=[target_time, reset_input], + outputs=[d.time], + ) + + target_qpos = wp.array(mjm.key_qpos[key], dtype=float) + wp.launch( + reset_qpos, + dim=(d.nworld, m.nq), + inputs=[target_qpos, reset_input], + outputs=[d.qpos], + ) + + target_qvel = wp.array(mjm.key_qvel[key], dtype=float) + wp.launch( + reset_qvel, + dim=(d.nworld, m.nv), + inputs=[target_qvel, reset_input], + outputs=[d.qvel], + ) + + target_act = wp.array(mjm.key_act[key], dtype=float) + wp.launch( + reset_activation, + dim=(d.nworld, m.na), + inputs=[target_act, reset_input], + outputs=[d.act], + ) + + target_mpos = wp.array(mjm.key_mpos[key], dtype=wp.vec3) + target_mquat = wp.array(mjm.key_mquat[key], dtype=wp.quat) + wp.launch( + reset_mocap, + dim=(d.nworld, m.nbody), + inputs=[ + m.body_mocapid, + target_mpos, + target_mquat, + reset_input, + ], + outputs=[d.mocap_pos, d.mocap_quat], + ) + + target_ctrl = wp.array(mjm.key_ctrl[key], dtype=float) + wp.launch( + reset_control, + dim=(d.nworld, m.nu), + inputs=[target_ctrl, reset_input], + outputs=[d.ctrl], + ) + + # kernel_analyzer: off @wp.kernel def _init_subtreemass( diff --git a/mujoco_warp/_src/io_test.py b/mujoco_warp/_src/io_test.py index 1a79df424..0155a64a4 100644 --- a/mujoco_warp/_src/io_test.py +++ b/mujoco_warp/_src/io_test.py @@ -1062,6 +1062,96 @@ def test_reset_data_world(self): _assert_eq(d.qvel.numpy()[0], 1.0, "qvel[0]") _assert_eq(d.qvel.numpy()[1], 2.0, "qvel[1]") + def test_reset_data_keyframe(self): + """Tests that reset_data_keyframe matches mj_resetDataKeyframe.""" + _MJCF = """ + + + + + + + + + + + + + + + + + + """ + reset_datafield = ["time", "qpos", "qvel", "act", "ctrl", "mocap_pos", "mocap_quat"] + key = 0 + + mjm, mjd, m, d = test_data.fixture(xml=_MJCF) + + # corrupt data + for arr in reset_datafield: + attr = getattr(d, arr) + if attr.dtype == float: + attr.fill_(wp.nan) + else: + attr.fill_(-1) + + mujoco.mj_resetDataKeyframe(mjm, mjd, key) + mjwarp.reset_data_keyframe(mjm, m, d, key) + + for arr in reset_datafield: + _assert_eq(getattr(d, arr).numpy()[0], getattr(mjd, arr), arr) + + def test_reset_data_keyframe_world(self): + """Tests per-world reset for reset_data_keyframe.""" + _MJCF = """ + + + + + + + + + + + + """ + key = 0 + mjm = mujoco.MjModel.from_xml_string(_MJCF) + m = mjwarp.put_model(mjm) + d = mjwarp.make_data(mjm, nworld=2) + + # nonzero values + qpos = wp.array(np.array([[1.0], [2.0]]), dtype=float) + + wp.copy(d.qpos, qpos) + + # reset both worlds + mjwarp.reset_data_keyframe(mjm, m, d, key) + + _assert_eq(d.qpos.numpy()[0], 0.5, "qpos[0]") + _assert_eq(d.qpos.numpy()[1], 0.5, "qpos[1]") + + wp.copy(d.qpos, qpos) + + # don't reset second world + reset10 = wp.array(np.array([True, False]), dtype=bool) + mjwarp.reset_data_keyframe(mjm, m, d, key, reset=reset10) + + _assert_eq(d.qpos.numpy()[0], 0.5, "qpos[0]") + _assert_eq(d.qpos.numpy()[1], 2.0, "qpos[1]") + + wp.copy(d.qpos, qpos) + + # don't reset both worlds + reset00 = wp.array(np.array([False, False], dtype=bool)) + mjwarp.reset_data_keyframe(mjm, m, d, key, reset=reset00) + + _assert_eq(d.qpos.numpy()[0], 1.0, "qpos[0]") + _assert_eq(d.qpos.numpy()[1], 2.0, "qpos[1]") + def test_sdf(self): """Tests that an SDF can be loaded.""" mjm, mjd, m, d = test_data.fixture("collision_sdf/cow.xml") From ac546fe76e2989c970e9fb5ee9c5023c55841ba3 Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Fri, 3 Jul 2026 18:34:51 +0200 Subject: [PATCH 2/5] Add keyframe info to Model dataclass; remove MjModel input from reset_data_keyframe --- mujoco_warp/_src/io.py | 70 ++++++++++++++++++------------------- mujoco_warp/_src/io_test.py | 11 +++--- mujoco_warp/_src/types.py | 16 +++++++++ 3 files changed, 55 insertions(+), 42 deletions(-) diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index 02fa245ec..f7b44ea2c 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -2577,7 +2577,6 @@ def reset_sleep( def reset_data_keyframe( - mjm: mujoco.MjModel, m: types.Model, d: types.Data, key: int, @@ -2585,23 +2584,20 @@ def reset_data_keyframe( ): """Reset data, set fields from specified keyframe; optionally by world. - Unlike `reset_data`, this function requires the host-side MuJoCo model as an argument - becasue the keyframe data are not tracked by `types.Model`. - Args: - mjm: The host-side MuJoCo model. m: The model containing kinematic and dynamic information (device). d: The data object containing the current state and output arrays (device). - key: The keyframe index to initialize the data with. An IndexError is raised if - key<0 or key>=mjm.nkey. + key: The keyframe index to initialize the data with. reset: Per-world bitmask. Reset if True. """ reset_data(m, d, reset) @wp.kernel(module="unique", enable_backward=False) def reset_time( + # Model: + key_time: wp.array[float], # In: - target_time: float, + key: int, reset_in: wp.array[bool], # Data out: time_out: wp.array[float], @@ -2612,12 +2608,14 @@ def reset_time( if not reset_in[worldid]: return - time_out[worldid] = target_time + time_out[worldid] = key_time[key] @wp.kernel(module="unique", enable_backward=False) def reset_qpos( + # Model: + key_qpos: wp.array2d[float], # In: - target_qpos: wp.array[float], + key: int, reset_in: wp.array[bool], # Data out: qpos_out: wp.array2d[float], @@ -2628,12 +2626,14 @@ def reset_qpos( if not reset_in[worldid]: return - qpos_out[worldid, qid] = target_qpos[qid] + qpos_out[worldid, qid] = key_qpos[key, qid] @wp.kernel(module="unique", enable_backward=False) def reset_qvel( + # Model: + key_qvel: wp.array2d[float], # In: - target_qvel: wp.array[float], + key: int, reset_in: wp.array[bool], # Data out: qvel_out: wp.array2d[float], @@ -2644,12 +2644,14 @@ def reset_qvel( if not reset_in[worldid]: return - qvel_out[worldid, vid] = target_qvel[vid] + qvel_out[worldid, vid] = key_qvel[key, vid] @wp.kernel(module="unique", enable_backward=False) def reset_activation( + # Model: + key_act: wp.array2d[float], # In: - target_act: wp.array[float], + key: int, reset_in: wp.array[bool], # Data out: act_out: wp.array2d[float], @@ -2660,16 +2662,16 @@ def reset_activation( if not reset_in[worldid]: return - act_out[worldid, aid] = target_act[aid] + act_out[worldid, aid] = key_act[key, aid] @wp.kernel(module="unique", enable_backward=False) def reset_mocap( # Model: body_mocapid: wp.array[int], + key_mpos: wp.array2d[wp.vec3], + key_mquat: wp.array2d[wp.quat], # In: - target_mpos: wp.array[wp.vec3], - target_mquat: wp.array[wp.quat], - # In: + key: int, reset_in: wp.array[bool], # Data out: mocap_pos_out: wp.array2d[wp.vec3], @@ -2684,13 +2686,15 @@ def reset_mocap( mocapid = body_mocapid[bodyid] if mocapid >= 0: - mocap_pos_out[worldid, mocapid] = target_mpos[mocapid] - mocap_quat_out[worldid, mocapid] = target_mquat[mocapid] + mocap_pos_out[worldid, mocapid] = key_mpos[key, mocapid] + mocap_quat_out[worldid, mocapid] = key_mquat[key, mocapid] @wp.kernel(module="unique", enable_backward=False) def reset_control( + # Model: + key_ctrl: wp.array2d[float], # In: - target_ctrl: wp.array[float], + key: int, reset_in: wp.array[bool], # Data out: ctrl_out: wp.array2d[float], @@ -2701,61 +2705,55 @@ def reset_control( if not reset_in[worldid]: return - ctrl_out[worldid, cid] = target_ctrl[cid] + ctrl_out[worldid, cid] = key_ctrl[key, cid] reset_input = reset or wp.ones(d.nworld, dtype=bool) - target_time = mjm.key_time[key] wp.launch( reset_time, dim=d.nworld, - inputs=[target_time, reset_input], + inputs=[m.key_time, key, reset_input], outputs=[d.time], ) - target_qpos = wp.array(mjm.key_qpos[key], dtype=float) wp.launch( reset_qpos, dim=(d.nworld, m.nq), - inputs=[target_qpos, reset_input], + inputs=[m.key_qpos, key, reset_input], outputs=[d.qpos], ) - target_qvel = wp.array(mjm.key_qvel[key], dtype=float) wp.launch( reset_qvel, dim=(d.nworld, m.nv), - inputs=[target_qvel, reset_input], + inputs=[m.key_qvel, key, reset_input], outputs=[d.qvel], ) - target_act = wp.array(mjm.key_act[key], dtype=float) wp.launch( reset_activation, dim=(d.nworld, m.na), - inputs=[target_act, reset_input], + inputs=[m.key_act, key, reset_input], outputs=[d.act], ) - target_mpos = wp.array(mjm.key_mpos[key], dtype=wp.vec3) - target_mquat = wp.array(mjm.key_mquat[key], dtype=wp.quat) wp.launch( reset_mocap, dim=(d.nworld, m.nbody), inputs=[ m.body_mocapid, - target_mpos, - target_mquat, + m.key_mpos, + m.key_mquat, + key, reset_input, ], outputs=[d.mocap_pos, d.mocap_quat], ) - target_ctrl = wp.array(mjm.key_ctrl[key], dtype=float) wp.launch( reset_control, dim=(d.nworld, m.nu), - inputs=[target_ctrl, reset_input], + inputs=[m.key_ctrl, key, reset_input], outputs=[d.ctrl], ) diff --git a/mujoco_warp/_src/io_test.py b/mujoco_warp/_src/io_test.py index 0155a64a4..3266ae257 100644 --- a/mujoco_warp/_src/io_test.py +++ b/mujoco_warp/_src/io_test.py @@ -1087,7 +1087,7 @@ def test_reset_data_keyframe(self): reset_datafield = ["time", "qpos", "qvel", "act", "ctrl", "mocap_pos", "mocap_quat"] key = 0 - mjm, mjd, m, d = test_data.fixture(xml=_MJCF) + mjm, mjd, m, d = test_data.fixture(xml=_MJCF, keyframe=key) # corrupt data for arr in reset_datafield: @@ -1097,8 +1097,7 @@ def test_reset_data_keyframe(self): else: attr.fill_(-1) - mujoco.mj_resetDataKeyframe(mjm, mjd, key) - mjwarp.reset_data_keyframe(mjm, m, d, key) + mjwarp.reset_data_keyframe(m, d, key) for arr in reset_datafield: _assert_eq(getattr(d, arr).numpy()[0], getattr(mjd, arr), arr) @@ -1129,7 +1128,7 @@ def test_reset_data_keyframe_world(self): wp.copy(d.qpos, qpos) # reset both worlds - mjwarp.reset_data_keyframe(mjm, m, d, key) + mjwarp.reset_data_keyframe(m, d, key) _assert_eq(d.qpos.numpy()[0], 0.5, "qpos[0]") _assert_eq(d.qpos.numpy()[1], 0.5, "qpos[1]") @@ -1138,7 +1137,7 @@ def test_reset_data_keyframe_world(self): # don't reset second world reset10 = wp.array(np.array([True, False]), dtype=bool) - mjwarp.reset_data_keyframe(mjm, m, d, key, reset=reset10) + mjwarp.reset_data_keyframe(m, d, key, reset=reset10) _assert_eq(d.qpos.numpy()[0], 0.5, "qpos[0]") _assert_eq(d.qpos.numpy()[1], 2.0, "qpos[1]") @@ -1147,7 +1146,7 @@ def test_reset_data_keyframe_world(self): # don't reset both worlds reset00 = wp.array(np.array([False, False], dtype=bool)) - mjwarp.reset_data_keyframe(mjm, m, d, key, reset=reset00) + mjwarp.reset_data_keyframe(m, d, key, reset=reset00) _assert_eq(d.qpos.numpy()[0], 1.0, "qpos[0]") _assert_eq(d.qpos.numpy()[1], 2.0, "qpos[1]") diff --git a/mujoco_warp/_src/types.py b/mujoco_warp/_src/types.py index 600489128..5e79244fd 100644 --- a/mujoco_warp/_src/types.py +++ b/mujoco_warp/_src/types.py @@ -1005,6 +1005,7 @@ class Model: nJten: number of non-zeros in sparse tendon Jacobian nwrap: number of wrap objects in all tendon paths nsensor: number of sensors + nkey: number of keyframes nmocap: number of mocap bodies nplugin: number of plugin instances nJmom: number of non-zeros in actuator_moment @@ -1302,6 +1303,13 @@ class Model: sensor_interval: sensor interval and phase (nsensor, 2) plugin: globally registered plugin slot number (nplugin,) plugin_attr: config attributes of geom plugin (nplugin, _NPLUGINATTR) + key_time: keyframe time (nkey,) + key_qpos: keyframe qpos (nkey, nq) + key_qvel: keyframe qvel (nkey, nv) + key_act: keyframe act (nkey, na) + key_mpos: keyframe mocap pos (nkey, nmocap, 3) + key_mquat: keyframe mocap quat (nkey, nmocap, 4) + key_ctrl: keyframe ctrl (nkey, nu) M_rownnz: number of non-zeros in each row of M (nv,) M_rowadr: index of each row in M (nv,) M_colind: column indices of non-zeros in M (nC,) @@ -1490,6 +1498,7 @@ class Model: nJten: int nwrap: int nsensor: int + nkey: int nmocap: int nplugin: int nJmom: int @@ -1786,6 +1795,13 @@ class Model: sensor_interval: array("nsensor", wp.vec2) plugin: array("nplugin", int) plugin_attr: array("nplugin", vec_pluginattr) + key_time: array("nkey", float) + key_qpos: array("nkey", "nq", float) + key_qvel: array("nkey", "nv", float) + key_act: array("nkey", "na", float) + key_mpos: array("nkey", "nmocap", wp.vec3) + key_mquat: array("nkey", "nmocap", wp.quat) + key_ctrl: array("nkey", "nu", float) M_rownnz: array("nv", int) M_rowadr: array("nv", int) M_colind: array("nC", int) From 75c483a6d82ad4a2a9df175d51a8da4c70c79de5 Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Fri, 3 Jul 2026 19:20:18 +0200 Subject: [PATCH 3/5] add out-of-bound validation for keyframe index --- mujoco_warp/_src/io.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index f7b44ea2c..befdc8cac 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -2589,7 +2589,13 @@ def reset_data_keyframe( d: The data object containing the current state and output arrays (device). key: The keyframe index to initialize the data with. reset: Per-world bitmask. Reset if True. + + Raises: + ValueError: If key<0 or key>=m.nkey. """ + if key < 0 or key >= m.nkey: + raise ValueError(f"key ({key}) must be in [0, {m.nkey}).") + reset_data(m, d, reset) @wp.kernel(module="unique", enable_backward=False) From 6cb9e2d6a30df63e8a34cfb2e779d55f9caead8e Mon Sep 17 00:00:00 2001 From: Sibo Wang-Chen Date: Mon, 13 Jul 2026 15:52:26 +0200 Subject: [PATCH 4/5] enable resetting each world to a different keyframe --- mujoco_warp/_src/io.py | 115 ++++++++++++++++++++++++------------ mujoco_warp/_src/io_test.py | 35 +++++++++++ 2 files changed, 111 insertions(+), 39 deletions(-) diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index befdc8cac..318384993 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -2579,7 +2579,7 @@ def reset_sleep( def reset_data_keyframe( m: types.Model, d: types.Data, - key: int, + key: int | wp.array, reset: Optional[wp.array] = None, ): """Reset data, set fields from specified keyframe; optionally by world. @@ -2587,33 +2587,72 @@ def reset_data_keyframe( Args: m: The model containing kinematic and dynamic information (device). d: The data object containing the current state and output arrays (device). - key: The keyframe index to initialize the data with. + key: The keyframe index to initialize the data with. Either an int, which + sets the same keyframe for every world, or a wp.array(dtype=int) of size + nworld, which sets a keyframe per world. In the latter case, worlds with a + keyframe index < 0 or >= m.nkey are not reset. reset: Per-world bitmask. Reset if True. Raises: - ValueError: If key<0 or key>=m.nkey. + ValueError: If key is an int and key<0 or key>=m.nkey. """ - if key < 0 or key >= m.nkey: - raise ValueError(f"key ({key}) must be in [0, {m.nkey}).") + # Resolve reset world mask + if reset is None: + reset_input = wp.ones(d.nworld, dtype=bool) + else: + reset_input = reset - reset_data(m, d, reset) + # Resolve target keyframe index + if isinstance(key, wp.array): + key_input = key + else: + if key < 0 or key >= m.nkey: + raise ValueError(f"key ({key}) must be in [0, {m.nkey}).") + key_input = wp.full(d.nworld, key, dtype=int) + # Update reset mask: if key is out of bounds, mark reset mask as False for that world + @wp.kernel(module="unique", enable_backward=False) + def update_mask_for_invalid_keys( + # Model: + nkey: int, + # In: + key_in: wp.array[int], + reset_in: wp.array[bool], + # Out: + mask_out: wp.array[bool], + ): + worldid = wp.tid() + key = key_in[worldid] + mask_out[worldid] = reset_in[worldid] and key >= 0 and key < nkey + + reset_input_with_invalid_keys_masked = wp.empty(d.nworld, dtype=bool) + wp.launch( + update_mask_for_invalid_keys, + dim=d.nworld, + inputs=[m.nkey, key_input, reset_input], + outputs=[reset_input_with_invalid_keys_masked], + ) + + # Call normal reset using updated mask + reset_data(m, d, reset_input_with_invalid_keys_masked) + + # Set time, qpos, qvel, act, mocap_pos, mocap_quat, ctrl from keyframes @wp.kernel(module="unique", enable_backward=False) def reset_time( # Model: key_time: wp.array[float], # In: - key: int, + key_in: wp.array[int], reset_in: wp.array[bool], # Data out: time_out: wp.array[float], ): worldid = wp.tid() - if wp.static(reset is not None): - if not reset_in[worldid]: - return + if not reset_in[worldid]: + return + key = key_in[worldid] time_out[worldid] = key_time[key] @wp.kernel(module="unique", enable_backward=False) @@ -2621,17 +2660,17 @@ def reset_qpos( # Model: key_qpos: wp.array2d[float], # In: - key: int, + key_in: wp.array[int], reset_in: wp.array[bool], # Data out: qpos_out: wp.array2d[float], ): worldid, qid = wp.tid() - if wp.static(reset is not None): - if not reset_in[worldid]: - return + if not reset_in[worldid]: + return + key = key_in[worldid] qpos_out[worldid, qid] = key_qpos[key, qid] @wp.kernel(module="unique", enable_backward=False) @@ -2639,17 +2678,17 @@ def reset_qvel( # Model: key_qvel: wp.array2d[float], # In: - key: int, + key_in: wp.array[int], reset_in: wp.array[bool], # Data out: qvel_out: wp.array2d[float], ): worldid, vid = wp.tid() - if wp.static(reset is not None): - if not reset_in[worldid]: - return + if not reset_in[worldid]: + return + key = key_in[worldid] qvel_out[worldid, vid] = key_qvel[key, vid] @wp.kernel(module="unique", enable_backward=False) @@ -2657,17 +2696,17 @@ def reset_activation( # Model: key_act: wp.array2d[float], # In: - key: int, + key_in: wp.array[int], reset_in: wp.array[bool], # Data out: act_out: wp.array2d[float], ): worldid, aid = wp.tid() - if wp.static(reset is not None): - if not reset_in[worldid]: - return + if not reset_in[worldid]: + return + key = key_in[worldid] act_out[worldid, aid] = key_act[key, aid] @wp.kernel(module="unique", enable_backward=False) @@ -2677,7 +2716,7 @@ def reset_mocap( key_mpos: wp.array2d[wp.vec3], key_mquat: wp.array2d[wp.quat], # In: - key: int, + key_in: wp.array[int], reset_in: wp.array[bool], # Data out: mocap_pos_out: wp.array2d[wp.vec3], @@ -2685,13 +2724,13 @@ def reset_mocap( ): worldid, bodyid = wp.tid() - if wp.static(reset is not None): - if not reset_in[worldid]: - return + if not reset_in[worldid]: + return mocapid = body_mocapid[bodyid] if mocapid >= 0: + key = key_in[worldid] mocap_pos_out[worldid, mocapid] = key_mpos[key, mocapid] mocap_quat_out[worldid, mocapid] = key_mquat[key, mocapid] @@ -2700,46 +2739,44 @@ def reset_control( # Model: key_ctrl: wp.array2d[float], # In: - key: int, + key_in: wp.array[int], reset_in: wp.array[bool], # Data out: ctrl_out: wp.array2d[float], ): worldid, cid = wp.tid() - if wp.static(reset is not None): - if not reset_in[worldid]: - return + if not reset_in[worldid]: + return + key = key_in[worldid] ctrl_out[worldid, cid] = key_ctrl[key, cid] - reset_input = reset or wp.ones(d.nworld, dtype=bool) - wp.launch( reset_time, dim=d.nworld, - inputs=[m.key_time, key, reset_input], + inputs=[m.key_time, key_input, reset_input_with_invalid_keys_masked], outputs=[d.time], ) wp.launch( reset_qpos, dim=(d.nworld, m.nq), - inputs=[m.key_qpos, key, reset_input], + inputs=[m.key_qpos, key_input, reset_input_with_invalid_keys_masked], outputs=[d.qpos], ) wp.launch( reset_qvel, dim=(d.nworld, m.nv), - inputs=[m.key_qvel, key, reset_input], + inputs=[m.key_qvel, key_input, reset_input_with_invalid_keys_masked], outputs=[d.qvel], ) wp.launch( reset_activation, dim=(d.nworld, m.na), - inputs=[m.key_act, key, reset_input], + inputs=[m.key_act, key_input, reset_input_with_invalid_keys_masked], outputs=[d.act], ) @@ -2750,8 +2787,8 @@ def reset_control( m.body_mocapid, m.key_mpos, m.key_mquat, - key, - reset_input, + key_input, + reset_input_with_invalid_keys_masked, ], outputs=[d.mocap_pos, d.mocap_quat], ) @@ -2759,7 +2796,7 @@ def reset_control( wp.launch( reset_control, dim=(d.nworld, m.nu), - inputs=[m.key_ctrl, key, reset_input], + inputs=[m.key_ctrl, key_input, reset_input_with_invalid_keys_masked], outputs=[d.ctrl], ) diff --git a/mujoco_warp/_src/io_test.py b/mujoco_warp/_src/io_test.py index 3266ae257..4ca5db496 100644 --- a/mujoco_warp/_src/io_test.py +++ b/mujoco_warp/_src/io_test.py @@ -1151,6 +1151,41 @@ def test_reset_data_keyframe_world(self): _assert_eq(d.qpos.numpy()[0], 1.0, "qpos[0]") _assert_eq(d.qpos.numpy()[1], 2.0, "qpos[1]") + def test_reset_data_keyframe_per_world(self): + """Tests reset_data_keyframe with a per-world keyframe array.""" + _MJCF = """ + + + + + + + + + + + + + """ + mjm = mujoco.MjModel.from_xml_string(_MJCF) + m = mjwarp.put_model(mjm) + d = mjwarp.make_data(mjm, nworld=4) + + qpos = wp.array(np.array([[1.0], [2.0], [3.0], [4.0]]), dtype=float) + wp.copy(d.qpos, qpos) + + # world 0 -> key 0 + # world 1 -> key 1 + # world 2 -> key < 0 (reset skipped) + # world 3 -> key >= nkey (reset skipped) + key = wp.array(np.array([0, 1, -1, 2]), dtype=int) + mjwarp.reset_data_keyframe(m, d, key) + + _assert_eq(d.qpos.numpy()[0], 0.5, "qpos[0]") + _assert_eq(d.qpos.numpy()[1], 0.6, "qpos[1]") + _assert_eq(d.qpos.numpy()[2], 3.0, "qpos[2]") + _assert_eq(d.qpos.numpy()[3], 4.0, "qpos[3]") + def test_sdf(self): """Tests that an SDF can be loaded.""" mjm, mjd, m, d = test_data.fixture("collision_sdf/cow.xml") From d55d0b6436369090f3df385df7e11956417a9b19 Mon Sep 17 00:00:00 2001 From: Sibo Wang-Chen Date: Mon, 13 Jul 2026 16:28:16 +0200 Subject: [PATCH 5/5] add type and shape checks for reset mask and target keys in reset_data and reset_data_keyframe --- mujoco_warp/_src/io.py | 19 ++++++++++++-- mujoco_warp/_src/io_test.py | 50 +++++++++++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+), 2 deletions(-) diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index 318384993..00b0d70e5 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -2241,6 +2241,9 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None): m: The model containing kinematic and dynamic information (device). d: The data object containing the current state and output arrays (device). reset: Per-world bitmask. Reset if True. + + Raises: + ValueError: If reset is specified but its shape is not (d.nworld,). """ sleep_enabled = bool(m.opt.enableflags & types.EnableBit.SLEEP) @@ -2470,7 +2473,14 @@ def reset_sleep( if elemid < nv: dof_awake_ind_out[worldid, elemid] = elemid - reset_input = reset or wp.ones(d.nworld, dtype=bool) + if reset is None: + reset_input = wp.ones(d.nworld, dtype=bool) + elif isinstance(reset, wp.array): + if reset.shape != (d.nworld,): + raise ValueError(f"reset array must have shape ({d.nworld},), got {reset.shape}.") + reset_input = reset + else: + raise ValueError(f"reset must be None or a wp.array, got {type(reset)}.") wp.launch(reset_xfrc_applied, dim=(d.nworld, m.nbody, 6), inputs=[reset_input], outputs=[d.xfrc_applied]) wp.launch( @@ -2595,6 +2605,7 @@ def reset_data_keyframe( Raises: ValueError: If key is an int and key<0 or key>=m.nkey. + ValueError: If key is a wp.array but its shape is not (d.nworld,). """ # Resolve reset world mask if reset is None: @@ -2604,11 +2615,15 @@ def reset_data_keyframe( # Resolve target keyframe index if isinstance(key, wp.array): + if key.shape != (d.nworld,): + raise ValueError(f"key array must have shape ({d.nworld},), got {key.shape}.") key_input = key - else: + elif isinstance(key, int): if key < 0 or key >= m.nkey: raise ValueError(f"key ({key}) must be in [0, {m.nkey}).") key_input = wp.full(d.nworld, key, dtype=int) + else: + raise ValueError(f"key must be an int or a wp.array, got {type(key)}.") # Update reset mask: if key is out of bounds, mark reset mask as False for that world @wp.kernel(module="unique", enable_backward=False) diff --git a/mujoco_warp/_src/io_test.py b/mujoco_warp/_src/io_test.py index 4ca5db496..09adcef65 100644 --- a/mujoco_warp/_src/io_test.py +++ b/mujoco_warp/_src/io_test.py @@ -1062,6 +1062,28 @@ def test_reset_data_world(self): _assert_eq(d.qvel.numpy()[0], 1.0, "qvel[0]") _assert_eq(d.qvel.numpy()[1], 2.0, "qvel[1]") + def test_reset_data_reset_invalid(self): + """Tests that reset_data validates the reset argument.""" + _MJCF = """ + + + + + + + + + """ + mjm = mujoco.MjModel.from_xml_string(_MJCF) + m = mjwarp.put_model(mjm) + d = mjwarp.make_data(mjm, nworld=2) + + with self.assertRaisesRegex(ValueError, "reset array must have shape"): + mjwarp.reset_data(m, d, reset=wp.array(np.array([True, False, True]), dtype=bool)) + + with self.assertRaisesRegex(ValueError, "reset must be None or a wp.array"): + mjwarp.reset_data(m, d, reset=[True, False]) + def test_reset_data_keyframe(self): """Tests that reset_data_keyframe matches mj_resetDataKeyframe.""" _MJCF = """ @@ -1186,6 +1208,34 @@ def test_reset_data_keyframe_per_world(self): _assert_eq(d.qpos.numpy()[2], 3.0, "qpos[2]") _assert_eq(d.qpos.numpy()[3], 4.0, "qpos[3]") + def test_reset_data_keyframe_key_invalid(self): + """Tests that reset_data_keyframe validates the key argument.""" + _MJCF = """ + + + + + + + + + + + + """ + mjm = mujoco.MjModel.from_xml_string(_MJCF) + m = mjwarp.put_model(mjm) + d = mjwarp.make_data(mjm, nworld=2) + + with self.assertRaisesRegex(ValueError, r"key \(-1\) must be in \[0, 1\)"): + mjwarp.reset_data_keyframe(m, d, -1) + with self.assertRaisesRegex(ValueError, r"key \(1\) must be in \[0, 1\)"): + mjwarp.reset_data_keyframe(m, d, 1) + with self.assertRaisesRegex(ValueError, "key array must have shape"): + mjwarp.reset_data_keyframe(m, d, wp.array(np.array([0, 0, 0]), dtype=int)) + with self.assertRaisesRegex(ValueError, "key must be an int or a wp.array"): + mjwarp.reset_data_keyframe(m, d, 0.5) + def test_sdf(self): """Tests that an SDF can be loaded.""" mjm, mjd, m, d = test_data.fixture("collision_sdf/cow.xml")