Import google-deepmind/mujoco_warp from GitHub.

PiperOrigin-RevId: 835230315
Change-Id: Ifa4f8f66c4e3de1a3074f0897daf93f88e469f06

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py

@@ -32,17 +32,18 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
+from mujoco.mjx.third_party.mujoco_warp._src.types import mat43
+from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
-from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel
+from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
 
 # TODO(team): improve compile time to enable backward pass
 wp.set_module_options({"enable_backward": False})
 
-MULTI_CONTACT_COUNT = 8
+vec_maxconpair = wp.types.vector(length=MJ_MAXCONPAIR, dtype=float)
-mat3c = wp.types.matrix(shape=(MULTI_CONTACT_COUNT, 3), dtype=float)
+mat_maxconpair = wp.types.matrix(shape=(MJ_MAXCONPAIR, 3), dtype=float)
-mat63 = wp.types.matrix(shape=(6, 3), dtype=float)
 
 _CONVEX_COLLISION_PAIRS = [
   (GeomType.HFIELD, GeomType.SPHERE),
@@ -260,9 +261,9 @@ def ccd_kernel_builder(
     contact_geomcollisionid_out: wp.array(dtype=int),
     nacon_out: wp.array(dtype=int),
   ) -> int:
-    points = mat3c()
+    points = mat43()
-    witness1 = mat3c()
+    witness1 = mat43()
-    witness2 = mat3c()
+    witness2 = mat43()
     geom1.margin = margin
     geom2.margin = margin
     if pairid[1] >= 0:
@@ -566,6 +567,33 @@ def ccd_kernel_builder(
 
       adr = hfield_adr[geom1_dataid]
 
+      hfield_contact_dist = vec_maxconpair()
+      hfield_contact_pos = mat_maxconpair()
+      hfield_contact_normal = mat_maxconpair()
+      min_dist = float(wp.inf)
+      min_normal = wp.vec3(wp.inf, wp.inf, wp.inf)
+      min_pos = wp.vec3(wp.inf, wp.inf, wp.inf)
+      min_id = int(-1)
+
+      # TODO(team): height field margin?
+      geom1.margin = margin
+      geom2.margin = margin
+
+      # EPA memory
+      epa_vert = epa_vert_in[tid]
+      epa_vert1 = epa_vert1_in[tid]
+      epa_vert2 = epa_vert2_in[tid]
+      epa_vert_index1 = epa_vert_index1_in[tid]
+      epa_vert_index2 = epa_vert_index2_in[tid]
+      epa_face = epa_face_in[tid]
+      epa_pr = epa_pr_in[tid]
+      epa_norm2 = epa_norm2_in[tid]
+      epa_index = epa_index_in[tid]
+      epa_map = epa_map_in[tid]
+      epa_horizon = epa_horizon_in[tid]
+
+      collision_pairid = collision_pairid_in[tid]
+
       # process all prisms in subgrid
       count = int(0)
       for r in range(rmin, rmax):
@@ -573,6 +601,13 @@ def ccd_kernel_builder(
         for c in range(cmin, cmax + 1):
           # 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,
+              )
+              continue
+
             # add vert
             x = dx * float(c) - size[0]
             y = dy * float(r + dr[i]) - size[1]
@@ -607,66 +642,253 @@ def ccd_kernel_builder(
               x1_ += prism[i]
             x1 += geom1.rot @ (x1_ / 6.0)
 
-            ncontact = eval_ccd_write_contact(
+            dist, ncontact, w1, w2, idx = ccd(
-              opt_ccd_tolerance,
+              opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
-              geom_type,
+              0.0,
-              naconmax_in,
+              ccd_iterations,
-              epa_vert_in,
-              epa_vert1_in,
-              epa_vert2_in,
-              epa_vert_index1_in,
-              epa_vert_index2_in,
-              epa_face_in,
-              epa_pr_in,
-              epa_norm2_in,
-              epa_index_in,
-              epa_map_in,
-              epa_horizon_in,
-              multiccd_polygon_in,
-              multiccd_clipped_in,
-              multiccd_pnormal_in,
-              multiccd_pdist_in,
-              multiccd_idx1_in,
-              multiccd_idx2_in,
-              multiccd_n1_in,
-              multiccd_n2_in,
-              multiccd_endvert_in,
-              multiccd_face1_in,
-              multiccd_face2_in,
               geom1,
               geom2,
-              geoms,
+              geomtype1,
-              worldid,
+              geomtype2,
-              tid,
-              margin,
-              gap,
-              condim,
-              friction,
-              solref,
-              solreffriction,
-              solimp,
               x1,
               geom2.pos,
-              count,
+              epa_vert,
-              collision_pairid_in[tid],
+              epa_vert1,
-              contact_dist_out,
+              epa_vert2,
-              contact_pos_out,
+              epa_vert_index1,
-              contact_frame_out,
+              epa_vert_index2,
-              contact_includemargin_out,
+              epa_face,
-              contact_friction_out,
+              epa_pr,
-              contact_solref_out,
+              epa_norm2,
-              contact_solreffriction_out,
+              epa_index,
-              contact_solimp_out,
+              epa_map,
-              contact_dim_out,
+              epa_horizon,
-              contact_geom_out,
-              contact_worldid_out,
-              contact_type_out,
-              contact_geomcollisionid_out,
-              nacon_out,
             )
-            count += ncontact
+
-            if count >= MJ_MAXCONPAIR:
+            if ncontact == 0:
-              return
+              continue
+
+            # cache contact information
+            hfield_contact_dist[count] = dist
+
+            pos = 0.5 * (w1 + w2)
+            hfield_contact_pos[count, 0] = pos[0]
+            hfield_contact_pos[count, 1] = pos[1]
+            hfield_contact_pos[count, 2] = pos[2]
+
+            frame = make_frame(w1 - w2)
+            normal = wp.vec3(frame[0, 0], frame[0, 1], frame[0, 2])
+            hfield_contact_normal[count, 0] = normal[0]
+            hfield_contact_normal[count, 1] = normal[1]
+            hfield_contact_normal[count, 2] = normal[2]
+
+            # contact with minimum distance
+            if dist < min_dist:
+              min_dist = dist
+              min_normal = normal
+              min_pos = pos
+              min_id = count
+
+            count += 1
+
+      # contact 0: minimum distance
+      write_contact(
+        naconmax_in,
+        0,
+        min_dist,
+        min_pos,
+        make_frame(min_normal),
+        margin,
+        gap,
+        condim,
+        friction,
+        solref,
+        solreffriction,
+        solimp,
+        geoms,
+        collision_pairid,
+        worldid,
+        contact_dist_out,
+        contact_pos_out,
+        contact_frame_out,
+        contact_includemargin_out,
+        contact_friction_out,
+        contact_solref_out,
+        contact_solreffriction_out,
+        contact_solimp_out,
+        contact_dim_out,
+        contact_geom_out,
+        contact_worldid_out,
+        contact_type_out,
+        contact_geomcollisionid_out,
+        nacon_out,
+      )
+
+      # TODO(team): routine for select subset of contacts
+      # TODO(team): if use_multiccd?
+      if wp.static(True):
+        MIN_DIST_TO_NEXT_CONTACT = 1.0e-3
+
+        # contact 1: furthest from minimum distance contact
+        id1 = int(-1)
+        dist1 = float(-wp.inf)
+        for i in range(count):
+          if i == min_id:
+            continue
+
+          hf_pos = wp.vec3(hfield_contact_pos[i, 0], hfield_contact_pos[i, 1], hfield_contact_pos[i, 2])
+          dist = wp.norm_l2(hf_pos - min_pos)
+
+          if dist > dist1:
+            id1 = i
+            dist1 = dist
+
+        if id1 == -1 or (0.0 < dist1 and dist1 < MIN_DIST_TO_NEXT_CONTACT):
+          return
+
+        pos1 = wp.vec3(hfield_contact_pos[id1, 0], hfield_contact_pos[id1, 1], hfield_contact_pos[id1, 2])
+        normal1 = wp.vec3(hfield_contact_normal[id1, 0], hfield_contact_normal[id1, 1], hfield_contact_normal[id1, 2])
+
+        write_contact(
+          naconmax_in,
+          1,
+          hfield_contact_dist[id1],
+          pos1,
+          make_frame(normal1),
+          margin,
+          gap,
+          condim,
+          friction,
+          solref,
+          solreffriction,
+          solimp,
+          geoms,
+          collision_pairid,
+          worldid,
+          contact_dist_out,
+          contact_pos_out,
+          contact_frame_out,
+          contact_includemargin_out,
+          contact_friction_out,
+          contact_solref_out,
+          contact_solreffriction_out,
+          contact_solimp_out,
+          contact_dim_out,
+          contact_geom_out,
+          contact_worldid_out,
+          contact_type_out,
+          contact_geomcollisionid_out,
+          nacon_out,
+        )
+
+        # contact 2: point furthest from min_pos - pos1 line
+        dist_min1 = wp.cross(min_normal, min_pos - pos1)
+
+        id2 = int(-1)
+        dist_12 = float(-wp.inf)
+        for i in range(count):
+          if i == min_id or i == id1:
+            continue
+
+          hf_pos = wp.vec3(hfield_contact_pos[i, 0], hfield_contact_pos[i, 1], hfield_contact_pos[i, 2])
+          dist = wp.abs(wp.dot(hf_pos - min_pos, dist_min1))
+
+          if dist > dist_12:
+            id2 = i
+            dist_12 = dist
+
+        if id2 == -1 or (0.0 < dist_12 and dist_12 < MIN_DIST_TO_NEXT_CONTACT):
+          return
+
+        pos2 = wp.vec3(hfield_contact_pos[id2, 0], hfield_contact_pos[id2, 1], hfield_contact_pos[id2, 2])
+        normal2 = wp.vec3(hfield_contact_normal[id2, 0], hfield_contact_normal[id2, 1], hfield_contact_normal[id2, 2])
+
+        write_contact(
+          naconmax_in,
+          2,
+          hfield_contact_dist[id2],
+          pos2,
+          make_frame(normal2),
+          margin,
+          gap,
+          condim,
+          friction,
+          solref,
+          solreffriction,
+          solimp,
+          geoms,
+          collision_pairid,
+          worldid,
+          contact_dist_out,
+          contact_pos_out,
+          contact_frame_out,
+          contact_includemargin_out,
+          contact_friction_out,
+          contact_solref_out,
+          contact_solreffriction_out,
+          contact_solimp_out,
+          contact_dim_out,
+          contact_geom_out,
+          contact_worldid_out,
+          contact_type_out,
+          contact_geomcollisionid_out,
+          nacon_out,
+        )
+
+        # contact 3: point furthest from other triangle edge
+        vec_min2 = wp.cross(min_normal, min_pos - pos2)
+        vec_12 = wp.cross(min_normal, pos1 - pos2)
+
+        id3 = int(-1)
+        dist3 = float(-wp.inf)
+        for i in range(count):
+          if i == min_id or i == id1 or i == id2:
+            continue
+
+          hf_pos = wp.vec3(hfield_contact_pos[i, 0], hfield_contact_pos[i, 1], hfield_contact_pos[i, 2])
+          dist = wp.abs(wp.dot(hf_pos - min_pos, vec_min2)) + wp.abs(wp.dot(pos1 - hf_pos, vec_12))
+
+          if dist > dist3:
+            id3 = i
+            dist3 = dist
+
+        if id3 == -1 or (0.0 < dist3 and dist3 < MIN_DIST_TO_NEXT_CONTACT):
+          return
+
+        pos3 = wp.vec3(hfield_contact_pos[id3, 0], hfield_contact_pos[id3, 1], hfield_contact_pos[id3, 2])
+        normal3 = wp.vec3(hfield_contact_normal[id3, 0], hfield_contact_normal[id3, 1], hfield_contact_normal[id3, 2])
+
+        write_contact(
+          naconmax_in,
+          3,
+          hfield_contact_dist[id3],
+          pos3,
+          make_frame(normal3),
+          margin,
+          gap,
+          condim,
+          friction,
+          solref,
+          solreffriction,
+          solimp,
+          geoms,
+          collision_pairid,
+          worldid,
+          contact_dist_out,
+          contact_pos_out,
+          contact_frame_out,
+          contact_includemargin_out,
+          contact_friction_out,
+          contact_solref_out,
+          contact_solreffriction_out,
+          contact_solimp_out,
+          contact_dim_out,
+          contact_geom_out,
+          contact_worldid_out,
+          contact_type_out,
+          contact_geomcollisionid_out,
+          nacon_out,
+        )
     else:
       eval_ccd_write_contact(
         opt_ccd_tolerance,

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py

@@ -27,9 +27,11 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
+from mujoco.mjx.third_party.mujoco_warp._src.types import mat23
+from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
-from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel
+from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
 
 wp.set_module_options({"enable_backward": False})
 
@@ -170,10 +172,6 @@ def _aabb_filter(
   return True
 
 
-mat23 = wp.types.matrix(shape=(2, 3), dtype=float)
-mat63 = wp.types.matrix(shape=(6, 3), dtype=float)
-
-
 # TODO(team): improve performance by precomputing bounding box
 @wp.func
 def _obb_filter(
@@ -232,7 +230,7 @@ def _obb_filter(
   return True
 
 
-def _broadphase_filter(m: Model):
+def _broadphase_filter(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int):
   @wp.func
   def func(
     # Model:
@@ -252,28 +250,28 @@ def _broadphase_filter(m: Model):
     # 4: aabb
     # 8: obb
 
-    aabb_id = worldid % geom_aabb.shape[0] if wp.static(m.geom_aabb.shape[0] > 1) else 0
+    aabb_id = worldid % ngeom_aabb if wp.static(ngeom_aabb > 1) else 0
     center1, center2 = geom_aabb[aabb_id, geom1, 0], geom_aabb[aabb_id, geom2, 0]
     size1, size2 = geom_aabb[aabb_id, geom1, 1], geom_aabb[aabb_id, geom2, 1]
 
-    rbound_id = worldid % geom_rbound.shape[0] if wp.static(m.geom_rbound.shape[0] > 1) else 0
+    rbound_id = worldid % ngeom_rbound if wp.static(ngeom_rbound > 1) else 0
     rbound1, rbound2 = geom_rbound[rbound_id, geom1], geom_rbound[rbound_id, geom2]
-    margin_id = worldid % geom_margin.shape[0] if wp.static(m.geom_margin.shape[0] > 1) else 0
+    margin_id = worldid % ngeom_margin if wp.static(ngeom_margin > 1) else 0
     margin1, margin2 = geom_margin[margin_id, geom1], geom_margin[margin_id, geom2]
     xpos1, xpos2 = geom_xpos_in[worldid, geom1], geom_xpos_in[worldid, geom2]
     xmat1, xmat2 = geom_xmat_in[worldid, geom1], geom_xmat_in[worldid, geom2]
 
     if rbound1 == 0.0 or rbound2 == 0.0:
-      if wp.static(m.opt.broadphase_filter & BroadphaseFilter.PLANE):
+      if wp.static(opt_broadphase_filter & BroadphaseFilter.PLANE):
         return _plane_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2, xmat1, xmat2)
     else:
-      if wp.static(m.opt.broadphase_filter & BroadphaseFilter.SPHERE):
+      if wp.static(opt_broadphase_filter & BroadphaseFilter.SPHERE):
         if not _sphere_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2):
           return False
-      if wp.static(m.opt.broadphase_filter & BroadphaseFilter.AABB):
+      if wp.static(opt_broadphase_filter & BroadphaseFilter.AABB):
         if not _aabb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2):
           return False
-      if wp.static(m.opt.broadphase_filter & BroadphaseFilter.OBB):
+      if wp.static(opt_broadphase_filter & BroadphaseFilter.OBB):
         if not _obb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2):
           return False
 
@@ -403,7 +401,7 @@ def _sap_range(
 
 
 @cache_kernel
-def _sap_broadphase(broadphase_filter):
+def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int):
   @nested_kernel(module="unique", enable_backward=False)
   def kernel(
     # Model:
@@ -461,7 +459,9 @@ def _sap_broadphase(broadphase_filter):
         continue
 
       if (
-        broadphase_filter(geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid)
+        wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin))(
+          geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid
+        )
         or pairid[1] >= 0
       ):
         _add_geom_pair(
@@ -578,9 +578,8 @@ def sap_broadphase(m: Model, d: Data):
   # estimate number of overlap checks
   # assumes each geom has 5 other geoms (batched over all worlds)
   nsweep = 5 * nworldgeom
-  broadphase_filter = _broadphase_filter(m)
   wp.launch(
-    kernel=_sap_broadphase(broadphase_filter),
+    kernel=_sap_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0]),
     dim=nsweep,
     inputs=[
       m.ngeom,
@@ -602,7 +601,7 @@ def sap_broadphase(m: Model, d: Data):
 
 
 @cache_kernel
-def _nxn_broadphase(broadphase_filter):
+def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int):
   @nested_kernel(module="unique", enable_backward=False)
   def kernel(
     # Model:
@@ -629,7 +628,9 @@ def _nxn_broadphase(broadphase_filter):
     geom2 = geom[1]
 
     if (
-      broadphase_filter(geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid)
+      wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin))(
+        geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid
+      )
       or nxn_pairid[elementid][1] >= 0
     ):
       _add_geom_pair(
@@ -663,9 +664,8 @@ def nxn_broadphase(m: Model, d: Data):
   The initial list of pairs is filtered at model creation time to exclude pairs based on
   `contype`/`conaffinity`, parent-child relationships, and explicit `<exclude>` tags.
   """
-  broadphase_filter = _broadphase_filter(m)
   wp.launch(
-    _nxn_broadphase(broadphase_filter),
+    _nxn_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0]),
     dim=(d.nworld, m.nxn_geom_pair_filtered.shape[0]),
     inputs=[
       m.geom_type,

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py

@@ -20,6 +20,8 @@ import warp as wp
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import Geom
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
+from mujoco.mjx.third_party.mujoco_warp._src.types import mat43
+from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
 
 # TODO(team): improve compile time to enable backward pass
 wp.set_module_options({"enable_backward": False})
@@ -32,13 +34,6 @@ MJ_MINVAL2 = MJ_MINVAL * MJ_MINVAL
 FACE_TOL = 0.99999872
 EDGE_TOL = 0.00159999931
 
-mat43 = wp.types.matrix(shape=(4, 3), dtype=float)
-mat63 = wp.types.matrix(shape=(6, 3), dtype=float)
-
-# box-box currently supports up to 8 contacts
-MULTI_CONTACT_COUNT = 8
-mat3c = wp.types.matrix(shape=(MULTI_CONTACT_COUNT, 3), dtype=float)
-
 
 @wp.struct
 class GJKResult:
@@ -155,30 +150,25 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint:
       vert_edgeadr = geom.graphadr + 2
       vert_globalid = geom.graphadr + 2 + numvert
       edge_localid = geom.graphadr + 2 + 2 * numvert
-      # hillclimb until no change
       prev = int(-1)
-      imax = int(0)
+      imax = wp.where(geom.index > -1, geom.index, 0)
-      if geom.index > -1:
-        imax = geom.index
-        sp.cached_index = geom.index
 
-      while True:
+      # hillclimb until no change
-        prev = int(imax)
+      while imax != prev:
-        i = int(geom.graph[vert_edgeadr + imax])
+        prev = imax
-        while geom.graph[edge_localid + i] >= 0:
+        i = geom.graph[vert_edgeadr + imax]
-          subidx = geom.graph[edge_localid + i]
+        subidx = geom.graph[edge_localid + i]
+        while subidx >= 0:
           idx = geom.graph[vert_globalid + subidx]
           dist = wp.dot(local_dir, geom.vert[geom.vertadr + idx])
-          if dist > max_dist:
+          imax = wp.where(dist > max_dist, subidx, imax)
-            max_dist = dist
+          max_dist = wp.where(dist > max_dist, dist, max_dist)
-            imax = int(subidx)
+          i += 1
-          i += int(1)
+          subidx = geom.graph[edge_localid + i]
-        if imax == prev:
+
-          break
       sp.cached_index = imax
-      imax = geom.graph[vert_globalid + imax]
+      sp.vertex_index = geom.graph[vert_globalid + imax]
-      sp.vertex_index = imax
+      sp.point = geom.vert[geom.vertadr + sp.vertex_index]
-      sp.point = geom.vert[geom.vertadr + imax]
 
     sp.point = geom.rot @ sp.point + geom.pos
   elif geomtype == GeomType.HFIELD:
@@ -1813,9 +1803,9 @@ def _polygon_clip(
   # Out:
   polygon_out: wp.array(dtype=wp.vec3),
   clipped_out: wp.array(dtype=wp.vec3),
-) -> Tuple[int, mat3c, mat3c]:
+) -> Tuple[int, mat43, mat43]:
-  witness1 = mat3c()
+  witness1 = mat43()
-  witness2 = mat3c()
+  witness2 = mat43()
 
   # clipping face needs to be at least a triangle
   if nface1 < 3:
@@ -1930,9 +1920,9 @@ def multicontact(
   geom2: Geom,
   geomtype1: int,
   geomtype2: int,
-) -> Tuple[int, mat3c, mat3c]:
+) -> Tuple[int, mat43, mat43]:
-  witness1 = mat3c()
+  witness1 = mat43()
-  witness2 = mat3c()
+  witness2 = mat43()
   witness1[0] = x1
   witness2[0] = x2
 

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py

@@ -38,21 +38,16 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
+from mujoco.mjx.third_party.mujoco_warp._src.types import mat43
+from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
-from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel
+from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
 
 wp.set_module_options({"enable_backward": False})
 
 
-class mat43f(wp.types.matrix(shape=(4, 3), dtype=wp.float32)):
-  pass
-
-
-mat63 = wp.types.matrix(shape=(6, 3), dtype=float)
-
-
 @wp.struct
 class Geom:
   pos: wp.vec3
@@ -169,7 +164,7 @@ def geom_collision_pair(
 
 
 @wp.func
-def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tuple[wp.vec4, mat43f, wp.vec3]:
+def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tuple[wp.vec4, mat43, wp.vec3]:
   """Core contact geometry calculation for plane-convex collision.
 
   Args:
@@ -185,7 +180,7 @@ def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tup
   _HUGE_VAL = 1e6
 
   contact_dist = wp.vec4(wp.inf)
-  contact_pos = mat43f()
+  contact_pos = mat43()
   contact_count = int(0)
 
   # get points in the convex frame

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py

@@ -26,7 +26,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
-from mujoco.mjx.third_party.mujoco_warp._src.types import vec8f
+from mujoco.mjx.third_party.mujoco_warp._src.types import vec8
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec8i
 from mujoco.mjx.third_party.mujoco_warp._src.util_misc import halton
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
@@ -54,7 +54,7 @@ class VolumeData:
   half_size: wp.vec3
   oct_aabb: wp.array2d(dtype=wp.vec3)
   oct_child: wp.array(dtype=vec8i)
-  oct_coeff: wp.array(dtype=vec8f)
+  oct_coeff: wp.array(dtype=vec8)
   valid: bool = False
 
 
@@ -79,7 +79,7 @@ def get_sdf_params(
   # Model:
   oct_child: wp.array(dtype=vec8i),
   oct_aabb: wp.array2d(dtype=wp.vec3),
-  oct_coeff: wp.array(dtype=vec8f),
+  oct_coeff: wp.array(dtype=vec8),
   plugin: wp.array(dtype=int),
   plugin_attr: wp.array(dtype=wp.vec3f),
   # In:
@@ -225,12 +225,12 @@ def user_sdf_grad(p: wp.vec3, attr: wp.vec3, sdf_type: int) -> wp.vec3:
 @wp.func
 def find_oct(
   oct_child: wp.array(dtype=vec8i), oct_aabb: wp.array2d(dtype=wp.vec3), p: wp.vec3, grad: bool
-) -> Tuple[int, Tuple[vec8f, vec8f, vec8f]]:
+) -> Tuple[int, Tuple[vec8, vec8, vec8]]:
   stack = int(0)
   niter = int(100)
-  rx = vec8f(0.0)
+  rx = vec8(0.0)
-  ry = vec8f(0.0)
+  ry = vec8(0.0)
-  rz = vec8f(0.0)
+  rz = vec8(0.0)
   eps = 1e-6
 
   while niter > 0:
@@ -621,7 +621,7 @@ def _sdf_narrowphase(
   nmeshface: int,
   oct_child: wp.array(dtype=vec8i),
   oct_aabb: wp.array2d(dtype=wp.vec3),
-  oct_coeff: wp.array(dtype=vec8f),
+  oct_coeff: wp.array(dtype=vec8),
   geom_type: wp.array(dtype=int),
   geom_condim: wp.array(dtype=int),
   geom_dataid: wp.array(dtype=int),

mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py

@@ -1843,7 +1843,7 @@ def make_constraint(m: types.Model, d: types.Data):
       if m.opt.cone == types.ConeType.PYRAMIDAL:
         wp.launch(
           _efc_contact_pyramidal,
-          dim=(d.naconmax, 2 * (m.condim_max - 1) if m.condim_max > 1 else 1),
+          dim=(d.naconmax, m.nmaxpyramid),
           inputs=[
             m.nv,
             m.opt.timestep,
@@ -1888,7 +1888,7 @@ def make_constraint(m: types.Model, d: types.Data):
       elif m.opt.cone == types.ConeType.ELLIPTIC:
         wp.launch(
           _efc_contact_elliptic,
-          dim=(d.naconmax, m.condim_max),
+          dim=(d.naconmax, m.nmaxcondim),
           inputs=[
             m.nv,
             m.opt.timestep,

mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py

@@ -75,14 +75,21 @@ def _qderiv_actuator_passive(
       else:
         bias = 0.0
 
-      if actuator_dyntype[actid] != DynType.NONE:
+      if bias == 0.0 and gain == 0.0:
-        act_first = actuator_actadr[actid]
+        continue
-        act_last = act_first + actuator_actnum[actid] - 1
-        vel = bias + gain * act_in[worldid, act_last]
-      else:
-        vel = bias + gain * ctrl_in[worldid, actid]
 
-      qderiv += actuator_moment_in[worldid, actid, dofiid] * actuator_moment_in[worldid, actid, dofjid] * vel
+      vel = bias
+      if actuator_dyntype[actid] != DynType.NONE:
+        if gain != 0.0:
+          act_first = actuator_actadr[actid]
+          act_last = act_first + actuator_actnum[actid] - 1
+          vel += gain * act_in[worldid, act_last]
+      else:
+        if gain != 0.0:
+          vel += gain * ctrl_in[worldid, actid]
+
+      if vel != 0.0:
+        qderiv += actuator_moment_in[worldid, actid, dofiid] * actuator_moment_in[worldid, actid, dofjid] * vel
 
   # TODO(team): fluid model derivative
 

mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py

@@ -42,7 +42,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec10f
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
-from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel
+from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
 
 wp.set_module_options({"enable_backward": False})
 
@@ -521,7 +521,8 @@ def fwd_position(m: Model, d: Data, factorize: bool = True):
 
 
 # TODO(team): sparse actuator_moment version
-def _actuator_velocity(m: Model, d: Data):
+@cache_kernel
+def _actuator_velocity(nv: int):
   @nested_kernel(module="unique", enable_backward=False)
   def actuator_velocity(
     # Data in:
@@ -531,22 +532,17 @@ def _actuator_velocity(m: Model, d: Data):
     actuator_velocity_out: wp.array2d(dtype=float),
   ):
     worldid, actid = wp.tid()
-    moment_tile = wp.tile_load(actuator_moment_in[worldid, actid], shape=wp.static(m.nv))
+    moment_tile = wp.tile_load(actuator_moment_in[worldid, actid], shape=wp.static(nv))
-    qvel_tile = wp.tile_load(qvel_in[worldid], shape=wp.static(m.nv))
+    qvel_tile = wp.tile_load(qvel_in[worldid], shape=wp.static(nv))
     moment_qvel_tile = wp.tile_map(wp.mul, moment_tile, qvel_tile)
     actuator_velocity_tile = wp.tile_reduce(wp.add, moment_qvel_tile)
     actuator_velocity_out[worldid, actid] = actuator_velocity_tile[0]
 
-  wp.launch_tiled(
+  return actuator_velocity
-    actuator_velocity,
-    dim=(d.nworld, m.nu),
-    inputs=[d.qvel, d.actuator_moment],
-    outputs=[d.actuator_velocity],
-    block_dim=m.block_dim.actuator_velocity,
-  )
 
 
-def _tendon_velocity(m: Model, d: Data):
+@cache_kernel
+def _tendon_velocity(nv: int):
   @nested_kernel(module="unique", enable_backward=False)
   def tendon_velocity(
     # Data in:
@@ -556,29 +552,34 @@ def _tendon_velocity(m: Model, d: Data):
     ten_velocity_out: wp.array2d(dtype=float),
   ):
     worldid, tenid = wp.tid()
-    ten_J_tile = wp.tile_load(ten_J_in[worldid, tenid], shape=wp.static(m.nv))
+    ten_J_tile = wp.tile_load(ten_J_in[worldid, tenid], shape=wp.static(nv))
-    qvel_tile = wp.tile_load(qvel_in[worldid], shape=wp.static(m.nv))
+    qvel_tile = wp.tile_load(qvel_in[worldid], shape=wp.static(nv))
     ten_J_qvel_tile = wp.tile_map(wp.mul, ten_J_tile, qvel_tile)
     ten_velocity_tile = wp.tile_reduce(wp.add, ten_J_qvel_tile)
     ten_velocity_out[worldid, tenid] = ten_velocity_tile[0]
 
-  wp.launch_tiled(
+  return tendon_velocity
-    tendon_velocity,
-    dim=(d.nworld, m.ntendon),
-    inputs=[d.qvel, d.ten_J],
-    outputs=[d.ten_velocity],
-    block_dim=m.block_dim.tendon_velocity,
-  )
 
 
 @event_scope
 def fwd_velocity(m: Model, d: Data):
   """Velocity-dependent computations."""
-  _actuator_velocity(m, d)
+  wp.launch_tiled(
+    _actuator_velocity(m.nv),
+    dim=(d.nworld, m.nu),
+    inputs=[d.qvel, d.actuator_moment],
+    outputs=[d.actuator_velocity],
+    block_dim=m.block_dim.actuator_velocity,
+  )
 
-  if m.ntendon > 0:
+  # TODO(team): sparse version
-    # TODO(team): sparse version
+  wp.launch_tiled(
-    _tendon_velocity(m, d)
+    _tendon_velocity(m.nv),
+    dim=(d.nworld, m.ntendon),
+    inputs=[d.qvel, d.ten_J],
+    outputs=[d.ten_velocity],
+    block_dim=m.block_dim.tendon_velocity,
+  )
 
   smooth.com_vel(m, d)
   passive.passive(m, d)

mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py

@@ -561,7 +561,7 @@ def _flex_elasticity(
   flex_elem: wp.array(dtype=int),
   flex_elemedge: wp.array(dtype=int),
   flexedge_length0: wp.array(dtype=float),
-  flex_stiffness: wp.array(dtype=float),
+  flex_stiffness: wp.array2d(dtype=float),
   flex_damping: wp.array(dtype=float),
   # Data in:
   flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
@@ -612,8 +612,8 @@ def _flex_elasticity(
   id = int(0)
   for ed1 in range(nedge):
     for ed2 in range(ed1, nedge):
-      metric[ed1, ed2] = flex_stiffness[21 * elemid + id]
+      metric[ed1, ed2] = flex_stiffness[elemid, id]
-      metric[ed2, ed1] = flex_stiffness[21 * elemid + id]
+      metric[ed2, ed1] = flex_stiffness[elemid, id]
       id += 1
 
   force = wp.mat(0.0, shape=(6, 3))
@@ -640,7 +640,7 @@ def _flex_bending(
   flex_vertbodyid: wp.array(dtype=int),
   flex_edge: wp.array(dtype=wp.vec2i),
   flex_edgeflap: wp.array(dtype=wp.vec2i),
-  flex_bending: wp.array(dtype=float),
+  flex_bending: wp.array2d(dtype=float),
   # Data in:
   flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
   # Data out:
@@ -664,7 +664,7 @@ def _flex_bending(
     return
 
   frc = wp.mat(0.0, shape=(4, 3))
-  if flex_bending[17 * edgeid + 16]:
+  if flex_bending[edgeid, 16]:
     v0 = flexvert_xpos_in[worldid, v[0]]
     v1 = flexvert_xpos_in[worldid, v[1]]
     v2 = flexvert_xpos_in[worldid, v[2]]
@@ -678,8 +678,8 @@ def _flex_bending(
   for i in range(nvert):
     for x in range(3):
       for j in range(nvert):
-        force[i, x] -= flex_bending[17 * edgeid + 4 * i + j] * flexvert_xpos_in[worldid, v[j]][x]
+        force[i, x] -= flex_bending[edgeid, 4 * i + j] * flexvert_xpos_in[worldid, v[j]][x]
-    force[i, x] -= flex_bending[17 * edgeid + 16] * frc[i, x]
+    force[i, x] -= flex_bending[edgeid, 16] * frc[i, x]
 
   for i in range(nvert):
     bodyid = flex_vertbodyid[flex_vertadr[f] + v[i]]

mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py

@@ -37,12 +37,12 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import SensorType
 from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec6
-from mujoco.mjx.third_party.mujoco_warp._src.types import vec8f
+from mujoco.mjx.third_party.mujoco_warp._src.types import vec8
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec8i
 from mujoco.mjx.third_party.mujoco_warp._src.util_misc import inside_geom
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
-from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel
+from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
 
 wp.set_module_options({"enable_backward": False})
 
@@ -482,8 +482,8 @@ def _sensor_pos(
   sensor_adr: wp.array(dtype=int),
   sensor_cutoff: wp.array(dtype=float),
   sensor_pos_adr: wp.array(dtype=int),
-  sensor_collision_start_adr: wp.array(dtype=int),
   rangefinder_sensor_adr: wp.array(dtype=int),
+  sensor_collision_start_adr: wp.array(dtype=int),
   collision_sensor_adr: wp.array(dtype=int),
   # Data in:
   time_in: wp.array(dtype=float),
@@ -844,8 +844,8 @@ def sensor_pos(m: Model, d: Data):
       m.sensor_adr,
       m.sensor_cutoff,
       m.sensor_pos_adr,
-      m.sensor_collision_start_adr,
       m.rangefinder_sensor_adr,
+      m.sensor_collision_start_adr,
       m.collision_sensor_adr,
       d.time,
       d.energy,
@@ -2089,7 +2089,7 @@ def _sensor_tactile(
   body_weldid: wp.array(dtype=int),
   oct_child: wp.array(dtype=vec8i),
   oct_aabb: wp.array2d(dtype=wp.vec3),
-  oct_coeff: wp.array(dtype=vec8f),
+  oct_coeff: wp.array(dtype=vec8),
   geom_type: wp.array(dtype=int),
   geom_bodyid: wp.array(dtype=int),
   geom_size: wp.array2d(dtype=wp.vec3),

mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py

@@ -35,7 +35,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import vec10
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec11
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
-from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel
+from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
 
 wp.set_module_options({"enable_backward": False})
 

mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py

@@ -27,7 +27,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import create_blocke
 from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import create_blocked_cholesky_solve_func
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
-from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel
+from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
 
 wp.set_module_options({"enable_backward": False})
 
@@ -1448,7 +1448,7 @@ def update_gradient_JTDAJ_sparse_tiled(tile_size: int, njmax: int):
 
 
 @cache_kernel
-def update_gradient_JTDAJ_dense_tiled(nv: int, tile_size: int, njmax: int):
+def update_gradient_JTDAJ_dense_tiled(nv_padded: int, tile_size: int, njmax: int):
   if njmax < tile_size:
     tile_size = njmax
 
@@ -1473,7 +1473,7 @@ def update_gradient_JTDAJ_dense_tiled(nv: int, tile_size: int, njmax: int):
 
     nefc = nefc_in[worldid]
 
-    sum_val = wp.tile_load(qM_in[worldid], shape=(nv, nv), bounds_check=False)
+    sum_val = wp.tile_load(qM_in[worldid], shape=(nv_padded, nv_padded), bounds_check=True)
 
     # Each tile processes one output tile by looping over all constraints
     for k in range(0, njmax, TILE_SIZE_K):
@@ -1483,7 +1483,7 @@ def update_gradient_JTDAJ_dense_tiled(nv: int, tile_size: int, njmax: int):
       # AD: leaving bounds-check disabled here because I'm not entirely sure that
       # everything always hits the fast path. The padding takes care of any
       #  potential OOB accesses.
-      J_ki = wp.tile_load(efc_J_in[worldid], shape=(TILE_SIZE_K, nv), offset=(k, 0), bounds_check=False)
+      J_ki = wp.tile_load(efc_J_in[worldid], shape=(TILE_SIZE_K, nv_padded), offset=(k, 0), bounds_check=False)
       J_kj = J_ki
 
       # state check
@@ -1499,7 +1499,7 @@ def update_gradient_JTDAJ_dense_tiled(nv: int, tile_size: int, njmax: int):
       active_tile = wp.tile_map(active_check, tid_tile, threshold_tile)
       D_k = wp.tile_map(wp.mul, active_tile, D_k)
 
-      J_ki = wp.tile_map(wp.mul, wp.tile_transpose(J_ki), wp.tile_broadcast(D_k, shape=(nv, TILE_SIZE_K)))
+      J_ki = wp.tile_map(wp.mul, wp.tile_transpose(J_ki), wp.tile_broadcast(D_k, shape=(nv_padded, TILE_SIZE_K)))
 
       sum_val += wp.tile_matmul(J_ki, J_kj)
 
@@ -1742,8 +1742,9 @@ def _update_gradient(m: types.Model, d: types.Data):
         outputs=[d.efc.h],
       )
     else:
+      nv_padded = d.efc.J.shape[2]
       wp.launch_tiled(
-        update_gradient_JTDAJ_dense_tiled(m.nv, types.TILE_SIZE_JTDAJ_DENSE, d.njmax),
+        update_gradient_JTDAJ_dense_tiled(nv_padded, types.TILE_SIZE_JTDAJ_DENSE, d.njmax),
         dim=d.nworld,
         inputs=[
           d.nefc,

mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py

@@ -27,7 +27,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import TileSet
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
-from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel
+from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
 
 wp.set_module_options({"enable_backward": False})
 

mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py

@@ -121,9 +121,9 @@ def event_scope(fn, name: str = ""):
   return wrapper
 
 
-# @kernel decorator to automatically set up modules based on nested
+# @nested_kernel decorator to automatically set up modules based on nested
 # function names
-def kernel(
+def nested_kernel(
   f: Optional[Callable] = None,
   *,
   enable_backward: Optional[bool] = None,
@@ -136,20 +136,20 @@ def kernel(
 
   Example::
 
-      @kernel
+      @nested_kernel
       def my_kernel(a: wp.array(dtype=float), b: wp.array(dtype=float)):
         tid = wp.tid()
         b[tid] = a[tid] + 1.0
 
 
-      @kernel(enable_backward=False)
+      @nested_kernel(enable_backward=False)
       def my_kernel_no_backward(a: wp.array(dtype=float, ndim=2), x: float):
         # the backward pass will not be generated
         i, j = wp.tid()
         a[i, j] = x
 
 
-      @kernel(module="unique")
+      @nested_kernel(module="unique")
       def my_kernel_unique_module(a: wp.array(dtype=float), b: wp.array(dtype=float)):
         # the kernel will be registered in new unique module created just for this
         # kernel and its dependent functions and structs
@@ -157,7 +157,7 @@ def kernel(
         b[tid] = a[tid] + 1.0
 
 
-      @kernel(enable_backward=False, module=None)
+      @neste_kernel(enable_backward=False, module=None)
       def my_kernel_with_args(a: wp.array(dtype=float), b: wp.array(dtype=float)):
         # can now use arguments even when module=None
         tid = wp.tid()

mjx/mujoco/mjx/warp/collision_driver.py

@@ -94,7 +94,7 @@ def _collision_shim(
     nxn_pairid_filtered: wp.array(dtype=wp.vec2i),
     oct_aabb: wp.array2d(dtype=wp.vec3),
     oct_child: wp.array(dtype=mjwp_types.vec8i),
-    oct_coeff: wp.array(dtype=mjwp_types.vec8f),
+    oct_coeff: wp.array(dtype=mjwp_types.vec8),
     pair_dim: wp.array(dtype=int),
     pair_friction: wp.array2d(dtype=mjwp_types.vec5),
     pair_gap: wp.array2d(dtype=float),

mjx/mujoco/mjx/warp/forward.py

@@ -106,7 +106,6 @@ def _forward_shim(
     cam_sensorsize: wp.array(dtype=wp.vec2),
     cam_targetbodyid: wp.array(dtype=int),
     collision_sensor_adr: wp.array(dtype=int),
-    condim_max: int,
     dof_Madr: wp.array(dtype=int),
     dof_armature: wp.array2d(dtype=float),
     dof_bodyid: wp.array(dtype=int),
@@ -129,7 +128,7 @@ def _forward_shim(
     eq_solref: wp.array2d(dtype=wp.vec2),
     eq_ten_adr: wp.array(dtype=int),
     eq_wld_adr: wp.array(dtype=int),
-    flex_bending: wp.array(dtype=float),
+    flex_bending: wp.array2d(dtype=float),
     flex_damping: wp.array(dtype=float),
     flex_dim: wp.array(dtype=int),
     flex_edge: wp.array(dtype=wp.vec2i),
@@ -138,7 +137,7 @@ def _forward_shim(
     flex_elem: wp.array(dtype=int),
     flex_elemedge: wp.array(dtype=int),
     flex_elemedgeadr: wp.array(dtype=int),
-    flex_stiffness: wp.array(dtype=float),
+    flex_stiffness: wp.array2d(dtype=float),
     flex_vertadr: wp.array(dtype=int),
     flex_vertbodyid: wp.array(dtype=int),
     flexedge_length0: wp.array(dtype=float),
@@ -229,8 +228,10 @@ def _forward_shim(
     ngravcomp: int,
     njnt: int,
     nlight: int,
+    nmaxcondim: int,
     nmaxmeshdeg: int,
     nmaxpolygon: int,
+    nmaxpyramid: int,
     nmeshface: int,
     nmocap: int,
     nrangefinder: int,
@@ -247,7 +248,7 @@ def _forward_shim(
     nxn_pairid_filtered: wp.array(dtype=wp.vec2i),
     oct_aabb: wp.array2d(dtype=wp.vec3),
     oct_child: wp.array(dtype=mjwp_types.vec8i),
-    oct_coeff: wp.array(dtype=mjwp_types.vec8f),
+    oct_coeff: wp.array(dtype=mjwp_types.vec8),
     pair_dim: wp.array(dtype=int),
     pair_friction: wp.array2d(dtype=mjwp_types.vec5),
     pair_gap: wp.array2d(dtype=float),
@@ -556,7 +557,6 @@ def _forward_shim(
   _m.cam_sensorsize = cam_sensorsize
   _m.cam_targetbodyid = cam_targetbodyid
   _m.collision_sensor_adr = collision_sensor_adr
-  _m.condim_max = condim_max
   _m.dof_Madr = dof_Madr
   _m.dof_armature = dof_armature
   _m.dof_bodyid = dof_bodyid
@@ -679,8 +679,10 @@ def _forward_shim(
   _m.ngravcomp = ngravcomp
   _m.njnt = njnt
   _m.nlight = nlight
+  _m.nmaxcondim = nmaxcondim
   _m.nmaxmeshdeg = nmaxmeshdeg
   _m.nmaxpolygon = nmaxpolygon
+  _m.nmaxpyramid = nmaxpyramid
   _m.nmeshface = nmeshface
   _m.nmocap = nmocap
   _m.nrangefinder = nrangefinder
@@ -1279,7 +1281,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m.cam_sensorsize,
       m.cam_targetbodyid,
       m._impl.collision_sensor_adr,
-      m._impl.condim_max,
       m.dof_Madr,
       m.dof_armature,
       m.dof_bodyid,
@@ -1402,8 +1403,10 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m.ngravcomp,
       m.njnt,
       m.nlight,
+      m._impl.nmaxcondim,
       m._impl.nmaxmeshdeg,
       m._impl.nmaxpolygon,
+      m._impl.nmaxpyramid,
       m.nmeshface,
       m.nmocap,
       m._impl.nrangefinder,
@@ -1895,7 +1898,6 @@ def _step_shim(
     cam_sensorsize: wp.array(dtype=wp.vec2),
     cam_targetbodyid: wp.array(dtype=int),
     collision_sensor_adr: wp.array(dtype=int),
-    condim_max: int,
     dof_Madr: wp.array(dtype=int),
     dof_armature: wp.array2d(dtype=float),
     dof_bodyid: wp.array(dtype=int),
@@ -1918,7 +1920,7 @@ def _step_shim(
     eq_solref: wp.array2d(dtype=wp.vec2),
     eq_ten_adr: wp.array(dtype=int),
     eq_wld_adr: wp.array(dtype=int),
-    flex_bending: wp.array(dtype=float),
+    flex_bending: wp.array2d(dtype=float),
     flex_damping: wp.array(dtype=float),
     flex_dim: wp.array(dtype=int),
     flex_edge: wp.array(dtype=wp.vec2i),
@@ -1927,7 +1929,7 @@ def _step_shim(
     flex_elem: wp.array(dtype=int),
     flex_elemedge: wp.array(dtype=int),
     flex_elemedgeadr: wp.array(dtype=int),
-    flex_stiffness: wp.array(dtype=float),
+    flex_stiffness: wp.array2d(dtype=float),
     flex_vertadr: wp.array(dtype=int),
     flex_vertbodyid: wp.array(dtype=int),
     flexedge_length0: wp.array(dtype=float),
@@ -2019,8 +2021,10 @@ def _step_shim(
     ngravcomp: int,
     njnt: int,
     nlight: int,
+    nmaxcondim: int,
     nmaxmeshdeg: int,
     nmaxpolygon: int,
+    nmaxpyramid: int,
     nmeshface: int,
     nmocap: int,
     nrangefinder: int,
@@ -2037,7 +2041,7 @@ def _step_shim(
     nxn_pairid_filtered: wp.array(dtype=wp.vec2i),
     oct_aabb: wp.array2d(dtype=wp.vec3),
     oct_child: wp.array(dtype=mjwp_types.vec8i),
-    oct_coeff: wp.array(dtype=mjwp_types.vec8f),
+    oct_coeff: wp.array(dtype=mjwp_types.vec8),
     pair_dim: wp.array(dtype=int),
     pair_friction: wp.array2d(dtype=mjwp_types.vec5),
     pair_gap: wp.array2d(dtype=float),
@@ -2347,7 +2351,6 @@ def _step_shim(
   _m.cam_sensorsize = cam_sensorsize
   _m.cam_targetbodyid = cam_targetbodyid
   _m.collision_sensor_adr = collision_sensor_adr
-  _m.condim_max = condim_max
   _m.dof_Madr = dof_Madr
   _m.dof_armature = dof_armature
   _m.dof_bodyid = dof_bodyid
@@ -2471,8 +2474,10 @@ def _step_shim(
   _m.ngravcomp = ngravcomp
   _m.njnt = njnt
   _m.nlight = nlight
+  _m.nmaxcondim = nmaxcondim
   _m.nmaxmeshdeg = nmaxmeshdeg
   _m.nmaxpolygon = nmaxpolygon
+  _m.nmaxpyramid = nmaxpyramid
   _m.nmeshface = nmeshface
   _m.nmocap = nmocap
   _m.nrangefinder = nrangefinder
@@ -3072,7 +3077,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.cam_sensorsize,
       m.cam_targetbodyid,
       m._impl.collision_sensor_adr,
-      m._impl.condim_max,
       m.dof_Madr,
       m.dof_armature,
       m.dof_bodyid,
@@ -3196,8 +3200,10 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.ngravcomp,
       m.njnt,
       m.nlight,
+      m._impl.nmaxcondim,
       m._impl.nmaxmeshdeg,
       m._impl.nmaxpolygon,
+      m._impl.nmaxpyramid,
       m.nmeshface,
       m.nmocap,
       m._impl.nrangefinder,

mjx/mujoco/mjx/warp/forward_test.py

@@ -155,7 +155,8 @@ class ForwardTest(parameterized.TestCase):
         mujoco.mj_fullM(m, qm, d.qM)
       else:
         qm = d.qM
-      tu.assert_eq(qm, dx._impl.qM, 'qM')
+      # mjwarp adds padding to qM
+      tu.assert_eq(qm, dx._impl.qM[: m.nv, : m.nv], 'qM')
       # qLD is fused in a cholesky factorize and solve, and not written to.
 
       tu.assert_contact_eq(d, dx, worldid=i)

mjx/mujoco/mjx/warp/types.py

@@ -105,14 +105,11 @@ class ModelWarp(PyTreeNode):
   M_colind: np.ndarray
   M_rowadr: np.ndarray
   M_rownnz: np.ndarray
-  actuator_moment_tiles_nu: Tuple[TileSet, ...]
-  actuator_moment_tiles_nv: Tuple[TileSet, ...]
   actuator_trntype_body_adr: np.ndarray
   block_dim: BlockDim
   body_fluid_ellipsoid: np.ndarray
   body_tree: Tuple[np.ndarray, ...]
   collision_sensor_adr: np.ndarray
-  condim_max: int
   dof_tri_col: np.ndarray
   dof_tri_row: np.ndarray
   eq_connect_adr: np.ndarray
@@ -158,9 +155,14 @@ class ModelWarp(PyTreeNode):
   nflexedge: int
   nflexelem: int
   nflexelemdata: int
+  nflexelemedge: int
   nflexvert: int
+  nmaxcondim: int
   nmaxmeshdeg: int
   nmaxpolygon: int
+  nmaxpyramid: int
+  noct: int
+  nplugin: int
   nrangefinder: int
   nsensorcollision: int
   nsensorcontact: int
@@ -522,8 +524,6 @@ _NDIM = {
         'actuator_gaintype': 1,
         'actuator_gear': 3,
         'actuator_lengthrange': 2,
-        'actuator_moment_tiles_nu': -1,
-        'actuator_moment_tiles_nv': -1,
         'actuator_trnid': 2,
         'actuator_trntype': 1,
         'actuator_trntype_body_adr': 1,
@@ -577,7 +577,6 @@ _NDIM = {
         'cam_sensorsize': 2,
         'cam_targetbodyid': 1,
         'collision_sensor_adr': 1,
-        'condim_max': 0,
         'dof_Madr': 1,
         'dof_armature': 2,
         'dof_bodyid': 1,
@@ -603,7 +602,7 @@ _NDIM = {
         'eq_type': 1,
         'eq_wld_adr': 1,
         'exclude_signature': 1,
-        'flex_bending': 1,
+        'flex_bending': 2,
         'flex_damping': 1,
         'flex_dim': 1,
         'flex_edge': 2,
@@ -612,7 +611,7 @@ _NDIM = {
         'flex_elem': 1,
         'flex_elemedge': 1,
         'flex_elemedgeadr': 1,
-        'flex_stiffness': 1,
+        'flex_stiffness': 2,
         'flex_vertadr': 1,
         'flex_vertbodyid': 1,
         'flex_vertnum': 1,
@@ -677,7 +676,7 @@ _NDIM = {
         'light_type': 2,
         'mapM2M': 1,
         'mat_rgba': 3,
-        'mat_texid': 3,
+        'mat_texid': 2,
         'mat_texrepeat': 3,
         'mesh_face': 2,
         'mesh_faceadr': 1,
@@ -711,6 +710,7 @@ _NDIM = {
         'nflexedge': 0,
         'nflexelem': 0,
         'nflexelemdata': 0,
+        'nflexelemedge': 0,
         'nflexvert': 0,
         'ngeom': 0,
         'ngravcomp': 0,
@@ -719,16 +719,22 @@ _NDIM = {
         'njnt': 0,
         'nlight': 0,
         'nmat': 0,
+        'nmaxcondim': 0,
         'nmaxmeshdeg': 0,
         'nmaxpolygon': 0,
+        'nmaxpyramid': 0,
+        'nmesh': 0,
         'nmeshface': 0,
         'nmeshgraph': 0,
+        'nmeshnormal': 0,
         'nmeshpoly': 0,
         'nmeshpolymap': 0,
         'nmeshpolyvert': 0,
         'nmeshvert': 0,
         'nmocap': 0,
+        'noct': 0,
         'npair': 0,
+        'nplugin': 0,
         'nq': 0,
         'nrangefinder': 0,
         'nsensor': 0,
@@ -748,8 +754,6 @@ _NDIM = {
         'oct_aabb': 3,
         'oct_child': 2,
         'oct_coeff': 2,
-        'opt__broadphase': 0,
-        'opt__broadphase_filter': 0,
         'opt__ccd_iterations': 0,
         'opt__ccd_tolerance': 1,
         'opt__cone': 0,
@@ -867,8 +871,6 @@ _NDIM = {
         'wrap_type': 1,
     },
     'Option': {
-        'broadphase': 0,
-        'broadphase_filter': 0,
         'ccd_iterations': 0,
         'ccd_tolerance': 1,
         'cone': 0,
@@ -1060,8 +1062,6 @@ _BATCH_DIM = {
         'actuator_gaintype': False,
         'actuator_gear': True,
         'actuator_lengthrange': False,
-        'actuator_moment_tiles_nu': False,
-        'actuator_moment_tiles_nv': False,
         'actuator_trnid': False,
         'actuator_trntype': False,
         'actuator_trntype_body_adr': False,
@@ -1115,7 +1115,6 @@ _BATCH_DIM = {
         'cam_sensorsize': False,
         'cam_targetbodyid': False,
         'collision_sensor_adr': False,
-        'condim_max': False,
         'dof_Madr': False,
         'dof_armature': True,
         'dof_bodyid': False,
@@ -1215,7 +1214,7 @@ _BATCH_DIM = {
         'light_type': True,
         'mapM2M': False,
         'mat_rgba': True,
-        'mat_texid': True,
+        'mat_texid': False,
         'mat_texrepeat': True,
         'mesh_face': False,
         'mesh_faceadr': False,
@@ -1249,6 +1248,7 @@ _BATCH_DIM = {
         'nflexedge': False,
         'nflexelem': False,
         'nflexelemdata': False,
+        'nflexelemedge': False,
         'nflexvert': False,
         'ngeom': False,
         'ngravcomp': False,
@@ -1257,16 +1257,22 @@ _BATCH_DIM = {
         'njnt': False,
         'nlight': False,
         'nmat': False,
+        'nmaxcondim': False,
         'nmaxmeshdeg': False,
         'nmaxpolygon': False,
+        'nmaxpyramid': False,
+        'nmesh': False,
         'nmeshface': False,
         'nmeshgraph': False,
+        'nmeshnormal': False,
         'nmeshpoly': False,
         'nmeshpolymap': False,
         'nmeshpolyvert': False,
         'nmeshvert': False,
         'nmocap': False,
+        'noct': False,
         'npair': False,
+        'nplugin': False,
         'nq': False,
         'nrangefinder': False,
         'nsensor': False,
@@ -1286,8 +1292,6 @@ _BATCH_DIM = {
         'oct_aabb': False,
         'oct_child': False,
         'oct_coeff': False,
-        'opt__broadphase': False,
-        'opt__broadphase_filter': False,
         'opt__ccd_iterations': False,
         'opt__ccd_tolerance': True,
         'opt__cone': False,
@@ -1405,8 +1409,6 @@ _BATCH_DIM = {
         'wrap_type': False,
     },
     'Option': {
-        'broadphase': False,
-        'broadphase_filter': False,
         'ccd_iterations': False,
         'ccd_tolerance': True,
         'cone': False,