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..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(
@@ -2576,6 +2586,236 @@ def reset_sleep(
sleep.update_sleep(m, d)
+def reset_data_keyframe(
+ m: types.Model,
+ d: types.Data,
+ key: int | wp.array,
+ reset: Optional[wp.array] = None,
+):
+ """Reset data, set fields from specified keyframe; optionally by world.
+
+ 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. 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 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:
+ reset_input = wp.ones(d.nworld, dtype=bool)
+ else:
+ reset_input = reset
+
+ # 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
+ 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)
+ 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_in: wp.array[int],
+ reset_in: wp.array[bool],
+ # Data out:
+ time_out: wp.array[float],
+ ):
+ worldid = wp.tid()
+
+ if not reset_in[worldid]:
+ return
+
+ key = key_in[worldid]
+ time_out[worldid] = key_time[key]
+
+ @wp.kernel(module="unique", enable_backward=False)
+ def reset_qpos(
+ # Model:
+ key_qpos: wp.array2d[float],
+ # In:
+ key_in: wp.array[int],
+ reset_in: wp.array[bool],
+ # Data out:
+ qpos_out: wp.array2d[float],
+ ):
+ worldid, qid = wp.tid()
+
+ 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)
+ def reset_qvel(
+ # Model:
+ key_qvel: wp.array2d[float],
+ # In:
+ key_in: wp.array[int],
+ reset_in: wp.array[bool],
+ # Data out:
+ qvel_out: wp.array2d[float],
+ ):
+ worldid, vid = wp.tid()
+
+ 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)
+ def reset_activation(
+ # Model:
+ key_act: wp.array2d[float],
+ # In:
+ key_in: wp.array[int],
+ reset_in: wp.array[bool],
+ # Data out:
+ act_out: wp.array2d[float],
+ ):
+ worldid, aid = wp.tid()
+
+ 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)
+ def reset_mocap(
+ # Model:
+ body_mocapid: wp.array[int],
+ key_mpos: wp.array2d[wp.vec3],
+ key_mquat: wp.array2d[wp.quat],
+ # In:
+ key_in: wp.array[int],
+ 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 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]
+
+ @wp.kernel(module="unique", enable_backward=False)
+ def reset_control(
+ # Model:
+ key_ctrl: wp.array2d[float],
+ # In:
+ key_in: wp.array[int],
+ reset_in: wp.array[bool],
+ # Data out:
+ ctrl_out: wp.array2d[float],
+ ):
+ worldid, cid = wp.tid()
+
+ if not reset_in[worldid]:
+ return
+
+ key = key_in[worldid]
+ ctrl_out[worldid, cid] = key_ctrl[key, cid]
+
+ wp.launch(
+ reset_time,
+ dim=d.nworld,
+ 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_input, reset_input_with_invalid_keys_masked],
+ outputs=[d.qpos],
+ )
+
+ wp.launch(
+ reset_qvel,
+ dim=(d.nworld, m.nv),
+ 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_input, reset_input_with_invalid_keys_masked],
+ outputs=[d.act],
+ )
+
+ wp.launch(
+ reset_mocap,
+ dim=(d.nworld, m.nbody),
+ inputs=[
+ m.body_mocapid,
+ m.key_mpos,
+ m.key_mquat,
+ key_input,
+ reset_input_with_invalid_keys_masked,
+ ],
+ outputs=[d.mocap_pos, d.mocap_quat],
+ )
+
+ wp.launch(
+ reset_control,
+ dim=(d.nworld, m.nu),
+ inputs=[m.key_ctrl, key_input, reset_input_with_invalid_keys_masked],
+ 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..09adcef65 100644
--- a/mujoco_warp/_src/io_test.py
+++ b/mujoco_warp/_src/io_test.py
@@ -1062,6 +1062,180 @@ 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 = """
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ """
+ reset_datafield = ["time", "qpos", "qvel", "act", "ctrl", "mocap_pos", "mocap_quat"]
+ key = 0
+
+ mjm, mjd, m, d = test_data.fixture(xml=_MJCF, keyframe=key)
+
+ # corrupt data
+ for arr in reset_datafield:
+ attr = getattr(d, arr)
+ if attr.dtype == float:
+ attr.fill_(wp.nan)
+ else:
+ attr.fill_(-1)
+
+ mjwarp.reset_data_keyframe(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(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(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(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_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_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")
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)