improve testing for batched Model fields

.github/workflows/ci.yml

@@ -69,6 +69,11 @@ jobs:
       - name: Test with pytest
         run: |
           pytest -n 4
+      - name: Out-of-bounds check for batched model fields
+        # TODO(team): currently only works on GPU, enable on CPU?
+        # https://nvidia.github.io/warp/user_guide/debugging.html#debug-mode-compilation
+        run: |
+          pytest -k io_test --debug_mode --verify_cuda
 
   kernel_analyzer:
     name: Kernel analyzer

conftest.py

@@ -28,6 +28,7 @@ def pytest_addoption(parser):
   )
   parser.addoption("--lineinfo", action="store_true", default=False, help="add lineinfo to warp kernel")
   parser.addoption("--optimization_level", action="store", default=None, type=int, help="set wp.config.optimization_level")
+  parser.addoption("--debug_mode", action="store_true", default=False, help="debug mode compilation")
 
 
 def pytest_configure(config):
@@ -38,3 +39,5 @@ def pytest_configure(config):
   if config.getoption("--lineinfo"):
     wp.config.lineinfo = True
   wp.config.optimization_level = config.getoption("--optimization_level")
+  if config.getoption("--debug_mode"):
+    wp.config.mode = "debug"

mujoco_warp/_src/io_test.py

@@ -15,6 +15,8 @@
 
 """Tests for io functions."""
 
+import dataclasses
+
 import mujoco
 import numpy as np
 import warp as wp
@@ -608,6 +610,39 @@ def test_eq_active(self, active, make_data):
 
     _assert_eq(d.eq_active.numpy()[0], mjd.eq_active, "eq_active")
 
+  def test_model_batched_fields(self):
+    """Test Model batched fields."""
+    nworld = 2
+    mjm, _, m, d = test_data.fixture("humanoid/humanoid.xml", keyframe=0, nworld=nworld)
+
+    for f in dataclasses.fields(m):
+      # TODO(team): test arrays that are warp only
+      if not hasattr(mjm, f.name):
+        continue
+      if isinstance(f.type, wp.array):
+        # get fields
+        arr = getattr(m, f.name)
+        mj_arr = getattr(mjm, f.name)
+
+        # check that field is not empty
+        if 0 in mj_arr.shape + arr.shape:
+          continue
+
+        # check for batched field
+        if hasattr(arr, "_is_batched") and arr._is_batched:
+          assert arr.shape[0] == 1
+
+          # reshape if necessary
+          if f.name in ("cam_mat0"):
+            mj_arr = mj_arr.reshape((-1, 3, 3))
+
+          # set batched field
+          setattr(m, f.name, wp.array(np.tile(mj_arr, (nworld,) + arr.shape[1:]), dtype=f.type.dtype))
+
+    mjwarp.forward(m, d)
+    mjwarp.reset_data(m, d)
+    mjwarp.forward(m, d)
+
 
 if __name__ == "__main__":
   wp.init()