Small change.

PiperOrigin-RevId: 718975676 Change-Id: I4124ee2b397a31bbc5f2c6b6d2b37d0330a6fb65
google-deepmind · Jan 23, 2025 · ec28132 · ec28132
1 parent 76b235c
commit ec28132
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mujoco_playground/_src/mjx_env.py b/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(len(data.mocap_pos), -1))
+    data = data.replace(mocap_pos=mocap_pos.reshape(model.nmocap, -1))
   if mocap_quat is not None:
-    data = data.replace(mocap_quat=mocap_quat.reshape(len(data.mocap_quat), -1))
+    data = data.replace(mocap_quat=mocap_quat.reshape(model.nmocap, -1))
   data = mjx.forward(model, data)
   return data