island discovery: dfs discovery with flood fill

mujoco_warp/__init__.py

@@ -56,6 +56,7 @@
 from mujoco_warp._src.io import set_const as set_const
 from mujoco_warp._src.io import set_const_0 as set_const_0
 from mujoco_warp._src.io import set_const_fixed as set_const_fixed
+from mujoco_warp._src.island import island as island
 from mujoco_warp._src.passive import passive as passive
 from mujoco_warp._src.ray import ray as ray
 from mujoco_warp._src.ray import rays as rays

mujoco_warp/_src/forward.py

@@ -20,6 +20,7 @@
 from mujoco_warp._src import collision_driver
 from mujoco_warp._src import constraint
 from mujoco_warp._src import derivative
+from mujoco_warp._src import island
 from mujoco_warp._src import math
 from mujoco_warp._src import passive
 from mujoco_warp._src import sensor
@@ -517,6 +518,8 @@ def fwd_position(m: Model, d: Data, factorize: bool = True):
   if m.opt.run_collision_detection:
     collision_driver.collision(m, d)
   constraint.make_constraint(m, d)
+  if not (m.opt.disableflags & DisableBit.ISLAND):
+    island.island(m, d)
   smooth.transmission(m, d)
 
 

mujoco_warp/_src/io.py

@@ -740,6 +740,9 @@ def make_data(
     "eq_active": wp.array(np.tile(mjm.eq_active0.astype(bool), (nworld, 1)), shape=(nworld, mjm.neq), dtype=bool),
     # flexedge
     "flexedge_J": None,
+    # island arrays
+    "nisland": None,
+    "tree_island": None,
   }
   for f in dataclasses.fields(types.Data):
     if f.name in d_kwargs:
