Add reset_data_keyframe to mirror mj_resetDataKeyframe in MuJoCo

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

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(