Add reset_data_keyframe to mirror mj_resetDataKeyframe in MuJoCo

[sibocw]

Currently, mujoco_warp.reset_data mirrors mujoco.mj_resetData, but there's no mujoco_warp.reset_data_keyframe that mirrors mujoco.mj_resetDataKeyframe.

This PR aims to implement it by repeating exactly what mj_resetDataKeyframe does on the CPU: calling mj_resetData, then re-setting values defined in the keyframe.

The complication is that nkey and key_{time,qpos,qvel,act,mpos,mquat,ctrl} are not currently included in mujoco_warp.Model. Therefore, the reset-to-keypoint function would have to either

• receive the MjModel object as an additional input (this is the state after d599780), or
• rely on keyframe-related attributes being added to Model (this is the current state of this PR). The dynamic attribute setting logic in put_model automatically populates these fields, so we don't need to touch put_model.

Would this be something useful?

PS. I guess we could in principle also allow different worlds to be reset to different keyframes, but IDK if this is a "real" use case.

[thowell]

@erikfrey what do we think about adding reset_data_keyframe to the api?

[erikfrey]

@thowell I think it's appropriate to add this to the mjwarp API. I bet that env frameworks would call it for smarter resetting.

[thowell]

@sibocw the implementation looks great thanks for putting this together. lets update the implementation to set keyframes per world.

instead of batching all of the keyframe fields in Model we can have reset_data_keyframe accept either an int or wp.array(dtype=int). if an int is provide this value is used to set all worlds to the same keyframe. if wp.array(dtype=int) is provided it should have size nworld and will set keyframes per world. in the latter case if the keyframe provided for the world is < 0 or >=nkey then we skip resetting the world.

please let us know if you have any questions. thanks!

[sibocw]

@thowell How about this? I also noticed that there's no type/shape check for the per-world reset bitmask in reset_data, so I added it, along with checks for the target key in reset_data_keyframe. Let me know if you want me to get rid of them.

Thanks!