update test_get_data_into_io_test_models (#1189)

* update test_get_data_into_io_test_models

* fix tendon_armature indexing

* fix more indexing bugs

* fix formatting

Author: thowell

mujoco_warp/_src/forward.py

@@ -647,7 +647,7 @@ def _actuator_force(
       act = act_in[worldid, act_last]
       act_dot = (ctrl - act) / wp.max(dynprm[0], MJ_MINVAL)
     elif dyntype == DynType.MUSCLE:
-      dynprm = actuator_dynprm[worldid, uid]
+      dynprm = actuator_dynprm[worldid % actuator_dynprm.shape[0], uid]
       act = act_in[worldid, act_last]
       act_dot = util_misc.muscle_dynamics(ctrl, act, dynprm)
     else:  # DynType.NONE

mujoco_warp/_src/io_test.py

@@ -25,6 +25,8 @@
 from absl.testing import parameterized
 
 import mujoco_warp as mjwarp
+from mujoco_warp import ConeType
+from mujoco_warp import IntegratorType
 from mujoco_warp import test_data
 from mujoco_warp._src import warp_util
 from mujoco_warp._src.io import set_length_range
@@ -272,53 +274,58 @@ def test_get_data_into(self, nworld, world_id):
         field,
       )
 
-  @parameterized.parameters(*_IO_TEST_MODELS)
-  def test_get_data_into_io_test_models(self, xml):
+  @parameterized.product(
+    xml=_IO_TEST_MODELS,
+    cone=list(ConeType),
+    integrator=list(IntegratorType),
+  )
+  def test_get_data_into_io_test_models(self, xml, cone, integrator):
     """Tests get_data_into for field coverage across diverse model types."""
-    mjm, mjd, _, d = test_data.fixture(xml)
-
-    # Create fresh MjData to verify get_data_into populates it correctly
-    mjd_result = mujoco.MjData(mjm)
-
-    mjwarp.get_data_into(mjd_result, mjm, d)
-
-    # Compare key fields, including flex/tendon data not covered by humanoid.xml
-    for field in [
-      "qpos",
-      "qvel",
-      "qacc",
-      "ctrl",
-      "act",
-      "flexvert_xpos",
-      "flexedge_length",
-      "flexedge_velocity",
-      "ten_length",
-      "ten_velocity",
-      "actuator_length",
-      "actuator_velocity",
-      "actuator_force",
-      "xpos",
-      "xquat",
-      "geom_xpos",
-      "tree_island",
-    ]:
-      if field == "tree_island" and d.nisland.numpy()[0] == 0:
-        continue
-      if getattr(mjd, field).size > 0:
+    mjm, _, m, d = test_data.fixture(xml, nworld=2, overrides={"opt.cone": cone, "opt.integrator": integrator})
+    mjwarp.step(m, d)
+
+    for world_id in range(2):
+      # Create reference MjData from warp data (resizes contact/efc fields internally)
+      mjd = mujoco.MjData(mjm)
+      mjwarp.get_data_into(mjd, mjm, d, world_id=world_id)
+
+      # Compare key fields, including flex/tendon data not covered by humanoid.xml
+      for field in [
+        "qpos",
+        "qvel",
+        "qacc",
+        "ctrl",
+        "act",
+        "flexvert_xpos",
+        "flexedge_length",
+        "flexedge_velocity",
+        "ten_length",
+        "ten_velocity",
+        "actuator_length",
+        "actuator_velocity",
+        "actuator_force",
+        "xpos",
+        "xquat",
+        "geom_xpos",
+        "tree_island",
+      ]:
+        if field == "tree_island" and d.nisland.numpy()[0] == 0:
+          continue
+        if getattr(mjd, field).size > 0:
+          _assert_eq(
+            getattr(mjd, field).reshape(-1),
+            getattr(d, field).numpy()[world_id].reshape(-1),
+            f"{field} (model: {xml}, world: {world_id})",
+          )
+
+      # flexedge_J
+      if xml == "flex/floppy.xml":
         _assert_eq(
-          getattr(mjd_result, field).reshape(-1),
-          getattr(mjd, field).reshape(-1),
-          f"{field} (model: {xml})",
+          mjd.flexedge_J.reshape(-1),
+          d.flexedge_J.numpy()[world_id].reshape(-1),
+          f"flexedge_J (world: {world_id})",
         )
 
-    # flexedge_J
-    if xml == "flex/floppy.xml":
-      _assert_eq(
-        mjd_result.flexedge_J.reshape(-1),
-        d.flexedge_J.numpy()[0].reshape(-1),
-        "flexedge_J",
-      )
-
   def test_ellipsoid_fluid_model(self):
     mjm = mujoco.MjModel.from_xml_string(
       """

mujoco_warp/_src/smooth.py

@@ -908,7 +908,7 @@ def _tendon_armature(
   if is_sparse:  # is_sparse is not batched
     madr_ij = dof_Madr[dofid]
 
-  armature = tendon_armature[worldid, tenid]
+  armature = tendon_armature[worldid % tendon_armature.shape[0], tenid]
 
   if armature == 0.0:
     return
@@ -1321,7 +1321,7 @@ def _cfrc_ext_equality(
   )
 
   id = efc_id_in[worldid, efcid]
-  eq_data_ = eq_data[worldid, id]
+  eq_data_ = eq_data[worldid % eq_data.shape[0], id]
   body_semantic = eq_objtype[id] == ObjType.BODY
 
   obj1 = eq_obj1id[id]
@@ -1342,7 +1342,7 @@ def _cfrc_ext_equality(
       else:
         offset = wp.vec3(eq_data_[3], eq_data_[4], eq_data_[5])
     else:
-      offset = site_pos[worldid, obj1]
+      offset = site_pos[worldid % site_pos.shape[0], obj1]
 
     # transform point on body1: local -> global
     pos = xmat_in[worldid, bodyid1] @ offset + xpos_in[worldid, bodyid1]
@@ -1364,7 +1364,7 @@ def _cfrc_ext_equality(
       else:
         offset = wp.vec3(eq_data_[0], eq_data_[1], eq_data_[2])
     else:
-      offset = site_pos[worldid, obj2]
+      offset = site_pos[worldid % site_pos.shape[0], obj2]
 
     # transform point on body2: local -> global
     pos = xmat_in[worldid, bodyid2] @ offset + xpos_in[worldid, bodyid2]
@@ -1559,7 +1559,7 @@ def _tendon_dot(
 ):
   worldid, tenid = wp.tid()
 
-  armature = tendon_armature[worldid, tenid]
+  armature = tendon_armature[worldid % tendon_armature.shape[0], tenid]
   if armature == 0.0:
     return
 
@@ -1717,7 +1717,7 @@ def _tendon_bias_coef(
 ):
   worldid, tenid, dofid = wp.tid()
 
-  armature = tendon_armature[worldid, tenid]
+  armature = tendon_armature[worldid % tendon_armature.shape[0], tenid]
   if armature == 0.0:
     return
 
@@ -1741,7 +1741,7 @@ def _tendon_bias_qfrc(
 ):
   worldid, tenid, dofid = wp.tid()
 
-  armature = tendon_armature[worldid, tenid]
+  armature = tendon_armature[worldid % tendon_armature.shape[0], tenid]
   if armature == 0.0:
     return
 
@@ -2078,7 +2078,7 @@ def _transmission(
     siteid = trnid[0]
     refid = trnid[1]
 
-    gear = actuator_gear[worldid, actid]
+    gear = actuator_gear[actuator_gear_id, actid]
     site_quat_id = worldid % site_quat.shape[0]
     gear_translation = wp.spatial_top(gear)
     gear_rotational = wp.spatial_bottom(gear)
@@ -2943,7 +2943,7 @@ def _spatial_geom_tendon(
   # get geom information
   geom_xpos = geom_xpos_in[worldid, wrap_objid_geom]
   geom_xmat = geom_xmat_in[worldid, wrap_objid_geom]
-  geomsize = geom_size[worldid, wrap_objid_geom][0]
+  geomsize = geom_size[worldid % geom_size.shape[0], wrap_objid_geom][0]
   geom_type = wrap_type[wrap_adr]
 
   # get body ids for site-geom-site instances