### Description <!-- Describe your request. --> ### Context <!-- What is the motivation for this request? --> https://github.com/google-deepmind/mujoco_warp/issues/319