Add kernel warning system for overflow detection

mujoco_warp/__init__.py

@@ -99,3 +99,7 @@
 from mujoco_warp._src.types import State as State
 from mujoco_warp._src.types import Statistic as Statistic
 from mujoco_warp._src.types import TrnType as TrnType
+from mujoco_warp._src.types import WarningType as WarningType
+from mujoco_warp._src.warning import check_warnings as check_warnings
+from mujoco_warp._src.warning import clear_warnings as clear_warnings
+from mujoco_warp._src.warning import get_warnings as get_warnings

mujoco_warp/_src/benchmark.py

@@ -25,6 +25,7 @@
 from mujoco_warp._src.types import Data
 from mujoco_warp._src.types import Model
 from mujoco_warp._src.util_misc import halton
+from mujoco_warp._src.warning import check_warnings
 
 
 def _sum(stack1, stack2):
@@ -145,6 +146,9 @@ def benchmark(
         wp.synchronize()
         run_end = time.perf_counter()
 
+        # Check and emit any overflow warnings, then clear flags
+        check_warnings(d, clear=True)
+
       time_vec[i] = run_end - run_beg
       if trace:
         trace = _sum(trace, tracer.trace())

mujoco_warp/_src/collision_convex.py

@@ -34,6 +34,7 @@
 from mujoco_warp._src.types import EnableBit
 from mujoco_warp._src.types import GeomType
 from mujoco_warp._src.types import Model
+from mujoco_warp._src.types import WarningType
 from mujoco_warp._src.types import mat43
 from mujoco_warp._src.types import mat63
 from mujoco_warp._src.types import vec5
@@ -154,6 +155,7 @@ def ccd_hfield_kernel_builder(
   geomtype2: int,
   gjk_iterations: int,
   epa_iterations: int,
+  warning_printf: bool,
 ):
   """Kernel builder for heightfield CCD collisions (no multiccd args)."""
 
@@ -232,6 +234,9 @@ def ccd_hfield_kernel(
     contact_type_out: wp.array(dtype=int),
     contact_geomcollisionid_out: wp.array(dtype=int),
     nacon_out: wp.array(dtype=int),
+    # Warning output:
+    warning_out: wp.array(dtype=int),
+    warning_info_out: wp.array2d(dtype=int),
   ):
     tid = wp.tid()
     if tid >= ncollision_in[0]:
@@ -403,10 +408,13 @@ def ccd_hfield_kernel(
         # add both triangles from this cell
         for i in range(2):
           if count >= MJ_MAXCONPAIR:
-            wp.printf(
-              "height field collision overflow, number of collisions >= %u - please adjust resolution: \n decrease the number of hfield rows/cols or modify size of colliding geom\n",
-              MJ_MAXCONPAIR,
-            )
+            if wp.static(warning_printf):
+              wp.printf(
+                "height field collision overflow, number of collisions >= %u - please adjust resolution: \n decrease the number of hfield rows/cols or modify size of colliding geom\n",
+                MJ_MAXCONPAIR,
+              )
+            wp.atomic_max(warning_out, int(WarningType.HFIELD_OVERFLOW), 1)
+            wp.atomic_max(warning_info_out, int(WarningType.HFIELD_OVERFLOW), 0, MJ_MAXCONPAIR)
             continue
 
           # add vert
@@ -453,6 +461,9 @@ def ccd_hfield_kernel(
             epa_pr,
             epa_norm2,
             epa_horizon,
+            wp.static(warning_printf),
+            warning_out,
+            warning_info_out,
           )
 
           if ncontact == 0:
@@ -692,6 +703,7 @@ def ccd_kernel_builder(
   gjk_iterations: int,
   epa_iterations: int,
   use_multiccd: bool,
+  warning_printf: bool,
 ):
   """Kernel builder for non-heightfield CCD collisions (no hfield args)."""
 
@@ -751,6 +763,9 @@ def eval_ccd_write_contact(
     contact_type_out: wp.array(dtype=int),
     contact_geomcollisionid_out: wp.array(dtype=int),
     nacon_out: wp.array(dtype=int),
+    # Warning output:
+    warning_out: wp.array(dtype=int),
+    warning_info_out: wp.array2d(dtype=int),
   ) -> int:
     points = mat43()
     witness1 = mat43()
@@ -781,6 +796,9 @@ def eval_ccd_write_contact(
       epa_pr_in[tid],
       epa_norm2_in[tid],
       epa_horizon_in[tid],
+      wp.static(warning_printf),
+      warning_out,
+      warning_info_out,
     )
 
     if dist >= 0.0 and pairid[1] == -1:
@@ -953,6 +971,9 @@ def ccd_kernel(
     contact_type_out: wp.array(dtype=int),
     contact_geomcollisionid_out: wp.array(dtype=int),
     nacon_out: wp.array(dtype=int),
+    # Warning output:
+    warning_out: wp.array(dtype=int),
+    warning_info_out: wp.array2d(dtype=int),
   ):
     tid = wp.tid()
     if tid >= ncollision_in[0]:
@@ -1064,6 +1085,8 @@ def ccd_kernel(
       contact_type_out,
       contact_geomcollisionid_out,
       nacon_out,
+      warning_out,
+      warning_info_out,
     )
 
   return ccd_kernel
@@ -1150,6 +1173,8 @@ def _pair_count(p1: int, p2: int) -> int:
     d.contact.type,
     d.contact.geomcollisionid,
     d.nacon,
+    d.warning,
+    d.warning_info,
   ]
 
   # Launch heightfield collision kernels (no multiccd args, 72 args total)
@@ -1158,7 +1183,7 @@ def _pair_count(p1: int, p2: int) -> int:
     g2 = geom_pair[1].value
     if (g1 == GeomType.HFIELD or g2 == GeomType.HFIELD) and _pair_count(g1, g2):
       wp.launch(
-        ccd_hfield_kernel_builder(g1, g2, m.opt.ccd_iterations, epa_iterations),
+        ccd_hfield_kernel_builder(g1, g2, m.opt.ccd_iterations, epa_iterations, m.opt.warning_printf),
         dim=d.naconmax,
         inputs=[
           m.opt.ccd_tolerance,
@@ -1249,7 +1274,7 @@ def _pair_count(p1: int, p2: int) -> int:
     g2 = geom_pair[1].value
     if g1 != GeomType.HFIELD and g2 != GeomType.HFIELD and _pair_count(g1, g2):
       wp.launch(
-        ccd_kernel_builder(g1, g2, m.opt.ccd_iterations, epa_iterations, use_multiccd),
+        ccd_kernel_builder(g1, g2, m.opt.ccd_iterations, epa_iterations, use_multiccd, m.opt.warning_printf),
         dim=d.naconmax,
         inputs=[
           m.opt.ccd_tolerance,

mujoco_warp/_src/collision_gjk.py

@@ -20,6 +20,7 @@
 
 from mujoco_warp._src.collision_primitive import Geom
 from mujoco_warp._src.types import GeomType
+from mujoco_warp._src.types import WarningType
 from mujoco_warp._src.types import mat43
 from mujoco_warp._src.types import mat63
 
@@ -571,6 +572,10 @@ def gjk(
   geomtype2: int,
   cutoff: float,
   is_discrete: bool,
+  warning_printf: bool,
+  # Data out:
+  warning_out: wp.array(dtype=int),
+  warning_info_out: wp.array2d(dtype=int),
 ) -> GJKResult:
   """Find distance within a tolerance between two geoms."""
   cutoff2 = cutoff * cutoff
@@ -667,7 +672,10 @@ def gjk(
     cnt += 1
 
   if cnt == gjk_iterations:
-    wp.printf("Warning: opt.ccd_iterations, currently set to %d, needs to be increased.\n", gjk_iterations)
+    if warning_printf:
+      wp.printf("Warning: opt.ccd_iterations, currently set to %d, needs to be increased.\n", gjk_iterations)
+    wp.atomic_max(warning_out, int(WarningType.GJK_ITERATIONS), 1)
+    wp.atomic_max(warning_info_out, int(WarningType.GJK_ITERATIONS), 0, gjk_iterations)
 
   result = GJKResult()
 
@@ -1218,6 +1226,10 @@ def _epa(
   geomtype1: int,
   geomtype2: int,
   is_discrete: bool,
+  warning_printf: bool,
+  # Data out:
+  warning_out: wp.array(dtype=int),
+  warning_info_out: wp.array2d(dtype=int),
 ) -> Tuple[float, wp.vec3, wp.vec3, int]:
   """Recover penetration data from two geoms in contact given an initial polytope."""
   upper = FLOAT_MAX
@@ -1288,7 +1300,10 @@ def _epa(
     pt.nhorizon = _add_edge(pt, face[1], face[2])
     pt.nhorizon = _add_edge(pt, face[2], face[0])
     if pt.nhorizon == -1:
-      wp.printf("Warning: EPA horizon = %d isn't large enough.\n", pt.horizon.shape[0])
+      if warning_printf:
+        wp.printf("Warning: EPA horizon = %d isn't large enough.\n", pt.horizon.shape[0])
+      wp.atomic_max(warning_out, int(WarningType.EPA_HORIZON), 1)
+      wp.atomic_max(warning_info_out, int(WarningType.EPA_HORIZON), 0, pt.horizon.shape[0])
       idx = -1
       break
 
@@ -1305,7 +1320,10 @@ def _epa(
         pt.nhorizon = _add_edge(pt, face[1], face[2])
         pt.nhorizon = _add_edge(pt, face[2], face[0])
         if pt.nhorizon == -1:
-          wp.printf("Warning: EPA horizon = %d isn't large enough.\n", pt.horizon.shape[0])
+          if warning_printf:
+            wp.printf("Warning: EPA horizon = %d isn't large enough.\n", pt.horizon.shape[0])
+          wp.atomic_max(warning_out, int(WarningType.EPA_HORIZON), 1)
+          wp.atomic_max(warning_info_out, int(WarningType.EPA_HORIZON), 0, pt.horizon.shape[0])
           idx = -1
           break
 
@@ -1333,7 +1351,10 @@ def _epa(
     cnt += 1
 
   if cnt == epa_iterations:
-    wp.printf("Warning: opt.ccd_iterations, currently set to %d, needs to be increased.\n", gjk_iterations)
+    if warning_printf:
+      wp.printf("Warning: opt.ccd_iterations, currently set to %d, needs to be increased.\n", gjk_iterations)
+    wp.atomic_max(warning_out, int(WarningType.GJK_ITERATIONS), 1)
+    wp.atomic_max(warning_info_out, int(WarningType.GJK_ITERATIONS), 0, gjk_iterations)
 
   # return from valid face
   if idx > -1:
@@ -2232,6 +2253,10 @@ def ccd(
   face_pr: wp.array(dtype=wp.vec3),
   face_norm2: wp.array(dtype=float),
   horizon: wp.array(dtype=int),
+  warning_printf: bool,
+  # Data out:
+  warning_out: wp.array(dtype=int),
+  warning_info_out: wp.array2d(dtype=int),
 ) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
   """General convex collision detection via GJK/EPA."""
   full_margin1 = 0.0
@@ -2257,7 +2282,21 @@ def ccd(
   # special handling for sphere and capsule (shrink to point and line respectively)
   if size1 + size2 > 0.0:
     cutoff += full_margin1 + full_margin2
-    result = gjk(tolerance, gjk_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, is_discrete)
+    result = gjk(
+      tolerance,
+      gjk_iterations,
+      geom1,
+      geom2,
+      x_1,
+      x_2,
+      geomtype1,
+      geomtype2,
+      cutoff,
+      is_discrete,
+      warning_printf,
+      warning_out,
+      warning_info_out,
+    )
 
     # shallow penetration, inflate contact
     if result.dist > tolerance:
@@ -2273,7 +2312,21 @@ def ccd(
     geom2.size = wp.vec3(size2, geom2.size[1], geom2.size[2])
     cutoff -= full_margin1 + full_margin2
 
-  result = gjk(tolerance, gjk_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, is_discrete)
+  result = gjk(
+    tolerance,
+    gjk_iterations,
+    geom1,
+    geom2,
+    x_1,
+    x_2,
+    geomtype1,
+    geomtype2,
+    cutoff,
+    is_discrete,
+    warning_printf,
+    warning_out,
+    warning_info_out,
+  )
 
   # no penetration depth to recover
   if result.dist > tolerance or result.dim < 2:
@@ -2349,7 +2402,20 @@ def ccd(
   if pt.status:
     return result.dist, 1, result.x1, result.x2, -1
 
-  dist, x1, x2, idx = _epa(tolerance, gjk_iterations, epa_iterations, pt, geom1, geom2, geomtype1, geomtype2, is_discrete)
+  dist, x1, x2, idx = _epa(
+    tolerance,
+    gjk_iterations,
+    epa_iterations,
+    pt,
+    geom1,
+    geom2,
+    geomtype1,
+    geomtype2,
+    is_discrete,
+    warning_printf,
+    warning_out,
+    warning_info_out,
+  )
   if idx == -1:
     return FLOAT_MAX, 0, wp.vec3(), wp.vec3(), -1
 

mujoco_warp/_src/collision_gjk_test.py

@@ -26,6 +26,7 @@
 from mujoco_warp._src.collision_primitive import Geom
 from mujoco_warp._src.types import MJ_MAX_EPAFACES
 from mujoco_warp._src.types import MJ_MAX_EPAHORIZON
+from mujoco_warp._src.types import NUM_WARNINGS
 
 
 def _geom_dist(
@@ -108,6 +109,9 @@ def _ccd_kernel(
     endvert: wp.array(dtype=wp.vec3),
     face1: wp.array(dtype=wp.vec3),
     face2: wp.array(dtype=wp.vec3),
+    # Data out:
+    warning_out: wp.array(dtype=int),
+    warning_info_out: wp.array2d(dtype=int),
     # Out:
     dist_out: wp.array(dtype=float),
     ncon_out: wp.array(dtype=int),
@@ -200,6 +204,9 @@ def _ccd_kernel(
       face_pr,
       face_norm2,
       horizon,
+      False,  # warning_printf
+      warning_out,
+      warning_info_out,
     )
 
     if wp.static(multiccd):
@@ -236,6 +243,8 @@ def _ccd_kernel(
   dist_out = wp.array(shape=(1,), dtype=float)
   ncon_out = wp.array(shape=(1,), dtype=int)
   pos_out = wp.array(shape=(2,), dtype=wp.vec3)
+  warning_out = wp.zeros(NUM_WARNINGS, dtype=int)
+  warning_info_out = wp.zeros((NUM_WARNINGS, 2), dtype=int)
   wp.launch(
     _ccd_kernel,
     dim=1,
@@ -280,6 +289,8 @@ def _ccd_kernel(
       multiccd_endvert,
       multiccd_face1,
       multiccd_face2,
+      warning_out,
+      warning_info_out,
     ],
     outputs=[
       dist_out,