Import google-deepmind/mujoco_warp from GitHub.

PiperOrigin-RevId: 857298931
Change-Id: Icbaddaea8140198c3673124d5994a01d3c92ed69

mjx/mujoco/mjx/_src/io.py

@@ -178,7 +178,7 @@ def _wp_to_np_type(wp_field: Any, name: str = '') -> Any:
     return wp.types.warp_type_to_np_dtype[wp_dtype](wp_field)
 
   # warp arrays
-  if isinstance(wp_field, wp.types.array):
+  if isinstance(wp_field, wp.array):
     return wp_field.numpy()
 
   # static
@@ -191,7 +191,7 @@ def _wp_to_np_type(wp_field: Any, name: str = '') -> Any:
   # tuples
   if isinstance(wp_field, tuple) and len(wp_field) == 0:
     return ()
-  if isinstance(wp_field, tuple) and isinstance(wp_field[0], wp.types.array):
+  if isinstance(wp_field, tuple) and isinstance(wp_field[0], wp.array):
     return tuple(f.numpy() for f in wp_field)
   if isinstance(wp_field, tuple) and isinstance(
       wp_field[0], mjwp_types.TileSet

mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py

@@ -15,6 +15,13 @@
 
 """Public API for MJWarp."""
 
+from importlib import metadata
+
+try:
+  __version__ = metadata.version("mujoco_warp")
+except metadata.PackageNotFoundError:
+  __version__ = "unknown"
+
 # isort: off
 from mujoco.mjx.third_party.mujoco_warp._src.forward import step as step
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model as Model

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

@@ -15,17 +15,12 @@
 
 """Utilities for benchmarking MuJoCo Warp."""
 
-import importlib
-import os
 import time
 from typing import Callable, Optional, Tuple
 
-import mujoco
 import numpy as np
 import warp as wp
 
-from mujoco.mjx.third_party.mujoco_warp._src import forward
-from mujoco.mjx.third_party.mujoco_warp._src import io
 from mujoco.mjx.third_party.mujoco_warp._src import warp_util
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
@@ -165,100 +160,3 @@ def benchmark(
     run_duration = np.sum(time_vec)
 
   return jit_duration, run_duration, trace, nacon, nefc, solver_niter, nsuccess
-
-
-class BenchmarkSuite:
-  """Base suite for all model benchmarks."""
-
-  path = ""
-  batch_size = -1
-  nconmax = -1
-  njmax = -1
-  nstep = 1000
-  param_names = ("function",)
-  params = (
-    "jit_duration",
-    "solver_niter_mean",
-    "solver_niter_p95",
-    "device_memory_allocated",
-    "step",
-    "step.forward",
-    "step.forward.fwd_position",
-    "step.forward.fwd_position.kinematics",
-    "step.forward.fwd_position.com_pos",
-    "step.forward.fwd_position.camlight",
-    "step.forward.fwd_position.crb",
-    "step.forward.fwd_position.tendon_armature",
-    "step.forward.fwd_position.collision",
-    "step.forward.fwd_position.make_constraint",
-    "step.forward.fwd_position.transmission",
-    "step.forward.sensor_pos",
-    "step.forward.fwd_velocity",
-    "step.forward.fwd_velocity.com_vel",
-    "step.forward.fwd_velocity.passive",
-    "step.forward.fwd_velocity.rne",
-    "step.forward.fwd_velocity.tendon_bias",
-    "step.forward.sensor_vel",
-    "step.forward.fwd_actuation",
-    "step.forward.fwd_acceleration",
-    "step.forward.fwd_acceleration.xfrc_accumulate",
-    "step.forward.sensor_acc",
-    "step.forward.solve",
-  )
-  number = 1
-  rounds = 1
-  sample_time = 0
-  repeat = 1
-  replay = ""
-
-  def setup_cache(self):
-    module = importlib.import_module(self.__module__)
-    path = os.path.join(os.path.realpath(os.path.dirname(module.__file__)), self.path)
-    mjm = mujoco.MjModel.from_xml_path(path)
-    mjd = mujoco.MjData(mjm)
-    ctrls = None
-
-    if self.replay:
-      keys = io.find_keys(mjm, self.replay)
-      if not keys:
-        raise ValueError(f"Key prefix not find: {self.replay}")
-      ctrls = io.make_trajectory(mjm, keys)
-      mujoco.mj_resetDataKeyframe(mjm, mjd, keys[0])
-
-    if mjm.nkey > 0:
-      mujoco.mj_resetDataKeyframe(mjm, mjd, 0)
-
-    # TODO(team): mj_forward call shouldn't be necessary, but it is
-    mujoco.mj_forward(mjm, mjd)
-
-    wp.init()
-    if os.environ.get("ASV_CACHE_KERNELS", "false").lower() == "false":
-      wp.clear_kernel_cache()
-
-    free_before = wp.get_device().free_memory
-    m = io.put_model(mjm)
-    d = io.put_data(mjm, mjd, self.batch_size, self.nconmax, self.njmax)
-    free_after = wp.get_device().free_memory
-
-    jit_duration, _, trace, _, _, solver_niter, _ = benchmark(forward.step, m, d, self.nstep, ctrls, True, False, True)
-    metrics = {
-      "jit_duration": jit_duration,
-      "solver_niter_mean": np.mean(solver_niter),
-      "solver_niter_p95": np.quantile(solver_niter, 0.95),
-      "device_memory_allocated": free_before - free_after,
-    }
-
-    def tree_flatten(d, parent_k=""):
-      ret = {}
-      steps = self.batch_size * 1000
-      for k, v in d.items():
-        k = parent_k + "." + k if parent_k else k
-        ret = ret | {k: 1e6 * v[0][0] / steps} | tree_flatten(v[1], k)
-      return ret
-
-    metrics = metrics | tree_flatten(trace)
-
-    return metrics
-
-  def track_metric(self, metrics, fn):
-    return metrics[fn]

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

@@ -32,9 +32,6 @@ def create_blocked_cholesky_func(block_size: int):
 
     It returns a lower-triangular matrix L such that A = L L^T.
     """
-    # TODO(team): remove conditional after mjwarp relies on >= 1.11
-    bleeding_edge_warp = wp.static(tuple(map(int, wp.__version__.split(".")[:2])) >= (1, 11))
-
     # Process the matrix in blocks along its leading dimension.
     for k in range(0, matrix_size, block_size):
       end = k + block_size
@@ -45,10 +42,7 @@ def create_blocked_cholesky_func(block_size: int):
 
       for j in range(0, k, block_size):
         L_block = wp.tile_load(L, shape=(block_size, block_size), offset=(k, j), storage="shared")
-        if bleeding_edge_warp:
+        wp.tile_matmul(L_block, wp.tile_transpose(L_block), A_kk_tile, alpha=-1.0)
-          wp.tile_matmul(L_block, wp.tile_transpose(L_block), A_kk_tile, alpha=-1.0)
-        else:
-          A_kk_tile -= wp.tile_matmul(L_block, wp.tile_transpose(L_block))
 
       # Compute the Cholesky factorization for the block
       L_kk_tile = wp.tile_cholesky(A_kk_tile)
@@ -61,15 +55,9 @@ def create_blocked_cholesky_func(block_size: int):
         for j in range(0, k, block_size):
           L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, j), storage="shared")
           L_2_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(k, j), storage="shared")
-          if bleeding_edge_warp:
+          wp.tile_matmul(L_tile, wp.tile_transpose(L_2_tile), A_ik_tile, alpha=-1.0)
-            wp.tile_matmul(L_tile, wp.tile_transpose(L_2_tile), A_ik_tile, alpha=-1.0)
-          else:
-            A_ik_tile -= wp.tile_matmul(L_tile, wp.tile_transpose(L_2_tile))
 
-        if bleeding_edge_warp:
+        wp.tile_lower_solve_inplace(L_kk_tile, wp.tile_transpose(A_ik_tile))
-          wp.tile_lower_solve_inplace(L_kk_tile, wp.tile_transpose(A_ik_tile))
-        else:
-          A_ik_tile = wp.tile_transpose(wp.tile_lower_solve(L_kk_tile, wp.tile_transpose(A_ik_tile)))
         wp.tile_store(L, A_ik_tile, offset=(i, k))
 
   return blocked_cholesky_func
@@ -91,8 +79,6 @@ def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int)
     Solves A x = b given the Cholesky factor L (A = L L^T) using blocked forward and backward
     substitution.
     """
-    # TODO(team): remove conditional after mjwarp relies on >= 1.11
-    bleeding_edge_warp = wp.static(tuple(map(int, wp.__version__.split(".")[:2])) >= (1, 11))
     rhs_tile = wp.tile_load(b, shape=(matrix_size_static, 1), offset=(0, 0), storage="shared", bounds_check=False)
 
     # Forward substitution: solve L y = b
@@ -101,17 +87,10 @@ def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int)
       for j in range(0, i, block_size):
         L_block = wp.tile_load(L, shape=(block_size, block_size), offset=(i, j), storage="shared")
         y_block = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
-        if bleeding_edge_warp:
+        wp.tile_matmul(L_block, y_block, rhs_view, alpha=-1.0)
-          wp.tile_matmul(L_block, y_block, rhs_view, alpha=-1.0)
-        else:
-          rhs_view -= wp.tile_matmul(L_block, y_block)
 
       L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, i), storage="shared")
-      if bleeding_edge_warp:
+      wp.tile_lower_solve_inplace(L_tile, rhs_view)
-        wp.tile_lower_solve_inplace(L_tile, rhs_view)
-      else:
-        rhs_tmp = wp.tile_lower_solve(L_tile, rhs_view)
-        wp.tile_assign(rhs_tile, rhs_tmp, offset=(i, 0))
 
     # Backward substitution: solve L^T x = y
     for i in range(matrix_size - block_size, -1, -block_size):
@@ -120,16 +99,10 @@ def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int)
       for j in range(i_end, matrix_size, block_size):
         L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(j, i), storage="shared")
         x_tile = wp.tile_load(x, shape=(block_size, 1), offset=(j, 0), storage="shared", bounds_check=False)
-        if bleeding_edge_warp:
+        wp.tile_matmul(wp.tile_transpose(L_tile), x_tile, tmp_tile, alpha=-1.0)
-          wp.tile_matmul(wp.tile_transpose(L_tile), x_tile, tmp_tile, alpha=-1.0)
-        else:
-          tmp_tile -= wp.tile_matmul(wp.tile_transpose(L_tile), x_tile)
       L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, i), storage="shared")
 