@@ -757,6 +760,10 @@ def make_data(
 
   d.flexedge_J = wp.zeros((nworld, 1, mjd.flexedge_J.size), dtype=float)
 
+  # island discovery arrays
+  d.nisland = wp.zeros((nworld,), dtype=int)
+  d.tree_island = wp.zeros((nworld, mjm.ntree), dtype=int)
+
   return d
 
 
@@ -891,6 +898,9 @@ def put_data(
     "actuator_moment": None,
     "flexedge_J": None,
     "nacon": None,
+    # island arrays
+    "nisland": None,
+    "tree_island": None,
   }
   for f in dataclasses.fields(types.Data):
     if f.name in d_kwargs:
@@ -918,6 +928,10 @@ def put_data(
 
   d.flexedge_J = wp.array(np.tile(mjd.flexedge_J.reshape(-1), (nworld, 1)).reshape((nworld, 1, -1)), dtype=float)
 
+  # island arrays
+  d.nisland = wp.zeros((nworld,), dtype=int)
+  d.tree_island = wp.zeros((nworld, mjm.ntree), dtype=int)
+
   if mujoco.mj_isSparse(mjm):
     ten_J = np.zeros((mjm.ntendon, mjm.nv))
     mujoco.mju_sparse2dense(ten_J, mjd.ten_J.reshape(-1), mjd.ten_J_rownnz, mjd.ten_J_rowadr, mjd.ten_J_colind.reshape(-1))
@@ -1138,6 +1152,12 @@ def get_data_into(
   # sensors
   result.sensordata[:] = d.sensordata.numpy()[world_id]
 
+  # islands
+  nisland = d.nisland.numpy()[world_id]
+  result.nisland = nisland
+  if nisland:
+    result.tree_island[:] = d.tree_island.numpy()[world_id]
+
 
 def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
   """Clear data, set defaults; optionally by world.

mujoco_warp/_src/io_test.py

@@ -124,11 +124,10 @@ def test_get_data_into_m(self):
   )
   def test_get_data_into(self, nworld, world_id):
     # keyframe=0: ncon=8, nefc=32
-    mjm, mjd, _, d = test_data.fixture("humanoid/humanoid.xml", keyframe=0, nworld=nworld)
+    mjm, mjd, m, d = test_data.fixture("humanoid/humanoid.xml", keyframe=0, nworld=nworld)
 
     # keyframe=2: ncon=0, nefc=0
     mujoco.mj_resetDataKeyframe(mjm, mjd, 2)
-    d.time.fill_(0.12345)
 
     # check that mujoco._functions._realloc_con_efc allocates for contact and efc
     mjwarp.get_data_into(mjd, mjm, d, world_id=world_id)
@@ -141,6 +140,7 @@ def test_get_data_into(self, nworld, world_id):
     self.assertEqual(d.ne.numpy()[world_id], mjd.ne)
     self.assertEqual(d.nf.numpy()[world_id], mjd.nf)
     self.assertEqual(d.nl.numpy()[world_id], mjd.nl)
+    self.assertEqual(d.nisland.numpy()[world_id], mjd.nisland)
     _assert_eq(d.time.numpy()[world_id], mjd.time, "time")
 
     for field in [
@@ -296,7 +296,10 @@ def test_get_data_into_io_test_models(self, xml):
       "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_result, field).reshape(-1),

mujoco_warp/_src/island.py

@@ -19,6 +19,7 @@
 from mujoco_warp._src.types import ConstraintType
 from mujoco_warp._src.types import EqType
 from mujoco_warp._src.types import ObjType
+from mujoco_warp._src.warp_util import event_scope
 
 
 @wp.kernel
@@ -44,7 +45,7 @@ def _tree_edges(
   # Out:
   tree_tree: wp.array3d(dtype=int),  # kernel_analyzer: off
 ):
-  """Find tree edges. Launch: (nworld, njmax)."""
+  """Find tree edges."""
   worldid, efcid = wp.tid()
 
   # skip if beyond active constraints
@@ -176,3 +177,92 @@ def tree_edges(m: types.Model, d: types.Data, tree_tree: wp.array3d(dtype=int)):
     ],
     outputs=[tree_tree],
   )
+
+
+@wp.kernel
+def _flood_fill(
+  # Model:
+  ntree: int,
+  # In:
+  tree_tree_in: wp.array3d(dtype=int),
+  labels_in: wp.array2d(dtype=int),
+  stack_in: wp.array2d(dtype=int),
+  # Data out:
+  nisland_out: wp.array(dtype=int),
+  # Out:
+  labels_out: wp.array2d(dtype=int),
+  stack_out: wp.array2d(dtype=int),
+):
+  """DFS flood fill to discover islands using tree_tree matrix."""
+  worldid = wp.tid()
+  nisland = int(0)
+
+  # iterate over trees
+  for i in range(ntree):
+    # already assigned
+    if labels_in[worldid, i] != -1:
+      continue
+
+    # check if tree has any edges
+    has_edge = int(0)
+    for j in range(ntree):
+      if tree_tree_in[worldid, i, j] != 0:
+        has_edge = 1
+        break
+    if has_edge == 0:
+      continue
+
+    # DFS: push i onto stack
+    nstack = int(0)
+    stack_out[worldid, nstack] = i
+    nstack = nstack + 1
+
+    while nstack > 0:
+      # pop v from stack
+      nstack = nstack - 1
+      v = stack_in[worldid, nstack]
+
+      # already assigned
+      if labels_in[worldid, v] != -1:
+        continue
+
+      # assign to current island
+      labels_out[worldid, v] = nisland
+
+      # push neighbors
+      for neighbor in range(ntree):
+        if tree_tree_in[worldid, v, neighbor] != 0:
+          if labels_in[worldid, neighbor] == -1 and nstack < ntree:
+            stack_out[worldid, nstack] = neighbor
+            nstack = nstack + 1
+
+    # island filled
+    nisland = nisland + 1
+
+  nisland_out[worldid] = nisland
+
+
+@event_scope
+def island(
+  m: types.Model,
+  d: types.Data,
+) -> None:
+  """Discover constraint islands."""
+  if m.ntree == 0:
+    d.nisland.zero_()
+    return
+
+  # Step 1: Find tree edges
+  tree_tree = wp.zeros((d.nworld, m.ntree, m.ntree), dtype=int)
+  tree_edges(m, d, tree_tree)
+
+  # Step 2: DFS flood fill
+  d.tree_island.fill_(-1)
+  stack_scratch = wp.zeros((d.nworld, m.ntree), dtype=int)
+
+  wp.launch(
+    _flood_fill,
+    dim=d.nworld,
+    inputs=[m.ntree, tree_tree, d.tree_island, stack_scratch],
+    outputs=[d.nisland, d.tree_island, stack_scratch],
+  )