Copybara import of the project:

--
6da20df by huy-ha <[email protected]>:

generalize mocap body initialization to multiple mocap bodies

COPYBARA_INTEGRATE_REVIEW=#33 from huy-ha:main 6da20df
PiperOrigin-RevId: 718955482
Change-Id: I923a27b2f7af679e0d0feae4b518751effdc678f

Author: huy-ha

mujoco_playground/_src/mjx_env.py

@@ -148,9 +148,9 @@ def init(
   if act is not None:
     data = data.replace(act=act)
   if mocap_pos is not None:
-    data = data.replace(mocap_pos=mocap_pos.reshape(1, -1))
+    data = data.replace(mocap_pos=mocap_pos.reshape(len(data.mocap_pos), -1))
   if mocap_quat is not None:
-    data = data.replace(mocap_quat=mocap_quat.reshape(1, -1))
+    data = data.replace(mocap_quat=mocap_quat.reshape(len(data.mocap_quat), -1))
   data = mjx.forward(model, data)
   return data