-      if bleeding_edge_warp:
+      wp.tile_upper_solve_inplace(wp.tile_transpose(L_tile), tmp_tile)
-        wp.tile_upper_solve_inplace(wp.tile_transpose(L_tile), tmp_tile)
-      else:
-        tmp_tile = wp.tile_upper_solve(wp.tile_transpose(L_tile), tmp_tile)
       wp.tile_store(x, tmp_tile, offset=(i, 0), bounds_check=False)
 
   return blocked_cholesky_solve_func

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

@@ -19,6 +19,7 @@ import warp as wp
 
 from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import ccd
 from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import multicontact
+from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import support
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import Geom
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import contact_params
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import geom_collision_pair
@@ -28,8 +29,8 @@ from mujoco.mjx.third_party.mujoco_warp._src.math import upper_trid_index
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAFACES
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAHORIZON
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXCONPAIR
-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 EnableBit
 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
@@ -85,8 +86,9 @@ assert _check_convex_collision_pairs(), "_CONVEX_COLLISION_PAIRS is in invalid o
 @wp.func
 def _hfield_filter(
   # Model:
+  geom_type: wp.array(dtype=int),
   geom_dataid: wp.array(dtype=int),
-  geom_aabb: wp.array3d(dtype=wp.vec3),
+  geom_size: wp.array2d(dtype=wp.vec3),
   geom_rbound: wp.array2d(dtype=float),
   geom_margin: wp.array2d(dtype=float),
   hfield_size: wp.array(dtype=wp.vec4),
@@ -109,6 +111,7 @@ def _hfield_filter(
   # geom info
   rbound_id = worldid % geom_rbound.shape[0]
   margin_id = worldid % geom_margin.shape[0]
+  size_id = worldid % geom_size.shape[0]
 
   pos1 = geom_xpos_in[worldid, g1]
   mat1 = geom_xmat_in[worldid, g1]
@@ -135,47 +138,23 @@ def _hfield_filter(
   mat2 = geom_xmat_in[worldid, g2]
   mat = mat1T @ mat2
 
-  # aabb for geom in height field frame
+  # create geom in height field frame for support function queries
-  xmax = -MJ_MAXVAL
+  geom2 = Geom()
-  ymax = -MJ_MAXVAL
+  geom2.pos = pos
-  zmax = -MJ_MAXVAL
+  geom2.rot = mat
-  xmin = MJ_MAXVAL
+  geom2.size = geom_size[size_id, g2]
-  ymin = MJ_MAXVAL
+  geom2.margin = 0.0  # margin handled separately
-  zmin = MJ_MAXVAL
+  geom2.index = -1
 
-  aabb_id = worldid % geom_aabb.shape[0]
+  geomtype2 = geom_type[g2]
-  center2 = geom_aabb[aabb_id, g2, 0]
-  size2 = geom_aabb[aabb_id, g2, 1]
 
-  pos += mat1T @ center2
+  # use support functions for tight AABB bounds
-
+  xmax = support(geom2, geomtype2, wp.vec3(1.0, 0.0, 0.0)).point[0]
-  sign = wp.vec2(-1.0, 1.0)
+  xmin = support(geom2, geomtype2, wp.vec3(-1.0, 0.0, 0.0)).point[0]
-
+  ymax = support(geom2, geomtype2, wp.vec3(0.0, 1.0, 0.0)).point[1]
-  for i in range(2):
+  ymin = support(geom2, geomtype2, wp.vec3(0.0, -1.0, 0.0)).point[1]
-    for j in range(2):
+  zmax = support(geom2, geomtype2, wp.vec3(0.0, 0.0, 1.0)).point[2]
-      for k in range(2):
+  zmin = support(geom2, geomtype2, wp.vec3(0.0, 0.0, -1.0)).point[2]
-        corner_local = wp.vec3(sign[i] * size2[0], sign[j] * size2[1], sign[k] * size2[2])
-        corner_hf = mat @ corner_local
-
-        if corner_hf[0] > xmax:
-          xmax = corner_hf[0]
-        if corner_hf[1] > ymax:
-          ymax = corner_hf[1]
-        if corner_hf[2] > zmax:
-          zmax = corner_hf[2]
-        if corner_hf[0] < xmin:
-          xmin = corner_hf[0]
-        if corner_hf[1] < ymin:
-          ymin = corner_hf[1]
-        if corner_hf[2] < zmin:
-          zmin = corner_hf[2]
-
-  xmax += pos[0]
-  xmin += pos[0]
-  ymax += pos[1]
-  ymin += pos[1]
-  zmax += pos[2]
-  zmin += pos[2]
 
   # box-box test
   if (
@@ -213,7 +192,6 @@ def ccd_hfield_kernel_builder(
     geom_solref: wp.array2d(dtype=wp.vec2),
     geom_solimp: wp.array2d(dtype=vec5),
     geom_size: wp.array2d(dtype=wp.vec3),
-    geom_aabb: wp.array3d(dtype=wp.vec3),
     geom_rbound: wp.array2d(dtype=float),
     geom_friction: wp.array2d(dtype=wp.vec3),
     geom_margin: wp.array2d(dtype=float),
@@ -253,16 +231,13 @@ def ccd_hfield_kernel_builder(
     collision_worldid_in: wp.array(dtype=int),
     ncollision_in: wp.array(dtype=int),
     # In:
-    epa_vert_in: wp.array2d(dtype=wp.vec3),
     epa_vert1_in: wp.array2d(dtype=wp.vec3),
     epa_vert2_in: wp.array2d(dtype=wp.vec3),
     epa_vert_index1_in: wp.array2d(dtype=int),
     epa_vert_index2_in: wp.array2d(dtype=int),
-    epa_face_in: wp.array2d(dtype=wp.vec3i),
+    epa_face_in: wp.array2d(dtype=int),
     epa_pr_in: wp.array2d(dtype=wp.vec3),
     epa_norm2_in: wp.array2d(dtype=float),
-    epa_index_in: wp.array2d(dtype=int),
-    epa_map_in: wp.array2d(dtype=int),
     epa_horizon_in: wp.array2d(dtype=int),
     # Data out:
     nacon_out: wp.array(dtype=int),
@@ -295,7 +270,7 @@ def ccd_hfield_kernel_builder(
 
     # height field filter
     no_hf_collision, xmin, xmax, ymin, ymax, zmin, zmax = _hfield_filter(
-      geom_dataid, geom_aabb, geom_rbound, geom_margin, hfield_size, geom_xpos_in, geom_xmat_in, worldid, g1, g2
+      geom_type, geom_dataid, geom_size, geom_rbound, geom_margin, hfield_size, geom_xpos_in, geom_xmat_in, worldid, g1, g2
     )
     if no_hf_collision:
       return
@@ -346,6 +321,18 @@ def ccd_hfield_kernel_builder(
       worldid,
     )
 
+    # transform geom2 into heightfield frame
+    hf_pos = geom_xpos_in[worldid, g1]
+    hf_mat = geom_xmat_in[worldid, g1]
+    hf_matT = wp.transpose(hf_mat)
+
+    geom2.pos = hf_matT @ (geom2.pos - hf_pos)
+    geom2.rot = hf_matT @ geom2.rot
+
+    # geom1 has identity pose
+    geom1.pos = wp.vec3(0.0, 0.0, 0.0)
+    geom1.rot = wp.identity(n=3, dtype=float)
+
     # see MuJoCo mjc_ConvexHField
     geom1_dataid = geom_dataid[g1]
 
@@ -388,7 +375,6 @@ def ccd_hfield_kernel_builder(
     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]
@@ -396,8 +382,6 @@ def ccd_hfield_kernel_builder(
     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]
@@ -405,8 +389,24 @@ def ccd_hfield_kernel_builder(
     # process all prisms in subgrid
     count = int(0)
     for r in range(rmin, rmax):
-      nvert = int(0)
+      # pre-initialize first 2 vertices
-      for c in range(cmin, cmax + 1):
+      for init_i in range(2):
+        x = dx * float(cmin) - size[0]
+        y = dy * float(r + dr[init_i]) - size[1]
+        z = hfield_data[adr + (r + dr[init_i]) * ncol + cmin] * size[2] + margin
+
+        prism[0] = prism[1]
+        prism[1] = prism[2]
+        prism[3] = prism[4]
+        prism[4] = prism[5]
+
+        prism[2, 0] = x
+        prism[5, 0] = x
+        prism[2, 1] = y
+        prism[5, 1] = y
+        prism[5, 2] = z
+
+      for c in range(cmin + 1, cmax + 1):
         # add both triangles from this cell
         for i in range(2):
           if count >= MJ_MAXCONPAIR:
@@ -432,11 +432,6 @@ def ccd_hfield_kernel_builder(
           prism[5, 1] = y
           prism[5, 2] = z
 
-          nvert += 1
-
-          if nvert <= 2:
-            continue
-
           # prism height test
           if prism[3, 2] < zmin and prism[4, 2] < zmin and prism[5, 2] < zmin:
             continue
@@ -461,7 +456,6 @@ def ccd_hfield_kernel_builder(
             geomtype2,
             x1,
             geom2.pos,
-            epa_vert,
             epa_vert1,
             epa_vert2,
             epa_vert_index1,
@@ -469,8 +463,6 @@ def ccd_hfield_kernel_builder(
             epa_face,
             epa_pr,
             epa_norm2,
-            epa_index,
-            epa_map,
             epa_horizon,
           )
 
@@ -480,13 +472,16 @@ def ccd_hfield_kernel_builder(
           # cache contact information
           hfield_contact_dist[count] = dist
 
-          pos = 0.5 * (w1 + w2)
+          # transform contact to global frame
+          pos_local = 0.5 * (w1 + w2)
+          pos = hf_mat @ pos_local + hf_pos
           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)
+          frame_local = make_frame(w1 - w2)
-          normal = wp.vec3(frame[0, 0], frame[0, 1], frame[0, 2])
+          normal_local = wp.vec3(frame_local[0, 0], frame_local[0, 1], frame_local[0, 2])
+          normal = hf_mat @ normal_local
           hfield_contact_normal[count, 0] = normal[0]
           hfield_contact_normal[count, 1] = normal[1]
           hfield_contact_normal[count, 2] = normal[2]
@@ -720,16 +715,13 @@ def ccd_kernel_builder(
     # Data in:
     naconmax_in: int,
     # In:
-    epa_vert_in: wp.array2d(dtype=wp.vec3),
     epa_vert1_in: wp.array2d(dtype=wp.vec3),
     epa_vert2_in: wp.array2d(dtype=wp.vec3),
     epa_vert_index1_in: wp.array2d(dtype=int),
     epa_vert_index2_in: wp.array2d(dtype=int),
-    epa_face_in: wp.array2d(dtype=wp.vec3i),
+    epa_face_in: wp.array2d(dtype=int),
     epa_pr_in: wp.array2d(dtype=wp.vec3),
     epa_norm2_in: wp.array2d(dtype=float),
-    epa_index_in: wp.array2d(dtype=int),
-    epa_map_in: wp.array2d(dtype=int),
     epa_horizon_in: wp.array2d(dtype=int),
     multiccd_polygon_in: wp.array2d(dtype=wp.vec3),
     multiccd_clipped_in: wp.array2d(dtype=wp.vec3),
@@ -795,7 +787,6 @@ def ccd_kernel_builder(
       geomtype2,
       x1,
       x2,
-      epa_vert_in[tid],
       epa_vert1_in[tid],
       epa_vert2_in[tid],
       epa_vert_index1_in[tid],
@@ -803,8 +794,6 @@ def ccd_kernel_builder(
       epa_face_in[tid],
       epa_pr_in[tid],
       epa_norm2_in[tid],
-      epa_index_in[tid],
-      epa_map_in[tid],
       epa_horizon_in[tid],
     )
 
@@ -939,16 +928,13 @@ def ccd_kernel_builder(
     collision_worldid_in: wp.array(dtype=int),
     ncollision_in: wp.array(dtype=int),
     # In:
-    epa_vert_in: wp.array2d(dtype=wp.vec3),
     epa_vert1_in: wp.array2d(dtype=wp.vec3),
     epa_vert2_in: wp.array2d(dtype=wp.vec3),
     epa_vert_index1_in: wp.array2d(dtype=int),
     epa_vert_index2_in: wp.array2d(dtype=int),
-    epa_face_in: wp.array2d(dtype=wp.vec3i),
+    epa_face_in: wp.array2d(dtype=int),
     epa_pr_in: wp.array2d(dtype=wp.vec3),
     epa_norm2_in: wp.array2d(dtype=float),
-    epa_index_in: wp.array2d(dtype=int),
-    epa_map_in: wp.array2d(dtype=int),
     epa_horizon_in: wp.array2d(dtype=int),
     multiccd_polygon_in: wp.array2d(dtype=wp.vec3),
     multiccd_clipped_in: wp.array2d(dtype=wp.vec3),
@@ -1040,7 +1026,6 @@ def ccd_kernel_builder(
       opt_ccd_tolerance,
       geom_type,
       naconmax_in,
-      epa_vert_in,
       epa_vert1_in,
       epa_vert2_in,
       epa_vert_index1_in,
@@ -1048,8 +1033,6 @@ def ccd_kernel_builder(
       epa_face_in,
       epa_pr_in,
       epa_norm2_in,
-      epa_index_in,
-      epa_map_in,
       epa_horizon_in,
       multiccd_polygon_in,
       multiccd_clipped_in,
@@ -1153,12 +1136,10 @@ def convex_narrowphase(m: Model, d: Data):
   epa_iterations = m.opt.ccd_iterations
 
   # set to true to enable multiccd
-  use_multiccd = False
+  use_multiccd = m.opt.enableflags & EnableBit.MULTICCD
   nmaxpolygon = m.nmaxpolygon if use_multiccd else 0
   nmaxmeshdeg = m.nmaxmeshdeg if use_multiccd else 0
 
-  # epa_vert: vertices in EPA polytope in Minkowski space
-  epa_vert = wp.empty(shape=(d.naconmax, 5 + epa_iterations), dtype=wp.vec3)
   # epa_vert1: vertices in EPA polytope in geom 1 space
   epa_vert1 = wp.empty(shape=(d.naconmax, 5 + epa_iterations), dtype=wp.vec3)
   # epa_vert2: vertices in EPA polytope in geom 2 space
@@ -1168,17 +1149,13 @@ def convex_narrowphase(m: Model, d: Data):
   # epa_vert_index2: vertex indices in EPA polytope for geom 2  (naconmax, 5 + CCDiter)
   epa_vert_index2 = wp.empty(shape=(d.naconmax, 5 + epa_iterations), dtype=int)
   # epa_face: faces of polytope represented by three indices
-  epa_face = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=wp.vec3i)
+  epa_face = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=int)
   # epa_pr: projection of origin on polytope faces
   epa_pr = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=wp.vec3)
   # epa_norm2: epa_pr * epa_pr
   epa_norm2 = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=float)
-  # epa_index: index of face in polytope map
-  epa_index = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=int)
-  # epa_map: status of faces in polytope
-  epa_map = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=int)
   # epa_horizon: index pair (i j) of edges on horizon
-  epa_horizon = wp.empty(shape=(d.naconmax, 2 * MJ_MAX_EPAHORIZON), dtype=int)
+  epa_horizon = wp.empty(shape=(d.naconmax, MJ_MAX_EPAHORIZON), dtype=int)
 
   # Contact outputs
   contact_outputs = [
@@ -1216,7 +1193,6 @@ def convex_narrowphase(m: Model, d: Data):
           m.geom_solref,
           m.geom_solimp,
           m.geom_size,
-          m.geom_aabb,
           m.geom_rbound,
           m.geom_friction,
           m.geom_margin,
@@ -1254,7 +1230,6 @@ def convex_narrowphase(m: Model, d: Data):
           d.collision_pairid,
           d.collision_worldid,
           d.ncollision,
-          epa_vert,
           epa_vert1,
           epa_vert2,
           epa_vert_index1,
@@ -1262,8 +1237,6 @@ def convex_narrowphase(m: Model, d: Data):
           epa_face,
           epa_pr,
           epa_norm2,
-          epa_index,
-          epa_map,
           epa_horizon,
         ],
         outputs=contact_outputs,
@@ -1342,7 +1315,6 @@ def convex_narrowphase(m: Model, d: Data):
           d.collision_pairid,
           d.collision_worldid,
           d.ncollision,
-          epa_vert,
           epa_vert1,
           epa_vert2,
           epa_vert_index1,
@@ -1350,8 +1322,6 @@ def convex_narrowphase(m: Model, d: Data):
           epa_face,
           epa_pr,
           epa_norm2,
-          epa_index,
-          epa_map,
           epa_horizon,
           multiccd_polygon,
           multiccd_clipped,

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

@@ -34,6 +34,13 @@ MIN_DIST = 1e-10
 FACE_TOL = 0.99999872
 EDGE_TOL = 0.00159999931
 
+# Bit flags for face status in EPA polytope.
+# Defined at module scope to avoid Warp's intermediate type issues with literals.
+# See: https://github.com/NVIDIA/warp/issues/485
+_FACE_DELETED_BIT = wp.constant(wp.uint32(0x80000000))
+_FACE_INVALID_BIT = wp.constant(wp.uint32(0x40000000))
+_FACE_INVALID_OR_DELETED_MASK = wp.constant(wp.uint32(0xC0000000))
+
 
 @wp.struct
 class GJKResult:
@@ -53,7 +60,6 @@ class Polytope:
   status: int
 
   # vertices in polytope
-  vert: wp.array(dtype=wp.vec3)
   vert1: wp.array(dtype=wp.vec3)
   vert2: wp.array(dtype=wp.vec3)
   vert_index1: wp.array(dtype=int)
@@ -61,16 +67,13 @@ class Polytope:
   nvert: int
 
   # faces in polytope
-  face: wp.array(dtype=wp.vec3i)
+  # 10 bits per each vertex index, while the last significant bits are for
+  # invalid and deleted face
+  face: wp.array(dtype=int)
   face_pr: wp.array(dtype=wp.vec3)
   face_norm2: wp.array(dtype=float)
-  face_index: wp.array(dtype=int)
   nface: int
 
-  # TODO(kbayes): look into if a linear map actually improves performance
-  face_map: wp.array(dtype=int)
-  nmap: int
-
   # edges that make up the horizon when adding new vertices to polytope
   horizon: wp.array(dtype=int)
   nhorizon: int
@@ -91,12 +94,12 @@ def _discrete_geoms(g1: int, g2: int) -> bool:
 
 
 @wp.func
-def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint:
+def support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint:
   sp = SupportPoint()
   sp.cached_index = -1
   sp.vertex_index = -1
   if geomtype == GeomType.SPHERE:
-    sp.point = geom.pos + (0.5 * geom.margin) * geom.size[0] * dir
+    sp.point = geom.pos + (geom.size[0] + 0.5 * geom.margin) * dir
     return sp
 
   local_dir = wp.transpose(geom.rot) @ dir
@@ -174,14 +177,13 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint:
   elif geomtype == GeomType.HFIELD:
     max_dist = float(FLOAT_MIN)
     # TODO(kbayes): Support edge prisms
-    sp.vertex_index = wp.where(local_dir[2] < 0, -2, -3)
+    sp.vertex_index = wp.where(dir[2] < 0, -2, -3)
     for i in range(6):
       vert = geom.hfprism[i]
-      dist = wp.dot(vert, local_dir)
+      dist = wp.dot(vert, dir)
       if dist > max_dist:
         max_dist = dist
         sp.point = vert
-    sp.point = geom.rot @ sp.point + geom.pos
 
   if geom.margin > 0.0:
     sp.point += dir * (0.5 * geom.margin)
@@ -195,16 +197,15 @@ def _attach_face(pt: Polytope, idx: int, v1: int, v2: int, v3: int) -> float:
     return 0.0
 
   # compute witness point v
-  r, ret = _project_origin_plane(pt.vert[v3], pt.vert[v2], pt.vert[v1])
+  r, ret = _project_origin_plane(pt.vert1[v3] - pt.vert2[v3], pt.vert1[v2] - pt.vert2[v2], pt.vert1[v1] - pt.vert2[v1])
   if ret:
     return 0.0
 
-  face = wp.vec3i(v1, v2, v3)
+  face = v1 + (v2 << 10) + (v3 << 20)
   pt.face[idx] = face
   pt.face_pr[idx] = r
 
   pt.face_norm2[idx] = wp.dot(r, r)
-  pt.face_index[idx] = -1
   return pt.face_norm2[idx]
 
 
@@ -212,18 +213,16 @@ def _attach_face(pt: Polytope, idx: int, v1: int, v2: int, v3: int) -> float:
 def _epa_support(
   pt: Polytope, idx: int, geom1: Geom, geom2: Geom, geom1_type: int, geom2_type: int, dir: wp.vec3
 ) -> Tuple[int, int]:
-  sp = _support(geom1, geom1_type, dir)
+  sp = support(geom1, geom1_type, dir)
   pt.vert1[idx] = sp.point
   pt.vert_index1[idx] = sp.vertex_index
   index1 = sp.cached_index
 
-  sp = _support(geom2, geom2_type, -dir)
+  sp = support(geom2, geom2_type, -dir)
   pt.vert2[idx] = sp.point
   pt.vert_index2[idx] = sp.vertex_index
   index2 = sp.cached_index
 
-  pt.vert[idx] = pt.vert1[idx] - pt.vert2[idx]
-
   return index1, index2
 
 
@@ -596,13 +595,13 @@ def gjk(
     dir_neg = x_k / wp.sqrt(xnorm)
 
     # compute kth support point in geom1
-    sp = _support(geom1, geomtype1, -dir_neg)
+    sp = support(geom1, geomtype1, -dir_neg)
     simplex1[n] = sp.point
     geom1.index = sp.cached_index
     simplex_index1[n] = sp.vertex_index
 
     # compute kth support point in geom2
-    sp = _support(geom2, geomtype2, dir_neg)
+    sp = support(geom2, geomtype2, dir_neg)
     simplex2[n] = sp.point
     geom2.index = sp.cached_index
     simplex_index2[n] = sp.vertex_index
@@ -766,11 +765,6 @@ def _replace_simplex3(pt: Polytope, v1: int, v2: int, v3: int) -> GJKResult:
   result = GJKResult()
 
   # reset GJK simplex
-  simplex = mat43()
-  simplex[0] = pt.vert[v1]
-  simplex[1] = pt.vert[v2]
-  simplex[2] = pt.vert[v3]
-
   simplex1 = mat43()
   simplex1[0] = pt.vert1[v1]
   simplex1[1] = pt.vert1[v2]
@@ -781,6 +775,11 @@ def _replace_simplex3(pt: Polytope, v1: int, v2: int, v3: int) -> GJKResult:
   simplex2[1] = pt.vert2[v2]
   simplex2[2] = pt.vert2[v3]
 
+  simplex = mat43()
+  simplex[0] = simplex1[0] - simplex2[0]
+  simplex[1] = simplex1[1] - simplex2[1]
+  simplex[2] = simplex1[2] - simplex2[2]
+
   simplex_index1 = wp.vec4i()
   simplex_index1[0] = pt.vert_index1[v1]
   simplex_index1[1] = pt.vert_index1[v2]
@@ -835,6 +834,11 @@ def _ray_triangle(v1: wp.vec3, v2: wp.vec3, v3: wp.vec3, v4: wp.vec3, v5: wp.vec
   return 0
 
 
+@wp.func
+def _get_edge(edge: int) -> wp.vec2i:
+  return wp.vec2i(edge & 0x3FF, (edge >> 10) & 0x3FF)
+
+
 @wp.func
 def _add_edge(pt: Polytope, e1: int, e2: int) -> int:
   n = pt.nhorizon
@@ -842,47 +846,30 @@ def _add_edge(pt: Polytope, e1: int, e2: int) -> int:
   if n < 0:
     return -1
 
+  edge = (wp.min(e1, e2) << 10) | wp.max(e1, e2)
+
   for i in range(n):
-    old_e1 = pt.horizon[2 * i + 0]
+    if edge == pt.horizon[i]:
-    old_e2 = pt.horizon[2 * i + 1]
+      pt.horizon[i] = pt.horizon[n - 1]
-    if (old_e1 == e1 and old_e2 == e2) or (old_e1 == e2 and old_e2 == e1):
-      pt.horizon[2 * i + 0] = pt.horizon[2 * (n - 1) + 0]
-      pt.horizon[2 * i + 1] = pt.horizon[2 * (n - 1) + 1]
       return n - 1
 
   # out of memory, force EPA to return early without contact
-  if n > pt.horizon.shape[0] - 2:
+  if n == pt.horizon.shape[0]:
     return -1
 
-  pt.horizon[2 * n + 0] = e1
+  pt.horizon[n] = edge
-  pt.horizon[2 * n + 1] = e2
   return n + 1
 
 
-@wp.func
-def _delete_face(pt: Polytope, face_id: int) -> int:
-  index = pt.face_index[face_id]
-  # delete from map
-  if index >= 0:
-    last_face = pt.face_map[pt.nmap - 1]
-    pt.face_map[index] = last_face
-    pt.face_index[last_face] = index
-    pt.nmap -= 1
-  # mark face as deleted from polytope
-  pt.face_index[face_id] = -2
-  return pt.nmap
-
-
 @wp.func
 def _epa_witness(
   pt: Polytope, geom1: Geom, geom2: Geom, geomtype1: int, geomtype2: int, face_idx: int
 ) -> Tuple[wp.vec3, wp.vec3, float]:
-  face = pt.face[face_idx]
+  face = _get_face_verts(pt.face[face_idx])
-
   # compute affine coordinates for witness points on plane defined by face
-  v1 = pt.vert[face[0]]
+  v1 = pt.vert1[face[0]] - pt.vert2[face[0]]
-  v2 = pt.vert[face[1]]
+  v2 = pt.vert1[face[1]] - pt.vert2[face[1]]
-  v3 = pt.vert[face[2]]
+  v3 = pt.vert1[face[2]] - pt.vert2[face[2]]
 
   coordinates = _tri_affine_coord(v1, v2, v3, pt.face_pr[face_idx])
   l1 = coordinates[0]
@@ -904,12 +891,12 @@ def _epa_witness(
   i3 = pt.vert_index1[face[2]]
   if geomtype1 == GeomType.HFIELD and (i1 != i2 or i1 != i3):
     # TODO(kbayes): Fix case where geom2 is near bottom of height field or "extreme" prism heights
-    n = geom1.rot[:, 2]
+    n = wp.vec3(0.0, 0.0, 1.0)
 
-    # height field prism vertices in global frame
+    # height field prism vertices
-    a = geom1.pos + geom1.rot @ geom1.hfprism[3]
+    a = geom1.hfprism[3]
-    b = geom1.pos + geom1.rot @ geom1.hfprism[4]
+    b = geom1.hfprism[4]
-    c = geom1.pos + geom1.rot @ geom1.hfprism[5]
+    c = geom1.hfprism[5]
 
     # TODO(kbayes): Support cases where geom2 is larger than the height field
     if geomtype2 == GeomType.CAPSULE or geomtype2 == GeomType.SPHERE:
@@ -917,13 +904,13 @@ def _epa_witness(
       margin = geom2.margin
       geom2.margin = 0.0
       geom2.size = wp.vec3(0.0, geom2.size[1], geom2.size[2])
-      sp = _support(geom2, geomtype2, x2)
+      sp = support(geom2, geomtype2, x2)
       x2 = sp.point - (0.5 * margin + radius) * n
       geom2.size[0] = radius
       geom2.margin = margin
     else:
       x2 = wp.normalize(x2)
-      sp = _support(geom2, geomtype2, x2)
+      sp = support(geom2, geomtype2, x2)
       x2 = sp.point
 
     coordinates2 = _tri_affine_coord(a, b, c, x2)
@@ -952,7 +939,6 @@ def _epa_witness(
 def _polytope2(
   # In:
   pt: Polytope,
-  dist: float,
   simplex: mat43,
   simplex1: mat43,
   simplex2: mat43,
@@ -985,9 +971,6 @@ def _polytope2(
   d3 = R @ d2
 
   # save vertices and get indices for each one
-  pt.vert[0] = simplex[0]
-  pt.vert[1] = simplex[1]
-
   pt.vert1[0] = simplex1[0]
   pt.vert1[1] = simplex1[1]
 
@@ -1030,19 +1013,16 @@ def _polytope2(
     return pt, _replace_simplex3(pt, 1, 4, 3)
 
   # check hexahedron is convex
-  if not _ray_triangle(simplex[0], simplex[1], pt.vert[2], pt.vert[3], pt.vert[4]):
+  v2 = pt.vert1[2] - pt.vert2[2]
+  v3 = pt.vert1[3] - pt.vert2[3]
+  v4 = pt.vert1[4] - pt.vert2[4]
+  if not _ray_triangle(simplex[0], simplex[1], v2, v3, v4):
     pt.status = 1
     return pt, GJKResult()
 
-  # populate face map
-  for i in range(6):
-    pt.face_map[i] = i
-    pt.face_index[i] = i
-
   # set polytope counts
   pt.nvert = 5
   pt.nface = 6
-  pt.nmap = 6
   pt.status = 0
   return pt, GJKResult()
 
@@ -1069,10 +1049,6 @@ def _polytope3(
     pt.status = 2
     return pt
 
-  pt.vert[0] = simplex[0]
-  pt.vert[1] = simplex[1]
-  pt.vert[2] = simplex[2]
-
   pt.vert1[0] = simplex1[0]
   pt.vert1[1] = simplex1[1]
   pt.vert1[2] = simplex1[2]
@@ -1095,8 +1071,8 @@ def _polytope3(
   v1 = simplex[0]
   v2 = simplex[1]
   v3 = simplex[2]
-  v4 = pt.vert[3]
+  v4 = pt.vert1[3] - pt.vert2[3]
-  v5 = pt.vert[4]
+  v5 = pt.vert1[4] - pt.vert2[4]
 
   # check that v4 is not contained in the 2-simplex
   if _tri_point_intersect(v1, v2, v3, v4):
@@ -1134,15 +1110,9 @@ def _polytope3(
     pt.status = 11
     return pt
 
-  # populate face map
-  for i in range(6):
-    pt.face_map[i] = i
-    pt.face_index[i] = i
-
   # set polytope counts
   pt.nvert = 5
   pt.nface = 6
-  pt.nmap = 6
   pt.status = 0
   return pt
 
@@ -1151,23 +1121,13 @@ def _polytope3(
 def _polytope4(
   # In:
   pt: Polytope,
-  dist: float,
   simplex: mat43,
   simplex1: mat43,
   simplex2: mat43,
   simplex_index1: wp.vec4i,
   simplex_index2: wp.vec4i,
-  geom1: Geom,
-  geom2: Geom,
-  geomtype1: int,
-  geomtype2: int,
 ) -> Tuple[Polytope, GJKResult]:
   """Create polytope for EPA given a 3-simplex from GJK."""
-  pt.vert[0] = simplex[0]
-  pt.vert[1] = simplex[1]
-  pt.vert[2] = simplex[2]
-  pt.vert[3] = simplex[3]
-
   pt.vert1[0] = simplex1[0]
   pt.vert1[1] = simplex1[1]
   pt.vert1[2] = simplex1[2]
@@ -1205,23 +1165,47 @@ def _polytope4(
     pt.status = -1
     return pt, _replace_simplex3(pt, 3, 2, 1)
 
-  if not _test_tetra(pt.vert[0], pt.vert[1], pt.vert[2], pt.vert[3]):
+  if not _test_tetra(simplex[0], simplex[1], simplex[2], simplex[3]):
     pt.status = 12
     return pt, GJKResult()
 
-  # populate face map
-  for i in range(4):
-    pt.face_map[i] = i
-    pt.face_index[i] = i
-
   # set polytope counts
   pt.nvert = 4
   pt.nface = 4
-  pt.nmap = 4
   pt.status = 0
   return pt, GJKResult()
 
 
+@wp.func
+def _get_face_verts(face: int) -> wp.vec3i:
+  """Return the three vertices of the face given by indices into the polytope vertex array."""
+  return wp.vec3i(face & 0x3FF, face >> 10 & 0x3FF, face >> 20 & 0x3FF)
+
+
+@wp.func
+def _delete_face(face: int) -> int:
+  """Return the face with the deleted bit enabled."""
+  return int(wp.uint32(face) | _FACE_DELETED_BIT)
+
+
+@wp.func
+def _is_face_deleted(face: int) -> bool:
+  """Return true if face is deleted."""
+  return bool(wp.uint32(face) & _FACE_DELETED_BIT)
+
+
+@wp.func
+def _invalidate_face(face: int) -> int:
+  """Return the face with the invalid bit enabled."""
+  return int(wp.uint32(face) | _FACE_INVALID_BIT)
+
+
+@wp.func
+def _is_invalid_face(face: int) -> bool:
+  """Return true if face is invalid or deleted."""
+  return bool(wp.uint32(face) & _FACE_INVALID_OR_DELETED_MASK)
+
+
 @wp.func
 def _epa(
   # In:
@@ -1242,18 +1226,22 @@ def _epa(
   pidx = int(-1)
   epsilon = wp.where(is_discrete, 1e-15, tolerance)
   cnt = int(1)
+  nvalid = pt.nface  # number of potential faces for expanding the polytope
 
+  # the face vertices are encoded in 10-bits that index the vertex array,
+  # so iterations must be cap to limit the number of generated vertices
+  # (one new vertex per iteration)
+  epa_iterations = wp.min(epa_iterations, 1000)
   for _ in range(epa_iterations):
     pidx = idx
     idx = int(-1)
+    lower2 = float(FLOAT_MAX)
 
     # find the face closest to the origin (lower bound for penetration depth)
-    lower2 = float(FLOAT_MAX)
+    for i in range(pt.nface):
-    for i in range(pt.nmap):
+      if not _is_invalid_face(pt.face[i]) and pt.face_norm2[i] < lower2:
-      face_idx = pt.face_map[i]
+        idx = i
-      if pt.face_norm2[face_idx] < lower2:
+        lower2 = pt.face_norm2[i]
-        idx = int(face_idx)
-        lower2 = float(pt.face_norm2[face_idx])
 
     # face not valid, return previous face
     if lower2 > upper2 or idx < 0:
@@ -1269,12 +1257,13 @@ def _epa(
     wi = pt.nvert
     face_pr_normalized = pt.face_pr[idx] / lower
     i1, i2 = _epa_support(pt, wi, geom1, geom2, geomtype1, geomtype2, face_pr_normalized)
+    w = pt.vert1[wi] - pt.vert2[wi]
     geom1.index = i1
     geom2.index = i2
     pt.nvert += 1
 
     # upper bound for kth iteration
-    upper_k = wp.dot(face_pr_normalized, pt.vert[wi])
+    upper_k = wp.dot(face_pr_normalized, w)
     if upper_k < upper:
       upper = upper_k
       upper2 = upper * upper
@@ -1292,45 +1281,51 @@ def _epa(
       if found_repeated:
         break
 
-    pt.nmap = _delete_face(pt, idx)
+    nvalid -= 1
-    pt.nhorizon = _add_edge(pt, pt.face[idx][0], pt.face[idx][1])
+    pt.face[idx] = _delete_face(pt.face[idx])
-    pt.nhorizon = _add_edge(pt, pt.face[idx][1], pt.face[idx][2])
+    face = _get_face_verts(pt.face[idx])
-    pt.nhorizon = _add_edge(pt, pt.face[idx][2], pt.face[idx][0])
+    pt.nhorizon = _add_edge(pt, face[0], face[1])
+    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])
       idx = -1
       break
 
     # compute horizon for w
     for i in range(pt.nface):
-      if pt.face_index[i] == -2:
+      if _is_face_deleted(pt.face[i]):
         continue
 
-      if wp.dot(pt.face_pr[i], pt.vert[wi]) - pt.face_norm2[i] > 1e-10:
+      if wp.dot(pt.face_pr[i], w) - pt.face_norm2[i] > 1e-10:
-        pt.nmap = _delete_face(pt, i)
+        nvalid = wp.where(_is_invalid_face(pt.face[i]), nvalid, nvalid - 1)
-        pt.nhorizon = _add_edge(pt, pt.face[i][0], pt.face[i][1])
+        pt.face[i] = _delete_face(pt.face[i])
-        pt.nhorizon = _add_edge(pt, pt.face[i][1], pt.face[i][2])
+        face = _get_face_verts(pt.face[i])
-        pt.nhorizon = _add_edge(pt, pt.face[i][2], pt.face[i][0])
+        pt.nhorizon = _add_edge(pt, face[0], face[1])
+        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])
           idx = -1
           break
 
     # insert w as new vertex and attach faces along the horizon
     for i in range(pt.nhorizon):
-      dist2 = _attach_face(pt, pt.nface, wi, pt.horizon[2 * i + 0], pt.horizon[2 * i + 1])
+      edge = _get_edge(pt.horizon[i])
+      dist2 = _attach_face(pt, pt.nface, wi, edge[0], edge[1])
       if dist2 == 0:
         idx = -1
         break
 
       pt.nface += 1
 
-      # store face in map
       if dist2 >= lower2 and dist2 <= upper2:
-        pt.face_map[pt.nmap] = pt.nface - 1
+        nvalid += 1
-        pt.face_index[pt.nface - 1] = pt.nmap
+      else:
-        pt.nmap += 1
+        pt.face[pt.nface - 1] = _invalidate_face(pt.face[pt.nface - 1])
 
     # no face candidates left
-    if pt.nmap == 0 or idx == -1:
+    if nvalid == 0 or idx == -1:
       break
 
     # clear horizon
@@ -1394,7 +1389,7 @@ def _polygon_quad(polygon: wp.array(dtype=wp.vec3), npolygon: int) -> wp.vec4i:
       if c == b:
         c = _next(npolygon, c)
         if d == c:
-          d == _next(npolygon, d)
+          d = _next(npolygon, d)
   return res
 
 
@@ -1889,7 +1884,7 @@ def _polygon_clip(
     for i in range(npolygon):
       # get edge PQ of the polygon
       P = polygon_out[i]
-      Q = wp.where(i < npolygon - 1, polygon_out[i + 1], polygon_out[0])
+      Q = polygon_out[(i + 1) % npolygon]
 
       # determine if P and Q are in the halfspace of the clipping edge
       inside1 = _halfspace(face1[e], pn[e], P)
@@ -1968,7 +1963,7 @@ def multicontact(
   epa_vert2: wp.array(dtype=wp.vec3),
   epa_vert_index1: wp.array(dtype=int),
   epa_vert_index2: wp.array(dtype=int),
-  face: wp.vec3i,
+  epa_face: int,
   x1: wp.vec3,
   x2: wp.vec3,
   geom1: Geom,
@@ -2000,6 +1995,8 @@ def multicontact(
     polymapnum = geom2.mesh_polymapnum
     polymap = geom2.mesh_polymap
 
+  face = _get_face_verts(epa_face)
+
   # get dimensions of features of geoms 1 and 2
   nface1, feature_index1, feature_vertex1 = _feature_dim(face, epa_vert_index1, epa_vert1)
   nface2, feature_index2, feature_vertex2 = _feature_dim(face, epa_vert_index2, epa_vert2)
@@ -2190,13 +2187,15 @@ def _inflate(
         break
 
     if is_side:
-      n = geom1.rot[:, 2]
+      n = wp.vec3(0.0, 0.0, 1.0)
-      sp = _support(geom2, geomtype2, x2)
+      sp = support(geom2, geomtype2, x2)
       x2 = sp.point - margin2 * n
 
+      # height field prism vertices
       a = geom1.hfprism[3]
       b = geom1.hfprism[4]
       c = geom1.hfprism[5]
+
       coordinates = _tri_affine_coord(a, b, c, x2)
       if coordinates[0] > 0 and coordinates[1] > 0 and coordinates[2] > 0:
         x1 = coordinates[0] * a + coordinates[1] * b + coordinates[2] * c
@@ -2231,16 +2230,13 @@ def ccd(
   geomtype2: int,
   x_1: wp.vec3,
   x_2: wp.vec3,
-  vert: wp.array(dtype=wp.vec3),
   vert1: wp.array(dtype=wp.vec3),
   vert2: wp.array(dtype=wp.vec3),
   vert_index1: wp.array(dtype=int),
   vert_index2: wp.array(dtype=int),
-  face: wp.array(dtype=wp.vec3i),
+  face: wp.array(dtype=int),
   face_pr: wp.array(dtype=wp.vec3),
   face_norm2: wp.array(dtype=float),
-  face_index: wp.array(dtype=int),
-  face_map: wp.array(dtype=int),
   horizon: wp.array(dtype=int),
 ) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
   """General convex collision detection via GJK/EPA."""
@@ -2291,10 +2287,8 @@ def ccd(
 
   pt = Polytope()
   pt.nface = 0
-  pt.nmap = 0
   pt.nvert = 0
   pt.nhorizon = 0
-  pt.vert = vert
   pt.vert1 = vert1
   pt.vert2 = vert2
   pt.vert_index1 = vert_index1
@@ -2302,14 +2296,11 @@ def ccd(
   pt.face = face
   pt.face_pr = face_pr
   pt.face_norm2 = face_norm2
-  pt.face_index = face_index
-  pt.face_map = face_map
   pt.horizon = horizon
 
   if result.dim == 2:
     pt, new_result = _polytope2(
       pt,
-      result.dist,
       result.simplex,
       result.simplex1,
       result.simplex2,
@@ -2330,16 +2321,11 @@ def ccd(
   elif result.dim == 4:
     pt, new_result = _polytope4(
       pt,
-      result.dist,
       result.simplex,
       result.simplex1,
       result.simplex2,
       result.simplex_index1,
       result.simplex_index2,
-      geom1,
-      geom2,
-      geomtype1,
-      geomtype2,
     )
     if pt.status == -1:
       result.simplex = new_result.simplex

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

@@ -822,39 +822,41 @@ def capsule_capsule_wrapper(
     cap2_axis,
     cap2.size[0],  # radius2
     cap2.size[1],  # half_length2
+    margin,
   )
 
-  write_contact(
+  for i in range(2):
-    naconmax_in,
+    write_contact(
-    0,
+      naconmax_in,
-    dist,
+      i,
-    pos,
+      dist[i],
-    make_frame(normal),
+      wp.vec3(pos[i, 0], pos[i, 1], pos[i, 2]),
-    margin,
+      make_frame(wp.vec3(normal[i, 0], normal[i, 1], normal[i, 2])),
-    gap,
+      margin,
-    condim,
+      gap,
-    friction,
+      condim,
-    solref,
+      friction,
-    solreffriction,
+      solref,
-    solimp,
+      solreffriction,
-    geoms,
+      solimp,
-    pairid,
+      geoms,
-    worldid,
+      pairid,
-    contact_dist_out,
+      worldid,
-    contact_pos_out,
+      contact_dist_out,
-    contact_frame_out,
+      contact_pos_out,
-    contact_includemargin_out,
+      contact_frame_out,
-    contact_friction_out,
+      contact_includemargin_out,
-    contact_solref_out,
+      contact_friction_out,
-    contact_solreffriction_out,
+      contact_solref_out,
-    contact_solimp_out,
+      contact_solreffriction_out,
-    contact_dim_out,
+      contact_solimp_out,
-    contact_geom_out,
+      contact_dim_out,
-    contact_worldid_out,
+      contact_geom_out,
-    contact_type_out,
+      contact_worldid_out,
-    contact_geomcollisionid_out,
+      contact_type_out,
-    nacon_out,
+      contact_geomcollisionid_out,
-  )
+      nacon_out,
+    )
 
 
 @wp.func

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

@@ -188,7 +188,8 @@ def capsule_capsule(
   cap2_axis: wp.vec3,
   cap2_radius: float,
   cap2_half_length: float,
-) -> Tuple[float, wp.vec3, wp.vec3]:
+  margin: float,
+) -> Tuple[wp.vec2, mat23f, mat23f]:
   """Core contact geometry calculation for capsule-capsule collision.
 
   Args:
@@ -200,28 +201,110 @@ def capsule_capsule(
     cap2_axis: Axis direction of the second capsule.
     cap2_radius: Radius of the second capsule.
     cap2_half_length: Half length of the second capsule.
+    margin: Collision margin for filtering contacts.
 
   Returns:
-    - Vector of contact distances.
+    - Vector of contact distances (wp.inf for invalid contacts).
     - Matrix of contact positions (one per row).
     - Matrix of contact normal vectors (one per row).
   """
-  # TODO(team): parallel axes case
+  contact_dist = wp.vec2(wp.inf, wp.inf)
+  contact_pos = mat23f()
+  contact_normal = mat23f()
 
-  # Calculate capsule segments
+  # calculate scaled axes and center difference
-  seg1 = cap1_axis * cap1_half_length
+  axis1 = cap1_axis * cap1_half_length
-  seg2 = cap2_axis * cap2_half_length
+  axis2 = cap2_axis * cap2_half_length
+  dif = cap1_pos - cap2_pos
 
-  # Find closest points between capsule centerlines
+  # compute matrix coefficients and determinant
-  pt1, pt2 = closest_segment_to_segment_points(
+  ma = wp.dot(axis1, axis1)
-    cap1_pos - seg1,
+  mb = -wp.dot(axis1, axis2)
-    cap1_pos + seg1,
+  mc = wp.dot(axis2, axis2)
-    cap2_pos - seg2,
+  u = -wp.dot(axis1, dif)
-    cap2_pos + seg2,
+  v = wp.dot(axis2, dif)
-  )
+  det = ma * mc - mb * mb
 
-  # Use sphere-sphere collision between closest points
+  # non-parallel axes: 1 contact
-  return sphere_sphere(pt1, cap1_radius, pt2, cap2_radius)
+  if wp.abs(det) >= MJ_MINVAL:
+    inv_det = 1.0 / det
+    x1 = (mc * u - mb * v) * inv_det
+    x2 = (ma * v - mb * u) * inv_det
+
+    if x1 > 1.0:
+      x1 = 1.0
+      x2 = (v - mb) / mc
+    elif x1 < -1.0:
+      x1 = -1.0
+      x2 = (v + mb) / mc
+
+    if x2 > 1.0:
+      x2 = 1.0
+      x1 = wp.clamp((u - mb) / ma, -1.0, 1.0)
+    elif x2 < -1.0:
+      x2 = -1.0
+      x1 = wp.clamp((u + mb) / ma, -1.0, 1.0)
+
+    # find nearest points
+    vec1 = cap1_pos + axis1 * x1
+    vec2 = cap2_pos + axis2 * x2
+
+    dist, pos, normal = sphere_sphere(vec1, cap1_radius, vec2, cap2_radius)
+    if dist <= margin:
+      contact_dist[0] = dist
+      contact_pos[0] = pos
+      contact_normal[0] = normal
+
+  # parallel axes: test all 4 endpoint pairs, keep first 2 that pass margin check
+  else:
+    contact_count = 0
+
+    # x1 = 1: test positive end of capsule 1
+    vec1 = cap1_pos + axis1
+    x2 = wp.clamp((v - mb) / mc, -1.0, 1.0)
+    vec2 = cap2_pos + axis2 * x2
+    dist, pos, normal = sphere_sphere(vec1, cap1_radius, vec2, cap2_radius)
+    if dist <= margin:
+      contact_dist[contact_count] = dist
+      contact_pos[contact_count] = pos
+      contact_normal[contact_count] = normal
+      contact_count += 1
+
+    # x1 = -1: test negative end of capsule 1
+    vec1 = cap1_pos - axis1
+    x2 = wp.clamp((v + mb) / mc, -1.0, 1.0)
+    vec2 = cap2_pos + axis2 * x2
+    dist, pos, normal = sphere_sphere(vec1, cap1_radius, vec2, cap2_radius)
+    if dist <= margin:
+      contact_dist[contact_count] = dist
+      contact_pos[contact_count] = pos
+      contact_normal[contact_count] = normal
+      contact_count += 1
+
+    # x2 = 1: test positive end of capsule 2
+    if contact_count < 2:
+      vec2 = cap2_pos + axis2
+      x1 = wp.clamp((u - mb) / ma, -1.0, 1.0)
+      vec1 = cap1_pos + axis1 * x1
+      dist, pos, normal = sphere_sphere(vec1, cap1_radius, vec2, cap2_radius)
+      if dist <= margin:
+        contact_dist[contact_count] = dist
+        contact_pos[contact_count] = pos
+        contact_normal[contact_count] = normal
+        contact_count += 1
+
+    # x2 = -1: test negative end of capsule 2
+    if contact_count < 2:
+      vec2 = cap2_pos - axis2
+      x1 = wp.clamp((u + mb) / ma, -1.0, 1.0)
+      vec1 = cap1_pos + axis1 * x1
+      dist, pos, normal = sphere_sphere(vec1, cap1_radius, vec2, cap2_radius)
+      if dist <= margin:
+        contact_dist[contact_count] = dist
+        contact_pos[contact_count] = pos
+        contact_normal[contact_count] = normal
+
+  return contact_dist, contact_pos, contact_normal
 
 
 @wp.func

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

@@ -66,7 +66,6 @@ class MeshData:
   mesh_faceadr: wp.array(dtype=int)
   mesh_face: wp.array(dtype=wp.vec3i)
   data_id: int
-  data_id: int
   pos: wp.vec3
   mat: wp.mat33
   pnt: wp.vec3
@@ -169,7 +168,7 @@ def grad_sphere(p: wp.vec3) -> wp.vec3:
   if c > 1e-9:
     return p / c
   else:
-    wp.vec3(0.0)
+    return wp.vec3(0.0)
 
 
 @wp.func

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

@@ -486,8 +486,8 @@ def implicit(m: Model, d: Data):
       qDeriv = wp.empty((d.nworld, 1, m.nM), dtype=float)
       qLD = wp.empty((d.nworld, 1, m.nC), dtype=float)
     else:
-      qDeriv = wp.empty((d.nworld, m.nv, m.nv), dtype=float)
+      qDeriv = wp.empty(d.qM.shape, dtype=float)
-      qLD = wp.empty((d.nworld, m.nv, m.nv), dtype=float)
+      qLD = wp.empty(d.qM.shape, dtype=float)
     qLDiagInv = wp.empty((d.nworld, m.nv), dtype=float)
     derivative.deriv_smooth_vel(m, d, qDeriv)
     qacc = wp.empty((d.nworld, m.nv), dtype=float)
@@ -956,6 +956,7 @@ def step1(m: Model, d: Data):
   # TODO(team): mj_checkPos
   # TODO(team): mj_checkVel
   fwd_position(m, d)
+  d.sensordata.zero_()
   sensor.sensor_pos(m, d)
 
   if energy:

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

@@ -14,6 +14,9 @@
 # ==============================================================================
 
 import dataclasses
+import importlib.metadata
+import re
+import warnings
 from typing import Any, Optional, Sequence, Union
 
 import mujoco
@@ -25,6 +28,24 @@ from mujoco.mjx.third_party.mujoco_warp._src import warp_util
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
 
 
+def _is_mujoco_dev() -> bool:
+  _DEV_VERSION_PATTERN = re.compile(r"^\d+\.\d+\.\d+.+")  # anything after x.y.z
+
+  version = getattr(__import__("mujoco"), "__version__", None)
+  if version and _DEV_VERSION_PATTERN.match(version):
+    return True
+
+  # fall back to metadata
+  dist_version = importlib.metadata.version("mujoco")
+  if _DEV_VERSION_PATTERN.match(dist_version):
+    return True
+
+  return False
+
+
+BLEEDING_EDGE_MUJOCO = _is_mujoco_dev()
+
+
 def _create_array(data: Any, spec: wp.array, sizes: dict[str, int]) -> Union[wp.array, None]:
   """Creates a warp array and populates it with data.
 
@@ -72,36 +93,38 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   warp_util.check_toolkit_driver()
 
   # model: check supported features in array types
-  for field, field_type in (
+  for field, field_type, mj_type in (
-    (mjm.actuator_trntype, types.TrnType),
+    (mjm.actuator_trntype, types.TrnType, mujoco.mjtTrn),
-    (mjm.actuator_dyntype, types.DynType),
+    (mjm.actuator_dyntype, types.DynType, mujoco.mjtDyn),
-    (mjm.actuator_gaintype, types.GainType),
+    (mjm.actuator_gaintype, types.GainType, mujoco.mjtGain),
-    (mjm.actuator_biastype, types.BiasType),
+    (mjm.actuator_biastype, types.BiasType, mujoco.mjtBias),
-    (mjm.eq_type, types.EqType),
+    (mjm.eq_type, types.EqType, mujoco.mjtEq),
-    (mjm.geom_type, types.GeomType),
+    (mjm.geom_type, types.GeomType, mujoco.mjtGeom),
-    (mjm.sensor_type, types.SensorType),
+    (mjm.sensor_type, types.SensorType, mujoco.mjtSensor),
-    (mjm.wrap_type, types.WrapType),
+    (mjm.wrap_type, types.WrapType, mujoco.mjtWrap),
   ):
     missing = ~np.isin(field, field_type)
     if missing.any():
-      raise NotImplementedError(f"{field_type.__name__}: {field[missing]} not supported.")
+      names = [mj_type(v).name for v in field[missing]]
+      raise NotImplementedError(f"{names} not supported.")
 
   # opt: check supported features in scalar types
-  for field, field_type in (
+  for field, field_type, mj_type in (
-    (mjm.opt.integrator, types.IntegratorType),
+    (mjm.opt.integrator, types.IntegratorType, mujoco.mjtIntegrator),
-    (mjm.opt.cone, types.ConeType),
+    (mjm.opt.cone, types.ConeType, mujoco.mjtCone),
-    (mjm.opt.solver, types.SolverType),
+    (mjm.opt.solver, types.SolverType, mujoco.mjtSolver),
   ):
     if field not in set(field_type):
-      raise NotImplementedError(f"{field_type.__name__} {field} is unsupported.")
+      raise NotImplementedError(f"{mj_type(field).name} is unsupported.")
 
   # opt: check supported features in scalar flag types
-  for field, field_type in (
+  for field, field_type, mj_type in (
-    (mjm.opt.disableflags, types.DisableBit),
+    (mjm.opt.disableflags, types.DisableBit, mujoco.mjtDisableBit),
-    (mjm.opt.enableflags, types.EnableBit),
+    (mjm.opt.enableflags, types.EnableBit, mujoco.mjtEnableBit),
   ):
-    if field & ~np.bitwise_or.reduce(field_type):
+    unsupported = field & ~np.bitwise_or.reduce(field_type)
-      raise NotImplementedError(f"{field_type.__name__} {field} is unsupported.")
+    if unsupported:
+      raise NotImplementedError(f"{mj_type(unsupported).name} is unsupported.")
 
   if ((mjm.flex_contype != 0) | (mjm.flex_conaffinity != 0)).any():
     raise NotImplementedError("Flex collisions are not implemented.")
@@ -153,6 +176,21 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     if not_implemented(objtype, objid, types.GeomType.BOX) and not_implemented(reftype, refid, types.GeomType.BOX):
       raise NotImplementedError(f"Collision sensors with box-box collisions are not implemented.")
 
+  def _check_friction(name: str, id_: int, condim: int, friction, checks):
+    for min_condim, indices in checks:
+      if condim >= min_condim:
+        for idx in indices:
+          if friction[idx] < types.MJ_MINMU:
+            warnings.warn(
+              f"{name} {id_}: friction[{idx}] ({friction[idx]}) < MJ_MINMU ({types.MJ_MINMU}) with condim={condim} may cause NaN"
+            )
+
+  for geomid in range(mjm.ngeom):
+    _check_friction("geom", geomid, mjm.geom_condim[geomid], mjm.geom_friction[geomid], [(3, [0]), (4, [1]), (6, [2])])
+
+  for pairid in range(mjm.npair):
+    _check_friction("pair", pairid, mjm.pair_dim[pairid], mjm.pair_friction[pairid], [(3, [0]), (4, [1, 2]), (6, [3, 4])])
+
   # create opt
   opt_kwargs = {f.name: getattr(mjm.opt, f.name, None) for f in dataclasses.fields(types.Option)}
   if hasattr(mjm.opt, "impratio"):
@@ -856,16 +894,22 @@ def put_data(
     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))
     d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), ten_J), dtype=float)
-    flexedge_J = np.zeros((mjm.nflexedge, mjm.nv))
-    mujoco.mju_sparse2dense(
-      flexedge_J, mjd.flexedge_J.reshape(-1), mjm.flexedge_J_rownnz, mjm.flexedge_J_rowadr, mjm.flexedge_J_colind.reshape(-1)
-    )
-    d.flexedge_J = wp.array(np.full((nworld, mjm.nflexedge, mjm.nv), flexedge_J), dtype=float)
   else:
     ten_J = mjd.ten_J.reshape((mjm.ntendon, mjm.nv))
     d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), ten_J), dtype=float)
-    flexedge_J = mjd.flexedge_J.reshape((mjm.nflexedge, mjm.nv))
+
-    d.flexedge_J = wp.array(np.full((nworld, mjm.nflexedge, mjm.nv), flexedge_J), dtype=float)
+  flexedge_J = np.zeros((mjm.nflexedge, mjm.nv))
+  if mjd.flexedge_J.size:
+    # TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml
+    if BLEEDING_EDGE_MUJOCO:
+      mujoco.mju_sparse2dense(
+        flexedge_J, mjd.flexedge_J.reshape(-1), mjm.flexedge_J_rownnz, mjm.flexedge_J_rowadr, mjm.flexedge_J_colind.reshape(-1)
+      )
+    else:
+      mujoco.mju_sparse2dense(
+        flexedge_J, mjd.flexedge_J.reshape(-1), mjd.flexedge_J_rownnz, mjd.flexedge_J_rowadr, mjd.flexedge_J_colind.reshape(-1)
+      )
+  d.flexedge_J = wp.array(np.full((nworld, mjm.nflexedge, mjm.nv), flexedge_J), dtype=float)
 
   # TODO(taylorhowell): sparse actuator_moment
   actuator_moment = np.zeros((mjm.nu, mjm.nv))
@@ -980,16 +1024,31 @@ def get_data_into(
   result.cdof[:] = d.cdof.numpy()[world_id]
   result.cinert[:] = d.cinert.numpy()[world_id]
   result.flexvert_xpos[:] = d.flexvert_xpos.numpy()[world_id]
-  result.flexedge_J[:] = d.flexedge_J.numpy()[world_id]
+  flexedge_J = d.flexedge_J.numpy()[world_id]
+  if result.flexedge_J.size:
+    # TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml
+    if BLEEDING_EDGE_MUJOCO:
+      mujoco.mju_dense2sparse(
+        result.flexedge_J.reshape(-1),
+        flexedge_J,
+        mjm.flexedge_J_rownnz,
+        mjm.flexedge_J_rowadr,
+        mjm.flexedge_J_colind.reshape(-1),
+      )
+    else:
+      mujoco.mju_dense2sparse(
+        result.flexedge_J.reshape(-1),
+        flexedge_J,
+        result.flexedge_J_rownnz,
+        result.flexedge_J_rowadr,
+        result.flexedge_J_colind.reshape(-1),
+      )
   result.flexedge_length[:] = d.flexedge_length.numpy()[world_id]
   result.flexedge_velocity[:] = d.flexedge_velocity.numpy()[world_id]
   result.actuator_length[:] = d.actuator_length.numpy()[world_id]
+  actuator_moment = d.actuator_moment.numpy()[world_id]
   mujoco.mju_dense2sparse(
-    result.actuator_moment,
+    result.actuator_moment, actuator_moment, result.moment_rownnz, result.moment_rowadr, result.moment_colind
-    d.actuator_moment.numpy()[world_id],
-    result.moment_rownnz,
-    result.moment_rowadr,
-    result.moment_colind,
   )
   result.crb[:] = d.crb.numpy()[world_id]
   result.qLDiagInv[:] = d.qLDiagInv.numpy()[world_id]

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

@@ -595,15 +595,15 @@ def _flex_elasticity(
   nedge = nvert * (nvert - 1) / 2
   edges = wp.where(
     dim == 3,
-    wp.types.matrix(0, 1, 1, 2, 2, 0, 2, 3, 0, 3, 1, 3, shape=(6, 2), dtype=int),
+    wp.matrix(0, 1, 1, 2, 2, 0, 2, 3, 0, 3, 1, 3, shape=(6, 2), dtype=int),
-    wp.types.matrix(1, 2, 2, 0, 0, 1, 0, 0, 0, 0, 0, 0, shape=(6, 2), dtype=int),
+    wp.matrix(1, 2, 2, 0, 0, 1, 0, 0, 0, 0, 0, 0, shape=(6, 2), dtype=int),
   )
   if timestep > 0.0 and not dsbl_damper:
     kD = flex_damping[f] / timestep
   else:
     kD = 0.0
 
-  gradient = wp.types.matrix(0.0, shape=(6, 6))
+  gradient = wp.matrix(0.0, shape=(6, 6))
   for e in range(nedge):
     vert0 = flex_elem[(dim + 1) * elemid + edges[e, 0]]
     vert1 = flex_elem[(dim + 1) * elemid + edges[e, 1]]
@@ -622,7 +622,7 @@ def _flex_elasticity(
     previous = deformed - vel * timestep
     elongation[e] = deformed * deformed - reference * reference + (deformed * deformed - previous * previous) * kD
 
-  metric = wp.types.matrix(0.0, shape=(6, 6))
+  metric = wp.matrix(0.0, shape=(6, 6))
   id = int(0)
   for ed1 in range(nedge):
     for ed2 in range(ed1, nedge):
@@ -630,7 +630,7 @@ def _flex_elasticity(
       metric[ed2, ed1] = flex_stiffness[elemid, id]
       id += 1
 
-  force = wp.types.matrix(0.0, shape=(6, 3))
+  force = wp.matrix(0.0, shape=(6, 3))
   for ed1 in range(nedge):
     for ed2 in range(nedge):
       for i in range(2):
@@ -684,7 +684,7 @@ def _flex_bending(
     flex_vertadr[f] + flex_edgeflap[edgeid][1],
   )
 
-  frc = wp.types.matrix(0.0, shape=(4, 3))
+  frc = wp.matrix(0.0, shape=(4, 3))
   if flex_bending[edgeid, 16]:
     v0 = flexvert_xpos_in[worldid, v[0]]
     v1 = flexvert_xpos_in[worldid, v[1]]
@@ -695,7 +695,7 @@ def _flex_bending(
     frc[3] = wp.cross(v1 - v0, v2 - v0)
     frc[0] = -(frc[1] + frc[2] + frc[3])
 
-  force = wp.types.matrix(0.0, shape=(nvert, 3))
+  force = wp.matrix(0.0, shape=(nvert, 3))
   for i in range(nvert):
     for x in range(3):
       for j in range(nvert):
@@ -786,24 +786,24 @@ def passive(m: Model, d: Data):
       ],
       outputs=[d.qfrc_spring],
     )
-  wp.launch(
+    wp.launch(
-    _flex_bending,
+      _flex_bending,
-    dim=(d.nworld, m.nflexedge),
+      dim=(d.nworld, m.nflexedge),
-    inputs=[
+      inputs=[
-      m.nflex,
+        m.nflex,
-      m.body_dofadr,
+        m.body_dofadr,
-      m.flex_dim,
+        m.flex_dim,
-      m.flex_vertadr,
+        m.flex_vertadr,
-      m.flex_edgeadr,
+        m.flex_edgeadr,
-      m.flex_edgenum,
+        m.flex_edgenum,
-      m.flex_vertbodyid,
+        m.flex_vertbodyid,
-      m.flex_edge,
+        m.flex_edge,
-      m.flex_edgeflap,
+        m.flex_edgeflap,
-      m.flex_bending,
+        m.flex_bending,
-      d.flexvert_xpos,
+        d.flexvert_xpos,
-    ],
+      ],
-    outputs=[d.qfrc_spring],
+      outputs=[d.qfrc_spring],
-  )
+    )
 
   gravcomp = m.ngravcomp and not (m.opt.disableflags & DisableBit.GRAVITY)
 

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

@@ -572,7 +572,7 @@ def _sensor_pos(
     elif sensortype == SensorType.FRAMEZAXIS:
       axis = 2
     vec3 = _frame_axis(
-      ximat_in, xmat_in, geom_xmat_in, site_xmat_in, cam_xmat_in, worldid, objid, objtype, refid, reftype, axis
+      xmat_in, ximat_in, geom_xmat_in, site_xmat_in, cam_xmat_in, worldid, objid, objtype, refid, reftype, axis
     )
     _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
   elif sensortype == SensorType.FRAMEQUAT:

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

@@ -247,32 +247,6 @@ def xfrc_accumulate(m: Model, d: Data, qfrc: wp.array2d(dtype=float)):
   apply_ft(m, d, d.xfrc_applied, qfrc, True)
 
 
-@wp.func
-def all_same(v0: wp.vec3, v1: wp.vec3) -> wp.bool:
-  dx = abs(v0[0] - v1[0])
-  dy = abs(v0[1] - v1[1])
-  dz = abs(v0[2] - v1[2])
-
-  return (
-    (dx <= 1.0e-9 or dx <= max(abs(v0[0]), abs(v1[0])) * 1.0e-9)
-    and (dy <= 1.0e-9 or dy <= max(abs(v0[1]), abs(v1[1])) * 1.0e-9)
-    and (dz <= 1.0e-9 or dz <= max(abs(v0[2]), abs(v1[2])) * 1.0e-9)
-  )
-
-
-@wp.func
-def any_different(v0: wp.vec3, v1: wp.vec3) -> wp.bool:
-  dx = abs(v0[0] - v1[0])
-  dy = abs(v0[1] - v1[1])
-  dz = abs(v0[2] - v1[2])
-
-  return (
-    (dx > 1.0e-9 and dx > max(abs(v0[0]), abs(v1[0])) * 1.0e-9)
-    or (dy > 1.0e-9 and dy > max(abs(v0[1]), abs(v1[1])) * 1.0e-9)
-    or (dz > 1.0e-9 and dz > max(abs(v0[2]), abs(v1[2])) * 1.0e-9)
-  )
-
-
 @wp.func
 def _decode_pyramid(
   njmax_in: int, pyramid: wp.array(dtype=float), efc_address: int, mu: vec5, condim: int

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

@@ -25,7 +25,7 @@ MJ_MAXIMP = mujoco.mjMAXIMP  # maximum constraint impedance
 MJ_MAXCONPAIR = mujoco.mjMAXCONPAIR
 MJ_MINMU = mujoco.mjMINMU  # minimum friction
 # maximum size (by number of edges) of an horizon in EPA algorithm
-MJ_MAX_EPAHORIZON = 12
+MJ_MAX_EPAHORIZON = 24
 # maximum average number of trianglarfaces EPA can insert at each iteration
 MJ_MAX_EPAFACES = 5
 
@@ -173,11 +173,13 @@ class EnableBit(enum.IntFlag):
   Attributes:
     ENERGY: energy computation
     INVDISCRETE: discrete-time inverse dynamics
+    MULTICCD: multiple contacts with CCD
   """
 
   ENERGY = mujoco.mjtEnableBit.mjENBL_ENERGY
   INVDISCRETE = mujoco.mjtEnableBit.mjENBL_INVDISCRETE
-  # unsupported: OVERRIDE, FWDINV, ISLAND, MULTICCD
+  MULTICCD = mujoco.mjtEnableBit.mjENBL_MULTICCD
+  # unsupported: OVERRIDE, FWDINV, ISLAND
 
 
 class TrnType(enum.IntEnum):

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

@@ -18,8 +18,6 @@ import inspect
 from typing import Callable, Optional
 
 import warp as wp
-from warp._src.context import Module
-from warp._src.context import get_module
 
 _STACK = None
 
@@ -127,7 +125,7 @@ def nested_kernel(
   f: Optional[Callable] = None,
   *,
   enable_backward: Optional[bool] = None,
-  module: Optional[Module] = None,
+  module: Optional[wp.Module] = None,
 ):
   """Decorator to register a Warp kernel from a Python function.
 
@@ -166,7 +164,7 @@ def nested_kernel(
   Args:
       f: The function to be registered as a kernel.
       enable_backward: If False, the backward pass will not be generated.
-      module: The :class:`warp.context.Module` to which the kernel belongs. Alternatively,
+      module: The :class:`warp.Module` to which the kernel belongs. Alternatively,
               if a string `"unique"` is provided, the kernel is assigned to a new module
               named after the kernel name and hash. If None, the module is inferred from
               the function's module.
@@ -182,7 +180,7 @@ def nested_kernel(
       qualname = func.__qualname__
       parts = [part for part in qualname.split(".") if part != "<locals>"]
       outer_functions = parts[:-1]
-      module_name = get_module(".".join([func.__module__] + outer_functions))
+      module_name = wp.get_module(".".join([func.__module__] + outer_functions))
     else:
       module_name = module
 
@@ -220,8 +218,7 @@ def cache_kernel(func):
 
 
 def check_toolkit_driver():
-  if wp._src.context.runtime is None:
+  wp.init()
-    wp._src.context.init()
   if wp.get_device().is_cuda:
-    if wp._src.context.runtime.toolkit_version < (12, 4) or wp._src.context.runtime.driver_version < (12, 4):
+    if not wp.is_conditional_graph_supported():
       RuntimeError("Minimum supported CUDA version: 12.4.")

mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name="mujoco-warp"
-version = "0.0.1"
+version = "0.0.2"
 # TODO(team): create a distribution list
 authors = [
     {name = "Newton Developers", email = "mujoco@deepmind.com"},
@@ -27,9 +27,9 @@ requires-python = ">=3.10"
 dependencies = [
     "absl-py",
     "etils[epath]",
-    "mujoco>=3.3.7",
+    "mujoco>=3.4.0",
     "numpy",
-    "warp-lang>=1.9.1",
+    "warp-lang>=1.11.0",
 ]
 
 [[tool.uv.index]]
@@ -48,7 +48,6 @@ mujoco = {index = "mujoco"}
 
 [project.optional-dependencies]
 dev = [
-    "asv",
     "pre-commit",
     "pytest",
     "pytest-xdist",
@@ -56,7 +55,7 @@ dev = [
     "pygls>=1.0.0,<2.0.0",
     "lsprotocol>=2023.0.1,<2024.0.0",
     "mujoco>=3.3.7.dev0",
-    "warp-lang>=1.9.1.dev0",
+    "warp-lang>=1.11.0.dev0",
 ]
 # TODO(team): cpu and cuda JAX optional dependencies are temporary, remove after we land MJX:Warp
 cpu = [
@@ -113,6 +112,10 @@ convention = "google"
 docstring-code-format = true
 docstring-code-line-length = 100
 
+[tool.pytest.ini_options]
+testpaths = ["mujoco_warp"]
+norecursedirs = ["benchmarks", "contrib"]
+
 [tool.setuptools]
 package-data = { "mujoco_warp" = ["test_data/**"] }
 

mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py

@@ -121,7 +121,6 @@ def _main(argv: Sequence[str]) -> None:
     mujoco.mj_resetDataKeyframe(mjm, mjd, keys[0])
   elif mjm.nkey > 0 and _KEYFRAME.value > -1:
     mujoco.mj_resetDataKeyframe(mjm, mjd, _KEYFRAME.value)
-  mujoco.mj_forward(mjm, mjd)
 
   if _ENGINE.value == EngineOptions.C:
     override_model(mjm, _OVERRIDE.value)

mjx/mujoco/mjx/warp/collision_driver.py

@@ -42,6 +42,7 @@ _e = mjwarp.Constraint(
     **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
 )
 
+
 @ffi.format_args_for_warp
 def _collision_shim(
     # Model
@@ -108,6 +109,7 @@ def _collision_shim(
     opt__ccd_iterations: int,
     opt__ccd_tolerance: wp.array(dtype=float),
     opt__disableflags: int,
+    opt__enableflags: int,
     opt__sdf_initpoints: int,
     opt__sdf_iterations: int,
     # Data
@@ -190,6 +192,7 @@ def _collision_shim(
   _m.opt.ccd_iterations = opt__ccd_iterations
   _m.opt.ccd_tolerance = opt__ccd_tolerance
   _m.opt.disableflags = opt__disableflags
+  _m.opt.enableflags = opt__enableflags
   _m.opt.sdf_initpoints = opt__sdf_initpoints
   _m.opt.sdf_iterations = opt__sdf_iterations
   _m.pair_dim = pair_dim
@@ -341,6 +344,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       m.opt._impl.ccd_iterations,
       m.opt._impl.ccd_tolerance,
       m.opt.disableflags,
+      m.opt.enableflags,
       m.opt._impl.sdf_initpoints,
       m.opt._impl.sdf_iterations,
       d._impl.naconmax,

mjx/mujoco/mjx/warp/ffi.py

@@ -150,7 +150,7 @@ def _format_arg(arg: Any, name: str, annotation: Any, verbose: bool):
         for i in range(len(arg))
     )
 
-  if not isinstance(annotation, wp.types.array):
+  if not isinstance(annotation, wp.array):
     if verbose:
       print(f'Skipping {name}: {arg}')
     return arg