Import google-deepmind/mujoco_warp from GitHub.

PiperOrigin-RevId: 844675206
Change-Id: I02f75561eac3520d7107708dcece10cc82c106cd

mjx/cuda_requirements.txt

@@ -16,8 +16,12 @@ jax-cuda12-pjrt==0.5.3; python_version >= '3.10' \
 jax-cuda12-pjrt==0.4.30; python_version == '3.9' \
     --hash=sha256:895d0198ad99638fcaf976c47592e2a543eef79ea15fabd24a402d055390c328 \
     --hash=sha256:c36fb1e0c236563bf3a87e70f4d1ab28a31d7cf5d722c9ede30c4172116e8bcb
-warp-lang==1.10.0 \
+warp-lang==1.10.1 \
-    --hash=sha256:428a6388ba8c9b3ded973226ddb5be59b16e3b3f28ac939a5036ad2d7cdc79ed \
+    --hash=sha256:0c6f44d4136cfc86316f5d35883a863b4dca2bc318331597e755ea872db0cd97 \
-    --hash=sha256:1bdff31e170b00c89fb9d8b647e906fefdcdcf8741925126d1c0be42783174fa \
+    --hash=sha256:2884b642f16b07b930b3605193c1b97d183eef80c5b0083d3a17473aee92138e \
-    --hash=sha256:4aa8eb63cae5ee0d6dbdbdfc305d124140ec975475c97d4458f412396fb39eab \
+    --hash=sha256:7c8f839a01042677d1f417d31abe9304f9319fca7041c2ed82937ed511eb2362 \
-    --hash=sha256:81f73055e76a6a3f2284cf2b5fe542a341156861101df9f291bd5b59925ff6e5
+    --hash=sha256:7addd14a913b50c406ba3b1abebf3a208d01591a1c382da799cff2806d234c5a \
+    --hash=sha256:5f5ce0147d48f86ccab66b834aa54c8adedbbb2ba86e45b2b904b182fa963509 \
+    --hash=sha256:2068e877dd2972d841c1eb63f98eadce326e77ba58f09e8474d0847276f47633 \
+    --hash=sha256:96d24694333e30eac888b2c2962bc5ca4b2369f56707e4c54a465fa37b7ae007 \
+    --hash=sha256:b77ffe935e06cb57e76f6d4dd8bae7d59294ce63977ea8a0a4fd8e30d77c0a7f

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

@@ -56,6 +56,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.smooth import com_pos as com_pos
 from mujoco.mjx.third_party.mujoco_warp._src.smooth import com_vel as com_vel
 from mujoco.mjx.third_party.mujoco_warp._src.smooth import crb as crb
 from mujoco.mjx.third_party.mujoco_warp._src.smooth import factor_m as factor_m
+from mujoco.mjx.third_party.mujoco_warp._src.smooth import flex as flex
 from mujoco.mjx.third_party.mujoco_warp._src.smooth import kinematics as kinematics
 from mujoco.mjx.third_party.mujoco_warp._src.smooth import rne as rne
 from mujoco.mjx.third_party.mujoco_warp._src.smooth import rne_postconstraint as rne_postconstraint

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

@@ -174,6 +174,7 @@ class BenchmarkSuite:
   batch_size = -1
   nconmax = -1
   njmax = -1
+  nstep = 1000
   param_names = ("function",)
   params = (
     "jit_duration",
@@ -239,7 +240,7 @@ class BenchmarkSuite:
     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, 1000, ctrls, True, False, True)
+    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),

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

@@ -23,9 +23,8 @@ def create_blocked_cholesky_func(block_size: int):
   @wp.func
   def blocked_cholesky_func(
     # In:
-    tid_block: int,
     A: wp.array(dtype=float, ndim=2),
-    active_matrix_size: int,
+    matrix_size: int,
     # Out:
     L: wp.array(dtype=float, ndim=2),
   ):
@@ -33,89 +32,57 @@ def create_blocked_cholesky_func(block_size: int):
 
     It returns a lower-triangular matrix L such that A = L L^T.
     """
-    num_threads_per_block = wp.block_dim()
+    # TODO(team): remove conditional after mjwarp relies on >= 1.11
-
+    bleeding_edge_warp = wp.static(wp.__version__ >= "1.11")
-    # Round up active_matrix_size to next multiple of block_size
-    n = ((active_matrix_size + block_size - 1) // block_size) * block_size
 
     # Process the matrix in blocks along its leading dimension.
-    for k in range(0, n, block_size):
+    for k in range(0, matrix_size, block_size):
       end = k + block_size
 
       # Load current diagonal block A[k:end, k:end]
       # and update with contributions from previously computed blocks.
       A_kk_tile = wp.tile_load(A, shape=(block_size, block_size), offset=(k, k), storage="shared")
-      # The following if pads the matrix if it is not divisible by block_size
-      if k + block_size > active_matrix_size or k + block_size > active_matrix_size:
-        num_tile_elements = block_size * block_size
-        num_iterations = (num_tile_elements + num_threads_per_block - 1) // num_threads_per_block
 
-        for i in range(num_iterations):
+      for j in range(0, k, block_size):
-          linear_index = tid_block + i * num_threads_per_block
+        L_block = wp.tile_load(L, shape=(block_size, block_size), offset=(k, j), storage="shared")
-          linear_index = linear_index % num_tile_elements
+        if bleeding_edge_warp:
-          row = linear_index // block_size
+          wp.tile_matmul(L_block, wp.tile_transpose(L_block), A_kk_tile, alpha=-1.0)
-          col = linear_index % block_size
+        else:
-          value = A_kk_tile[row, col]
+          A_kk_tile -= wp.tile_matmul(L_block, wp.tile_transpose(L_block))
-          if k + row >= active_matrix_size or k + col >= active_matrix_size:
-            value = wp.where(row == col, float(1), float(0))
-          A_kk_tile[row, col] = value
-
-      if k > 0:
-        for j in range(0, k, block_size):
-          L_block = wp.tile_load(L, shape=(block_size, block_size), offset=(k, j))
-          L_block_T = wp.tile_transpose(L_block)
-          L_L_T_block = wp.tile_matmul(L_block, L_block_T)
-          A_kk_tile -= L_L_T_block
 
       # Compute the Cholesky factorization for the block
       L_kk_tile = wp.tile_cholesky(A_kk_tile)
       wp.tile_store(L, L_kk_tile, offset=(k, k))
 
       # Process the blocks below the current block
-      for i in range(end, n, block_size):
+      for i in range(end, matrix_size, block_size):
         A_ik_tile = wp.tile_load(A, shape=(block_size, block_size), offset=(i, k), storage="shared")
-        # The following if pads the matrix if it is not divisible by block_size
-        if i + block_size > active_matrix_size or k + block_size > active_matrix_size:
-          num_tile_elements = block_size * block_size
-          num_iterations = (num_tile_elements + num_threads_per_block - 1) // num_threads_per_block
 
-          for ii in range(num_iterations):
+        for j in range(0, k, block_size):
-            linear_index = tid_block + ii * num_threads_per_block
+          L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, j), storage="shared")
-            linear_index = linear_index % num_tile_elements
+          L_2_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(k, j), storage="shared")
-            row = linear_index // block_size
+          if bleeding_edge_warp:
-            col = linear_index % block_size
+            wp.tile_matmul(L_tile, wp.tile_transpose(L_2_tile), A_ik_tile, alpha=-1.0)
-            value = A_ik_tile[row, col]
+          else:
-            if i + row >= active_matrix_size or k + col >= active_matrix_size:
+            A_ik_tile -= wp.tile_matmul(L_tile, wp.tile_transpose(L_2_tile))
-              value = wp.where(i + row == k + col, float(1), float(0))
-            A_ik_tile[row, col] = value
 
-        if k > 0:
+        if bleeding_edge_warp:
-          for j in range(0, k, block_size):
+          wp.tile_lower_solve_inplace(L_kk_tile, wp.tile_transpose(A_ik_tile))
-            L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, j))
+        else:
-            L_2_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(k, j))
+          A_ik_tile = wp.tile_transpose(wp.tile_lower_solve(L_kk_tile, wp.tile_transpose(A_ik_tile)))
-            L_T_tile = wp.tile_transpose(L_2_tile)
+        wp.tile_store(L, A_ik_tile, offset=(i, k))
-            L_L_T_tile = wp.tile_matmul(L_tile, L_T_tile)
-            A_ik_tile -= L_L_T_tile
-
-        t = wp.tile_transpose(A_ik_tile)
-        tmp = wp.tile_lower_solve(L_kk_tile, t)
-        sol_tile = wp.tile_transpose(tmp)
-
-        wp.tile_store(L, sol_tile, offset=(i, k))
 
   return blocked_cholesky_func
 
 
 @lru_cache(maxsize=None)
-def create_blocked_cholesky_solve_func(block_size: int):
+def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int):
   @wp.func
   def blocked_cholesky_solve_func(
     # In:
-    tid_block: int,
     L: wp.array(dtype=float, ndim=2),
     b: wp.array(dtype=float, ndim=2),
-    tmp: wp.array(dtype=float, ndim=2),
+    matrix_size: int,
-    active_matrix_size: int,
     # Out:
     x: wp.array(dtype=float, ndim=2),
   ):
@@ -124,94 +91,45 @@ def create_blocked_cholesky_solve_func(block_size: int):
     Solves A x = b given the Cholesky factor L (A = L L^T) using blocked forward and backward
     substitution.
     """
-    num_threads_per_block = wp.block_dim()
+    # TODO(team): remove conditional after mjwarp relies on >= 1.11
-
+    bleeding_edge_warp = wp.static(wp.__version__ >= "1.11")
-    # Round up active_matrix_size to next multiple of block_size
+    rhs_tile = wp.tile_load(b, shape=(matrix_size_static, 1), offset=(0, 0), storage="shared", bounds_check=False)
-    n = ((active_matrix_size + block_size - 1) // block_size) * block_size
 
     # Forward substitution: solve L y = b
-    for i in range(0, n, block_size):
+    for i in range(0, matrix_size, block_size):
-      i_end = i + block_size
+      rhs_view = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(i, 0))
-      rhs_tile = wp.tile_load(b, shape=(block_size, 1), offset=(i, 0))
+      for j in range(0, i, block_size):
-      if i > 0:
+        L_block = wp.tile_load(L, shape=(block_size, block_size), offset=(i, j), storage="shared")
-        for j in range(0, i, block_size):
+        y_block = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
-          L_block = wp.tile_load(L, shape=(block_size, block_size), offset=(i, j))
+        if bleeding_edge_warp:
-          y_block = wp.tile_load(tmp, shape=(block_size, 1), offset=(j, 0))
+          wp.tile_matmul(L_block, y_block, rhs_view, alpha=-1.0)
-          Ly_block = wp.tile_matmul(L_block, y_block)
+        else:
-          rhs_tile -= Ly_block
+          rhs_view -= wp.tile_matmul(L_block, y_block)
-      L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, i))
 
-      # The following if pads the matrix if it is not divisible by block_size
+      L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, i), storage="shared")
-      if i + block_size > active_matrix_size:
+      if bleeding_edge_warp:
-        num_tile_elements = block_size * block_size
+        wp.tile_lower_solve_inplace(L_tile, rhs_view)
-        num_iterations = (num_tile_elements + num_threads_per_block - 1) // num_threads_per_block
+      else:
-
+        rhs_tmp = wp.tile_lower_solve(L_tile, rhs_view)
-        for ii in range(num_iterations):
+        wp.tile_assign(rhs_tile, rhs_tmp, offset=(i, 0))
-          linear_index = tid_block + ii * num_threads_per_block
-          linear_index = linear_index % num_tile_elements
-          row = linear_index // block_size
-          col = linear_index % block_size
-          value = L_tile[row, col]
-          if i + row >= active_matrix_size or i + col >= active_matrix_size:
-            value = wp.where(row == col, float(1), float(0))
-          L_tile[row, col] = value
-
-        # Handle rhs
-        num_tile_elements = block_size
-        num_iterations = (num_tile_elements + num_threads_per_block - 1) // num_threads_per_block
-
-        for ii in range(num_iterations):
-          linear_index = tid_block + ii * num_threads_per_block
-          linear_index = linear_index % num_tile_elements
-          value = rhs_tile[linear_index, 0]
-          if i + linear_index >= active_matrix_size:
-            value = float(0)
-          rhs_tile[linear_index, 0] = value
-
-      y_tile = wp.tile_lower_solve(L_tile, rhs_tile)
-      wp.tile_store(tmp, y_tile, offset=(i, 0))
 
     # Backward substitution: solve L^T x = y
-    for i in range(n - block_size, -1, -block_size):
+    for i in range(matrix_size - block_size, -1, -block_size):
       i_end = i + block_size
-      rhs_tile = wp.tile_load(tmp, shape=(block_size, 1), offset=(i, 0))
+      tmp_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(i, 0))
-      if i_end < n:
+      for j in range(i_end, matrix_size, block_size):
-        for j in range(i_end, n, block_size):
+        L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(j, i), storage="shared")
-          L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(j, i))
+        x_tile = wp.tile_load(x, shape=(block_size, 1), offset=(j, 0), storage="shared", bounds_check=False)
-          L_T_tile = wp.tile_transpose(L_tile)
+        if bleeding_edge_warp:
-          x_tile = wp.tile_load(x, shape=(block_size, 1), offset=(j, 0))
+          wp.tile_matmul(wp.tile_transpose(L_tile), x_tile, tmp_tile, alpha=-1.0)
-          L_T_x_tile = wp.tile_matmul(L_T_tile, x_tile)
+        else:
-          rhs_tile -= L_T_x_tile
+          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))
+      L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, i), storage="shared")
 
-      # The following if pads the matrix if it is not divisible by block_size
+      if bleeding_edge_warp:
-      if i + block_size > active_matrix_size:
+        wp.tile_upper_solve_inplace(wp.tile_transpose(L_tile), tmp_tile)
-        num_tile_elements = block_size * block_size
+      else:
-        num_iterations = (num_tile_elements + num_threads_per_block - 1) // num_threads_per_block
+        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)
-        for ii in range(num_iterations):
-          linear_index = tid_block + ii * num_threads_per_block
-          linear_index = linear_index % num_tile_elements
-          row = linear_index // block_size
-          col = linear_index % block_size
-          value = L_tile[row, col]
-          if i + row >= active_matrix_size or i + col >= active_matrix_size:
-            value = wp.where(row == col, float(1), float(0))
-          L_tile[row, col] = value
-
-        # Handle rhs
-        num_tile_elements = block_size
-        num_iterations = (num_tile_elements + num_threads_per_block - 1) // num_threads_per_block
-
-        for ii in range(num_iterations):
-          linear_index = tid_block + ii * num_threads_per_block
-          linear_index = linear_index % num_tile_elements
-          value = rhs_tile[linear_index, 0]
-          if i + linear_index >= active_matrix_size:
-            value = float(0)
-          rhs_tile[linear_index, 0] = value
-
-      x_tile = wp.tile_upper_solve(wp.tile_transpose(L_tile), rhs_tile)
-      wp.tile_store(x, x_tile, offset=(i, 0))
 
   return blocked_cholesky_solve_func

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

@@ -195,8 +195,8 @@ def _hfield_filter(
 def ccd_kernel_builder(
   geomtype1: int,
   geomtype2: int,
-  ccd_iterations: int,
+  gjk_iterations: int,
-  is_hfield: bool,
+  epa_iterations: int,
   use_multiccd: bool,
 ):
   @wp.func
@@ -274,7 +274,8 @@ def ccd_kernel_builder(
     dist, ncontact, w1, w2, idx = ccd(
       opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
       cutoff,
-      ccd_iterations,
+      gjk_iterations,
+      epa_iterations,
       geom1,
       geom2,
       geomtype1,
@@ -484,7 +485,7 @@ def ccd_kernel_builder(
     worldid = collision_worldid_in[tid]
 
     # height field filter
-    if wp.static(is_hfield):
+    if wp.static(geomtype1 == GeomType.HFIELD.value):
       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
       )
@@ -538,7 +539,7 @@ def ccd_kernel_builder(
     )
 
     # see MuJoCo mjc_ConvexHField
-    if wp.static(is_hfield):
+    if wp.static(geomtype1 == GeomType.HFIELD.value):
       geom1_dataid = geom_dataid[g1]
 
       # height field subgrid
@@ -645,7 +646,8 @@ def ccd_kernel_builder(
             dist, ncontact, w1, w2, idx = ccd(
               opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
               0.0,
-              ccd_iterations,
+              gjk_iterations,
+              epa_iterations,
               geom1,
               geom2,
               geomtype1,
@@ -967,34 +969,41 @@ def convex_narrowphase(m: Model, d: Data):
   kernel for each type of convex collision pair present in the model, avoiding unnecessary
   computations for non-existent pair types.
   """
-  if not any(m.geom_pair_type_count[upper_trid_index(len(GeomType), g[0].value, g[1].value)] for g in _CONVEX_COLLISION_PAIRS):
+
+  def _pair_count(p1: int, p2: int) -> int:
+    return m.geom_pair_type_count[upper_trid_index(len(GeomType), p1, p2)]
+
+  # no convex collisions, early return
+  if not any(_pair_count(g[0].value, g[1].value) for g in _CONVEX_COLLISION_PAIRS):
     return
 
+  epa_iterations = m.opt.ccd_iterations
+
   # set to true to enable multiccd
   use_multiccd = False
   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 + m.opt.ccd_iterations), dtype=wp.vec3)
+  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 + m.opt.ccd_iterations), dtype=wp.vec3)
+  epa_vert1 = wp.empty(shape=(d.naconmax, 5 + epa_iterations), dtype=wp.vec3)
   # epa_vert2: vertices in EPA polytope in geom 2 space
-  epa_vert2 = wp.empty(shape=(d.naconmax, 5 + m.opt.ccd_iterations), dtype=wp.vec3)
+  epa_vert2 = wp.empty(shape=(d.naconmax, 5 + epa_iterations), dtype=wp.vec3)
   # epa_vert_index1: vertex indices in EPA polytope for geom 1
-  epa_vert_index1 = wp.empty(shape=(d.naconmax, 5 + m.opt.ccd_iterations), dtype=int)
+  epa_vert_index1 = wp.empty(shape=(d.naconmax, 5 + epa_iterations), dtype=int)
   # epa_vert_index2: vertex indices in EPA polytope for geom 2  (naconmax, 5 + CCDiter)
-  epa_vert_index2 = wp.empty(shape=(d.naconmax, 5 + m.opt.ccd_iterations), dtype=int)
+  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 * m.opt.ccd_iterations), dtype=wp.vec3i)
+  epa_face = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=wp.vec3i)
   # epa_pr: projection of origin on polytope faces
-  epa_pr = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * m.opt.ccd_iterations), dtype=wp.vec3)
+  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 * m.opt.ccd_iterations), dtype=float)
+  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 * m.opt.ccd_iterations), dtype=int)
+  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 * m.opt.ccd_iterations), dtype=int)
+  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)
   # multiccd_polygon: clipped contact surface
@@ -1023,9 +1032,9 @@ def convex_narrowphase(m: Model, d: Data):
   for geom_pair in _CONVEX_COLLISION_PAIRS:
     g1 = geom_pair[0].value
     g2 = geom_pair[1].value
-    if m.geom_pair_type_count[upper_trid_index(len(GeomType), g1, g2)]:
+    if _pair_count(g1, g2):
       wp.launch(
-        ccd_kernel_builder(g1, g2, m.opt.ccd_iterations, g1 == GeomType.HFIELD, use_multiccd),
+        ccd_kernel_builder(g1, g2, m.opt.ccd_iterations, epa_iterations, use_multiccd),
         dim=d.naconmax,
         inputs=[
           m.opt.ccd_tolerance,

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

@@ -18,7 +18,6 @@ from typing import Tuple
 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
@@ -28,7 +27,8 @@ wp.set_module_options({"enable_backward": False})
 
 FLOAT_MIN = -1e30
 FLOAT_MAX = 1e30
-MJ_MINVAL2 = MJ_MINVAL * MJ_MINVAL
+MINVAL = 1e-15
+MIN_DIST = 1e-10
 
 # TODO(kbayes): write out formulas to derive these constants
 FACE_TOL = 0.99999872
@@ -122,7 +122,7 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint:
     res = wp.vec3(0.0, 0.0, 0.0)
     # set result in XY plane: support on circle
     d = wp.sqrt(local_dir[0] * local_dir[0] + local_dir[1] * local_dir[1])
-    if d > MJ_MINVAL:
+    if d > MINVAL:
       scl = geom.size[0] / d
       res[0] = local_dir[0] * scl
       res[1] = local_dir[1] * scl
@@ -173,6 +173,8 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint:
     sp.point = geom.rot @ sp.point + geom.pos
   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)
     for i in range(6):
       vert = geom.hfprism[i]
       dist = wp.dot(vert, local_dir)
@@ -241,7 +243,7 @@ def _linear_combine(n: int, coefs: wp.vec4, mat: mat43) -> wp.vec3:
 
 @wp.func
 def _almost_equal(v1: wp.vec3, v2: wp.vec3) -> bool:
-  return wp.abs(v1[0] - v2[0]) < MJ_MINVAL and wp.abs(v1[1] - v2[1]) < MJ_MINVAL and wp.abs(v1[2] - v2[2]) < MJ_MINVAL
+  return wp.abs(v1[0] - v2[0]) < MINVAL and wp.abs(v1[1] - v2[1]) < MINVAL and wp.abs(v1[2] - v2[2]) < MINVAL
 
 
 @wp.func
@@ -291,7 +293,7 @@ def _project_origin_plane(v1: wp.vec3, v2: wp.vec3, v3: wp.vec3) -> Tuple[wp.vec
   nn = wp.dot(n, n)
   if nn == 0:
     return z, 1
-  if nv != 0 and nn > MJ_MINVAL:
+  if nv != 0 and nn > MINVAL:
     v = (nv / nn) * n
     return v, 0
 
@@ -301,7 +303,7 @@ def _project_origin_plane(v1: wp.vec3, v2: wp.vec3, v3: wp.vec3) -> Tuple[wp.vec
   nn = wp.dot(n, n)
   if nn == 0:
     return z, 1
-  if nv != 0 and nn > MJ_MINVAL:
+  if nv != 0 and nn > MINVAL:
     v = (nv / nn) * n
     return v, 0
 
@@ -579,13 +581,14 @@ def gjk(
   simplex_index1 = wp.vec4i()
   simplex_index2 = wp.vec4i()
   n = int(0)
+  cnt = int(1)
   coordinates = wp.vec4()  # barycentric coordinates
   epsilon = wp.where(is_discrete, 0.0, 0.5 * tolerance * tolerance)
 
   # set initial guess
   x_k = x1_0 - x2_0
 
-  for k in range(gjk_iterations):
+  for _ in range(gjk_iterations):
     xnorm = wp.dot(x_k, x_k)
     # TODO(kbayes): determine new constant here
     if xnorm < 1e-12:
@@ -663,6 +666,11 @@ def gjk(
     if n == 4:
       break
 
+    cnt += 1
+
+  if cnt == gjk_iterations:
+    wp.printf("Warning: opt.ccd_iterations, currently set to %d, needs to be increased.\n", gjk_iterations)
+
   result = GJKResult()
 
   # compute the approximate witness points
@@ -751,7 +759,7 @@ def _tri_point_intersect(v1: wp.vec3, v2: wp.vec3, v3: wp.vec3, p: wp.vec3) -> b
   pr[0] = v1[0] * l1 + v2[0] * l2 + v3[0] * l3
   pr[1] = v1[1] * l1 + v2[1] * l2 + v3[1] * l3
   pr[2] = v1[2] * l1 + v2[2] * l2 + v3[2] * l3
-  return wp.norm_l2(pr - p) < MJ_MINVAL
+  return wp.norm_l2(pr - p) < MINVAL
 
 
 @wp.func
@@ -867,36 +875,66 @@ def _delete_face(pt: Polytope, face_id: int) -> int:
 
 
 @wp.func
-def _epa_witness(pt: Polytope, face_idx: int) -> Tuple[wp.vec3, wp.vec3]:
+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]
+
   # compute affine coordinates for witness points on plane defined by face
-  v1 = pt.vert[pt.face[face_idx][0]]
+  v1 = pt.vert[face[0]]
-  v2 = pt.vert[pt.face[face_idx][1]]
+  v2 = pt.vert[face[1]]
-  v3 = pt.vert[pt.face[face_idx][2]]
+  v3 = pt.vert[face[2]]
 
   coordinates = _tri_affine_coord(v1, v2, v3, pt.face_pr[face_idx])
   l1 = coordinates[0]
   l2 = coordinates[1]
   l3 = coordinates[2]
 
-  # face on geom 1
-  v1 = pt.vert1[pt.face[face_idx][0]]
-  v2 = pt.vert1[pt.face[face_idx][1]]
-  v3 = pt.vert1[pt.face[face_idx][2]]
-  x1 = wp.vec3()
-  x1[0] = v1[0] * l1 + v2[0] * l2 + v3[0] * l3
-  x1[1] = v1[1] * l1 + v2[1] * l2 + v3[1] * l3
-  x1[2] = v1[2] * l1 + v2[2] * l2 + v3[2] * l3
-
   # face on geom 2
-  v1 = pt.vert2[pt.face[face_idx][0]]
+  v1 = pt.vert2[face[0]]
-  v2 = pt.vert2[pt.face[face_idx][1]]
+  v2 = pt.vert2[face[1]]
-  v3 = pt.vert2[pt.face[face_idx][2]]
+  v3 = pt.vert2[face[2]]
   x2 = wp.vec3()
   x2[0] = v1[0] * l1 + v2[0] * l2 + v3[0] * l3
   x2[1] = v1[1] * l1 + v2[1] * l2 + v3[1] * l3
   x2[2] = v1[2] * l1 + v2[2] * l2 + v3[2] * l3
 
-  return x1, x2
+  # correct witness points for hfield geoms
+  i1 = pt.vert_index1[face[0]]
+  i2 = pt.vert_index1[face[1]]
+  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]
+    a = geom1.hfprism[3]
+    b = geom1.hfprism[4]
+    c = geom1.hfprism[5]
+    x2 = wp.normalize(x2)
+
+    # TODO(kbayes): Support cases where geom2 is larger than the height field
+    sp = _support(geom2, geomtype2, x2)
+    x2 = sp.point
+
+    coordinates2 = _tri_affine_coord(a, b, c, x2)
+    if coordinates2[0] > 0 and coordinates2[1] > 0 and coordinates2[2] > 0:
+      x1 = coordinates[0] * a + coordinates[1] * b + coordinates[2] * c
+    else:
+      p = c
+      p = wp.where(coordinates[1] > 0, b, p)
+      p = wp.where(coordinates[0] > 0, a, p)
+      x1 = x2 - wp.dot(x2 - p, n) * n
+    return x1, x2, -wp.norm_l2(x1 - x2)
+
+  # face on geom 1
+  v1 = pt.vert1[face[0]]
+  v2 = pt.vert1[face[1]]
+  v3 = pt.vert1[face[2]]
+  x1 = wp.vec3()
+  x1[0] = v1[0] * l1 + v2[0] * l2 + v3[0] * l3
+  x1[1] = v1[1] * l1 + v2[1] * l2 + v3[1] * l3
+  x1[2] = v1[2] * l1 + v2[2] * l2 + v3[2] * l3
+
+  return x1, x2, -wp.sqrt(pt.face_norm2[face_idx])
 
 
 @wp.func
@@ -926,7 +964,7 @@ def _polytope2(
       index = i
 
   # cross product with best coordinate axis
-  e = wp.vec(0.0, 0.0, 0.0)
+  e = wp.vec3(0.0, 0.0, 0.0)
   e[index] = 1.0
   d1 = wp.cross(e, diff)
 
@@ -956,27 +994,27 @@ def _polytope2(
   _epa_support(pt, 4, geom1, geom2, geomtype1, geomtype2, d3 / wp.norm_l2(d3))
 
   # build hexahedron
-  if _attach_face(pt, 0, 0, 2, 3) < MJ_MINVAL:
+  if _attach_face(pt, 0, 0, 2, 3) < MIN_DIST:
     pt.status = -1
     return pt, _replace_simplex3(pt, 0, 2, 3)
 
-  if _attach_face(pt, 1, 0, 4, 2) < MJ_MINVAL2:
+  if _attach_face(pt, 1, 0, 4, 2) < MIN_DIST:
     pt.status = -1
     return pt, _replace_simplex3(pt, 0, 4, 2)
 
-  if _attach_face(pt, 2, 0, 3, 4) < MJ_MINVAL2:
+  if _attach_face(pt, 2, 0, 3, 4) < MIN_DIST:
     pt.status = -1
     return pt, _replace_simplex3(pt, 0, 3, 4)
 
-  if _attach_face(pt, 3, 1, 3, 2) < MJ_MINVAL2:
+  if _attach_face(pt, 3, 1, 3, 2) < MIN_DIST:
     pt.status = -1
     return pt, _replace_simplex3(pt, 1, 3, 2)
 
-  if _attach_face(pt, 4, 1, 2, 4) < MJ_MINVAL2:
+  if _attach_face(pt, 4, 1, 2, 4) < MIN_DIST:
     pt.status = -1
     return pt, _replace_simplex3(pt, 1, 2, 4)
 
-  if _attach_face(pt, 5, 1, 4, 3) < MJ_MINVAL2:
+  if _attach_face(pt, 5, 1, 4, 3) < MIN_DIST:
     pt.status = -1
     return pt, _replace_simplex3(pt, 1, 4, 3)
 
@@ -1016,7 +1054,7 @@ def _polytope3(
   """Create polytope for EPA given a 2-simplex from GJK."""
   # get normals in both directions
   n = wp.cross(simplex[1] - simplex[0], simplex[2] - simplex[0])
-  if wp.norm_l2(n) < MJ_MINVAL:
+  if wp.norm_l2(n) < MINVAL:
     pt.status = 2
     return pt
 
@@ -1066,22 +1104,22 @@ def _polytope3(
     return pt
 
   # create hexahedron for EPA
-  if _attach_face(pt, 0, 4, 0, 1) < MJ_MINVAL2:
+  if _attach_face(pt, 0, 4, 0, 1) < MIN_DIST:
     pt.status = 6
     return pt
-  if _attach_face(pt, 1, 4, 2, 0) < MJ_MINVAL2:
+  if _attach_face(pt, 1, 4, 2, 0) < MIN_DIST:
     pt.status = 7
     return pt
-  if _attach_face(pt, 2, 4, 1, 2) < MJ_MINVAL2:
+  if _attach_face(pt, 2, 4, 1, 2) < MIN_DIST:
     pt.status = 8
     return pt
-  if _attach_face(pt, 3, 3, 1, 0) < MJ_MINVAL2:
+  if _attach_face(pt, 3, 3, 1, 0) < MIN_DIST:
     pt.status = 9
     return pt
-  if _attach_face(pt, 4, 3, 0, 2) < MJ_MINVAL2:
+  if _attach_face(pt, 4, 3, 0, 2) < MIN_DIST:
     pt.status = 10
     return pt
-  if _attach_face(pt, 5, 3, 2, 1) < MJ_MINVAL2:
+  if _attach_face(pt, 5, 3, 2, 1) < MIN_DIST:
     pt.status = 11
     return pt
 
@@ -1140,19 +1178,19 @@ def _polytope4(
   pt.vert_index2[3] = simplex_index2[3]
 
   # if the origin is on a face, replace the 3-simplex with a 2-simplex
-  if _attach_face(pt, 0, 0, 1, 2) < MJ_MINVAL2:
+  if _attach_face(pt, 0, 0, 1, 2) < MIN_DIST:
     pt.status = -1
     return pt, _replace_simplex3(pt, 0, 1, 2)
 
-  if _attach_face(pt, 1, 0, 3, 1) < MJ_MINVAL2:
+  if _attach_face(pt, 1, 0, 3, 1) < MIN_DIST:
     pt.status = -1
     return pt, _replace_simplex3(pt, 0, 3, 1)
 
-  if _attach_face(pt, 2, 0, 2, 3) < MJ_MINVAL2:
+  if _attach_face(pt, 2, 0, 2, 3) < MIN_DIST:
     pt.status = -1
     return pt, _replace_simplex3(pt, 0, 2, 3)
 
-  if _attach_face(pt, 3, 3, 2, 1) < MJ_MINVAL2:
+  if _attach_face(pt, 3, 3, 2, 1) < MIN_DIST:
     pt.status = -1
     return pt, _replace_simplex3(pt, 3, 2, 1)
 
@@ -1177,6 +1215,7 @@ def _polytope4(
 def _epa(
   # In:
   tolerance: float,
+  gjk_iterations: int,
   epa_iterations: int,
   pt: Polytope,
   geom1: Geom,
@@ -1191,9 +1230,10 @@ def _epa(
   idx = int(-1)
   pidx = int(-1)
   epsilon = wp.where(is_discrete, 1e-15, tolerance)
+  cnt = int(1)
 
-  for k in range(epa_iterations):
+  for _ in range(epa_iterations):
-    pidx = int(idx)
+    pidx = idx
     idx = int(-1)
 
     # find the face closest to the origin (lower bound for penetration depth)
@@ -1283,11 +1323,15 @@ def _epa(
 
     # clear horizon
     pt.nhorizon = 0
+    cnt += 1
+
+  if cnt == epa_iterations:
+    wp.printf("Warning: opt.ccd_iterations, currently set to %d, needs to be increased.\n", gjk_iterations)
 
   # return from valid face
   if idx > -1:
-    x1, x2 = _epa_witness(pt, idx)
+    x1, x2, dist = _epa_witness(pt, geom1, geom2, geomtype1, geomtype2, idx)
-    return -wp.sqrt(pt.face_norm2[idx]), x1, x2, idx
+    return dist, x1, x2, idx
   return 0.0, wp.vec3(), wp.vec3(), -1
 
 
@@ -1696,40 +1740,40 @@ def _box_edge_normals(
 def _box_face(mat: wp.mat33, pos: wp.vec3, size: wp.vec3, idx: int, face_out: wp.array(dtype=wp.vec3)) -> int:
   # compute global coordinates of the box face and face normal
   if idx == 0:  # right
-    face_out[0] = mat @ wp.vec(size[0], size[1], size[2]) + pos
+    face_out[0] = mat @ wp.vec3(size[0], size[1], size[2]) + pos
-    face_out[1] = mat @ wp.vec(size[0], size[1], -size[2]) + pos
+    face_out[1] = mat @ wp.vec3(size[0], size[1], -size[2]) + pos
-    face_out[2] = mat @ wp.vec(size[0], -size[1], -size[2]) + pos
+    face_out[2] = mat @ wp.vec3(size[0], -size[1], -size[2]) + pos
-    face_out[3] = mat @ wp.vec(size[0], -size[1], size[2]) + pos
+    face_out[3] = mat @ wp.vec3(size[0], -size[1], size[2]) + pos
     return 4
   if idx == 1:  # left
-    face_out[0] = mat @ wp.vec(-size[0], size[1], -size[2]) + pos
+    face_out[0] = mat @ wp.vec3(-size[0], size[1], -size[2]) + pos
-    face_out[1] = mat @ wp.vec(-size[0], size[1], size[2]) + pos
+    face_out[1] = mat @ wp.vec3(-size[0], size[1], size[2]) + pos
-    face_out[2] = mat @ wp.vec(-size[0], -size[1], size[2]) + pos
+    face_out[2] = mat @ wp.vec3(-size[0], -size[1], size[2]) + pos
-    face_out[3] = mat @ wp.vec(-size[0], -size[1], -size[2]) + pos
+    face_out[3] = mat @ wp.vec3(-size[0], -size[1], -size[2]) + pos
     return 4
   if idx == 2:  # top
-    face_out[0] = mat @ wp.vec(-size[0], size[1], -size[2]) + pos
+    face_out[0] = mat @ wp.vec3(-size[0], size[1], -size[2]) + pos
-    face_out[1] = mat @ wp.vec(size[0], size[1], -size[2]) + pos
+    face_out[1] = mat @ wp.vec3(size[0], size[1], -size[2]) + pos
-    face_out[2] = mat @ wp.vec(size[0], size[1], size[2]) + pos
+    face_out[2] = mat @ wp.vec3(size[0], size[1], size[2]) + pos
-    face_out[3] = mat @ wp.vec(-size[0], size[1], size[2]) + pos
+    face_out[3] = mat @ wp.vec3(-size[0], size[1], size[2]) + pos
     return 4
   if idx == 3:  # bottom
-    face_out[0] = mat @ wp.vec(-size[0], -size[1], size[2]) + pos
+    face_out[0] = mat @ wp.vec3(-size[0], -size[1], size[2]) + pos
-    face_out[1] = mat @ wp.vec(size[0], -size[1], size[2]) + pos
+    face_out[1] = mat @ wp.vec3(size[0], -size[1], size[2]) + pos
-    face_out[2] = mat @ wp.vec(size[0], -size[1], -size[2]) + pos
+    face_out[2] = mat @ wp.vec3(size[0], -size[1], -size[2]) + pos
-    face_out[3] = mat @ wp.vec(-size[0], -size[1], -size[2]) + pos
+    face_out[3] = mat @ wp.vec3(-size[0], -size[1], -size[2]) + pos
     return 4
   if idx == 4:  # front
-    face_out[0] = mat @ wp.vec(-size[0], size[1], size[2]) + pos
+    face_out[0] = mat @ wp.vec3(-size[0], size[1], size[2]) + pos
-    face_out[1] = mat @ wp.vec(size[0], size[1], size[2]) + pos
+    face_out[1] = mat @ wp.vec3(size[0], size[1], size[2]) + pos
-    face_out[2] = mat @ wp.vec(size[0], -size[1], size[2]) + pos
+    face_out[2] = mat @ wp.vec3(size[0], -size[1], size[2]) + pos
-    face_out[3] = mat @ wp.vec(-size[0], -size[1], size[2]) + pos
+    face_out[3] = mat @ wp.vec3(-size[0], -size[1], size[2]) + pos
     return 4
   if idx == 5:  # back
-    face_out[0] = mat @ wp.vec(size[0], size[1], -size[2]) + pos
+    face_out[0] = mat @ wp.vec3(size[0], size[1], -size[2]) + pos
-    face_out[1] = mat @ wp.vec(-size[0], size[1], -size[2]) + pos
+    face_out[1] = mat @ wp.vec3(-size[0], size[1], -size[2]) + pos
-    face_out[2] = mat @ wp.vec(-size[0], -size[1], -size[2]) + pos
+    face_out[2] = mat @ wp.vec3(-size[0], -size[1], -size[2]) + pos
-    face_out[3] = mat @ wp.vec(size[0], -size[1], -size[2]) + pos
+    face_out[3] = mat @ wp.vec3(size[0], -size[1], -size[2]) + pos
     return 4
   return 0
 
@@ -1791,7 +1835,6 @@ def _plane_intersect(pn: wp.vec3, pd: float, a: wp.vec3, b: wp.vec3) -> Tuple[fl
 @wp.func
 def _polygon_clip(
   # In:
-  prune: bool,
   plane_normal: wp.array(dtype=wp.vec3),
   plane_dist: wp.array(dtype=float),
   face1: wp.array(dtype=wp.vec3),
@@ -1871,7 +1914,7 @@ def _polygon_clip(
   if npolygon < 1:
     return 0, witness1, witness2
 
-  if prune and npolygon > 4:
+  if npolygon > 4:
     quad = _polygon_quad(polygon_out, npolygon)
     for i in range(4):
       witness2[i] = polygon_out[quad[i]]
@@ -2102,27 +2145,20 @@ def multicontact(
         face2,
       )
 
-  # TODO(kbayes): this approximates the contact direction, by scaling the face normal by the
-  # single contact direction's magnitude. This is effective, but polygonClip should compute
-  # this for each contact point.
-  approx_dir = wp.vec3()
-
   # face1 is an edge; clip face1 against face2
   if is_edge_contact_geom1:
     approx_dir = wp.norm_l2(dir) * n2[j]
-    return _polygon_clip(False, plane_normal, plane_dist, face2, nface2, face1, nface1, n2[j], approx_dir, polygon, clipped)
+    return _polygon_clip(plane_normal, plane_dist, face2, nface2, face1, nface1, n2[j], approx_dir, polygon, clipped)
 
   # face2 is an edge; clip face2 against face1
   if is_edge_contact_geom2:
     approx_dir = -wp.norm_l2(dir) * n1[j]
-    return _polygon_clip(False, plane_normal, plane_dist, face1, nface1, face2, nface2, n1[j], approx_dir, polygon, clipped)
+    return _polygon_clip(plane_normal, plane_dist, face1, nface1, face2, nface2, n1[j], approx_dir, polygon, clipped)
 
   # face-face collision
   approx_dir = wp.norm_l2(dir) * n2[j]
 
-  # don't prune box-box collisions (expect up to 8 contacts)
+  return _polygon_clip(plane_normal, plane_dist, face1, nface1, face2, nface2, n1[i], approx_dir, polygon, clipped)
-  prune = not (geomtype1 == GeomType.BOX and geomtype2 == GeomType.BOX)
-  return _polygon_clip(prune, plane_normal, plane_dist, face1, nface1, face2, nface2, n1[i], approx_dir, polygon, clipped)
 
 
 @wp.func
@@ -2142,7 +2178,8 @@ def ccd(
   # In:
   tolerance: float,
   cutoff: float,
-  ccd_iterations: int,
+  gjk_iterations: int,
+  epa_iterations: int,
   geom1: Geom,
   geom2: Geom,
   geomtype1: int,
@@ -2176,7 +2213,8 @@ def ccd(
     geom1.margin = 0.0
     geom1.size = wp.vec3(0.0, geom1.size[1], geom1.size[2])
 
-  if geomtype2 == GeomType.SPHERE or geomtype2 == GeomType.CAPSULE:
+  # TODO(kbayes): support gjk margin trick with height fields
+  if geomtype1 != GeomType.HFIELD and (geomtype2 == GeomType.SPHERE or geomtype2 == GeomType.CAPSULE):
     size2 = geom2.size[0]
     full_margin2 = size2 + 0.5 * geom2.margin
     geom2.margin = 0.0
@@ -2185,7 +2223,7 @@ def ccd(
   # special handling for sphere and capsule (shrink to point and line respectively)
   if size1 + size2 > 0.0:
     cutoff += full_margin1 + full_margin2
-    result = gjk(tolerance, ccd_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, is_discrete)
+    result = gjk(tolerance, gjk_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, is_discrete)
 
     # shallow penetration, inflate contact
     if result.dist > tolerance:
@@ -2201,7 +2239,7 @@ def ccd(
     geom2.size = wp.vec3(size2, geom2.size[1], geom2.size[2])
     cutoff -= full_margin1 + full_margin2
 
-  result = gjk(tolerance, ccd_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, is_discrete)
+  result = gjk(tolerance, gjk_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, is_discrete)
 
   # no penetration depth to recover
   if result.dist > tolerance or result.dim < 2:
@@ -2287,7 +2325,7 @@ def ccd(
   if pt.status:
     return result.dist, 1, result.x1, result.x2, -1
 
-  dist, x1, x2, idx = _epa(tolerance, ccd_iterations, pt, geom1, geom2, geomtype1, geomtype2, is_discrete)
+  dist, x1, x2, idx = _epa(tolerance, gjk_iterations, epa_iterations, pt, geom1, geom2, geomtype1, geomtype2, is_discrete)
   if idx == -1:
     return FLOAT_MAX, 0, wp.vec3(), wp.vec3(), -1
   return dist, 1, x1, x2, idx

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

@@ -1035,7 +1035,7 @@ def plane_box_wrapper(
   dist, pos, normal = plane_box(plane.normal, plane.pos, box.pos, box.rot, box.size)
   frame = make_frame(normal)
 
-  for i in range(4):
+  for i in range(8):
     write_contact(
       naconmax_in,
       i,
@@ -1562,16 +1562,9 @@ assert _check_primitive_collisions(), "_PRIMITIVE_COLLISIONS is in invalid order
 
 
 @cache_kernel
-def _create_narrowphase_kernel(primitive_collisions_types, primitive_collisions_func):
+def _primitive_narrowphase(primitive_collisions_types, primitive_collisions_func):
-  # AD: no unique here:
+  @nested_kernel(module="unique", enable_backward=False)
-  # * we expect this generator to be called only once per model, so no repeated compilation
+  def primitive_narrowphase(
-  # * module="unique" is generating problems because it uses the function name as the key
-  #   that in turn will cause multiple kernels to be generated with the same name
-  #   this is mostly problematic in cases like the UTs where we don't clear the kernel cache
-  #   between different tests.
-
-  @nested_kernel(enable_backward=False)
-  def _primitive_narrowphase(
     # Model:
     geom_type: wp.array(dtype=int),
     geom_condim: wp.array(dtype=int),
@@ -1719,20 +1712,11 @@ def _create_narrowphase_kernel(primitive_collisions_types, primitive_collisions_
           nacon_out,
         )
 
-  return _primitive_narrowphase
+  return primitive_narrowphase
 
 
-def _primitive_narrowphase_builder(m: Model):
+_PRIMITIVE_COLLISION_TYPES = []
-  _primitive_collisions_types = []
+_PRIMITIVE_COLLISION_FUNC = []
-  _primitive_collisions_func = []
-
-  for types, func in _PRIMITIVE_COLLISIONS.items():
-    idx = upper_trid_index(len(GeomType), types[0].value, types[1].value)
-    if m.geom_pair_type_count[idx] and types not in _primitive_collisions_types:
-      _primitive_collisions_types.append(types)
-      _primitive_collisions_func.append(func)
-
-  return _create_narrowphase_kernel(_primitive_collisions_types, _primitive_collisions_func)
 
 
 @event_scope
@@ -1751,10 +1735,17 @@ def primitive_narrowphase(m: Model, d: Data):
   the specific primitive collision types present in the model, avoiding
   unnecessary checks for non-existent collision pairs.
   """
-  # we need to figure out how to keep the overhead of this small - not launching anything
+  # TODO(team): keep the overhead of this small - not launching anything
   # for pair types without collisions, as well as updating the launch dimensions.
+
+  for types, func in _PRIMITIVE_COLLISIONS.items():
+    idx = upper_trid_index(len(GeomType), types[0].value, types[1].value)
+    if m.geom_pair_type_count[idx] and types not in _PRIMITIVE_COLLISION_TYPES:
+      _PRIMITIVE_COLLISION_TYPES.append(types)
+      _PRIMITIVE_COLLISION_FUNC.append(func)
+
   wp.launch(
-    _primitive_narrowphase_builder(m),
+    _primitive_narrowphase(_PRIMITIVE_COLLISION_TYPES, _PRIMITIVE_COLLISION_FUNC),
     dim=d.naconmax,
     inputs=[
       m.geom_type,

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

@@ -316,7 +316,7 @@ def plane_box(
   box_pos: wp.vec3,
   box_rot: wp.mat33,
   box_size: wp.vec3,
-) -> Tuple[wp.vec4, mat43f, wp.vec3]:
+) -> Tuple[vec8f, mat83f, wp.vec3]:
   """Core contact geometry calculation for plane-box collision.
 
   Args:
@@ -325,42 +325,36 @@ def plane_box(
     box_pos: Center position of the box.
     box_rot: Rotation matrix of the box.
     box_size: Half-extents of the box along each axis.
+    margin: Collision tolerance.
 
   Returns:
     - Vector of contact distances (wp.inf for unpopulated contacts).
     - Matrix of contact positions (one per row).
     - Contact normal vector.
   """
-  corner = wp.vec3()
   center_dist = wp.dot(box_pos - plane_pos, plane_normal)
 
-  dist = wp.vec4(wp.inf)
+  dist = vec8f(wp.inf)
-  pos = mat43f()
+  pos = mat83f()
 
   # test all corners, pick bottom 4
-  ncontact = int(0)
   for i in range(8):
     # get corner in local coordinates
-    corner.x = wp.where(i & 1, box_size.x, -box_size.x)
+    corner = wp.vec3(
-    corner.y = wp.where(i & 2, box_size.y, -box_size.y)
+      wp.where(i & 1, box_size[0], -box_size[0]),
-    corner.z = wp.where(i & 4, box_size.z, -box_size.z)
+      wp.where(i & 2, box_size[1], -box_size[1]),
+      wp.where(i & 4, box_size[2], -box_size[2]),
+    )
 
     # get corner in global coordinates relative to box center
     corner = box_rot * corner
 
-    # compute distance to plane, skip if too far or pointing up
+    # compute distance to plane
     ldist = wp.dot(plane_normal, corner)
-    if center_dist + ldist > 0 or ldist > 0:
-      continue
-
     cdist = center_dist + ldist
 
-    dist[ncontact] = cdist
+    dist[i] = cdist
-    pos[ncontact] = corner + box_pos - 0.5 * plane_normal * cdist
+    pos[i] = corner + box_pos - 0.5 * plane_normal * cdist
-    ncontact += 1
-
-    if ncontact >= 4:
-      break
 
   return dist, pos, plane_normal
 
@@ -586,8 +580,7 @@ def box_box(
     box2_pos: Center position of the second box.
     box2_rot: Rotation matrix of the second box.
     box2_size: Half-extents of the second box along each axis.
-    margin: Distance threshold for early contact generation (default: 0.0).
+    margin: Collision tolerance.
-            When positive, contacts are generated before boxes overlap.
 
   Returns:
     - Vector of contact distances (wp.inf for unpopulated contacts).
@@ -1111,10 +1104,8 @@ def capsule_box(
   halfaxis = axis * capsule_half_length  # halfaxis is the capsule direction
   axisdir = wp.int32(halfaxis[0] > 0.0) + 2 * wp.int32(halfaxis[1] > 0.0) + 4 * wp.int32(halfaxis[2] > 0.0)
 
-  bestdistmax = 2.0 * (capsule_radius + capsule_half_length + box_size[0] + box_size[1] + box_size[2])
-
   # keep track of closest point
-  bestdist = wp.float32(bestdistmax)
+  bestdist = wp.float32(1.0e32)
   bestsegmentpos = wp.float32(-12)
 
   # cltype: encoded collision configuration
@@ -1253,7 +1244,7 @@ def capsule_box(
 
   p = wp.vec2(pos.x, pos.y)
   dd = wp.vec2(halfaxis.x, halfaxis.y)
-  s = wp.vec2(box_size.x, box_size.y)
+  s = wp.vec2(box_size[0], box_size[1])
   secondpos = wp.float32(-4.0)
 
   uu = dd.x * s.y

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

@@ -38,6 +38,7 @@ def _zero_constraint_counts(
   ne_weld_out: wp.array(dtype=int),
   ne_jnt_out: wp.array(dtype=int),
   ne_ten_out: wp.array(dtype=int),
+  ne_flex_out: wp.array(dtype=int),
 ):
   worldid = wp.tid()
 
@@ -47,6 +48,7 @@ def _zero_constraint_counts(
   ne_weld_out[worldid] = 0
   ne_jnt_out[worldid] = 0
   ne_ten_out[worldid] = 0
+  ne_flex_out[worldid] = 0
   nf_out[worldid] = 0
   nl_out[worldid] = 0
   nefc_out[worldid] = 0
@@ -485,6 +487,81 @@ def _efc_equality_tendon(
   )
 
 
+@wp.kernel
+def _efc_equality_flex(
+  # Model:
+  nv: int,
+  opt_timestep: wp.array(dtype=float),
+  flexedge_length0: wp.array(dtype=float),
+  flexedge_invweight0: wp.array(dtype=float),
+  eq_solref: wp.array2d(dtype=wp.vec2),
+  eq_solimp: wp.array2d(dtype=vec5),
+  eq_flex_adr: wp.array(dtype=int),
+  # Data in:
+  qvel_in: wp.array2d(dtype=float),
+  flexedge_J_in: wp.array3d(dtype=float),
+  flexedge_length_in: wp.array2d(dtype=float),
+  njmax_in: int,
+  # In:
+  refsafe_in: int,
+  # Data out:
+  nefc_out: wp.array(dtype=int),
+  efc_type_out: wp.array2d(dtype=int),
+  efc_id_out: wp.array2d(dtype=int),
+  efc_J_out: wp.array3d(dtype=float),
+  efc_pos_out: wp.array2d(dtype=float),
+  efc_margin_out: wp.array2d(dtype=float),
+  efc_D_out: wp.array2d(dtype=float),
+  efc_vel_out: wp.array2d(dtype=float),
+  efc_aref_out: wp.array2d(dtype=float),
+  efc_frictionloss_out: wp.array2d(dtype=float),
+  ne_flex_out: wp.array(dtype=int),
+):
+  worldid, eqflexid, edgeid = wp.tid()
+  eqid = eq_flex_adr[eqflexid]
+
+  wp.atomic_add(ne_flex_out, worldid, 1)
+  efcid = wp.atomic_add(nefc_out, worldid, 1)
+
+  if efcid >= njmax_in:
+    return
+
+  pos = flexedge_length_in[worldid, edgeid] - flexedge_length0[edgeid]
+  solref = eq_solref[worldid % eq_solref.shape[0], eqid]
+  solimp = eq_solimp[worldid % eq_solimp.shape[0], eqid]
+
+  Jqvel = float(0.0)
+  for i in range(nv):
+    J = flexedge_J_in[worldid, edgeid, i]
+    efc_J_out[worldid, efcid, i] = J
+    Jqvel += J * qvel_in[worldid, i]
+
+  _update_efc_row(
+    worldid,
+    opt_timestep[worldid % opt_timestep.shape[0]],
+    refsafe_in,
+    efcid,
+    pos,
+    pos,
+    flexedge_invweight0[edgeid],
+    solref,
+    solimp,
+    0.0,
+    Jqvel,
+    0.0,
+    ConstraintType.EQUALITY,
+    eqid,
+    efc_type_out,
+    efc_id_out,
+    efc_pos_out,
+    efc_margin_out,
+    efc_D_out,
+    efc_vel_out,
+    efc_aref_out,
+    efc_frictionloss_out,
+  )
+
+
 @wp.kernel
 def _efc_friction_dof(
   # Model:
@@ -1135,7 +1212,7 @@ def _efc_contact_pyramidal(
   # Model:
   nv: int,
   opt_timestep: wp.array(dtype=float),
-  opt_impratio: wp.array(dtype=float),
+  opt_impratio_invsqrt: wp.array(dtype=float),
   body_parentid: wp.array(dtype=int),
   body_rootid: wp.array(dtype=int),
   body_invweight0: wp.array2d(dtype=wp.vec2),
@@ -1200,9 +1277,8 @@ def _efc_contact_pyramidal(
       contact_efc_address_out[conid, dimid] = -1
       return
 
-    opt_timestep_id = worldid % opt_timestep.shape[0]
+    timestep = opt_timestep[worldid % opt_timestep.shape[0]]
-    timestep = opt_timestep[opt_timestep_id]
+    impratio_invsqrt = opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
-    impratio = opt_impratio[opt_timestep_id]
     contact_efc_address_out[conid, dimid] = efcid
 
     geom = geom_in[conid]
@@ -1223,7 +1299,7 @@ def _efc_contact_pyramidal(
       fri0 = friction[0]
       frii = friction[dimid2 - 1]
       invweight = invweight + fri0 * fri0 * invweight
-      invweight = invweight * 2.0 * fri0 * fri0 / impratio
+      invweight = invweight * 2.0 * fri0 * fri0 * impratio_invsqrt * impratio_invsqrt
 
     Jqvel = float(0.0)
     for i in range(nv):
@@ -1306,7 +1382,7 @@ def _efc_contact_elliptic(
   # Model:
   nv: int,
   opt_timestep: wp.array(dtype=float),
-  opt_impratio: wp.array(dtype=float),
+  opt_impratio_invsqrt: wp.array(dtype=float),
   body_parentid: wp.array(dtype=int),
   body_rootid: wp.array(dtype=int),
   body_invweight0: wp.array2d(dtype=wp.vec2),
@@ -1370,9 +1446,8 @@ def _efc_contact_elliptic(
       contact_efc_address_out[conid, dimid] = -1
       return
 
-    opt_timestep_id = worldid % opt_timestep.shape[0]
+    timestep = opt_timestep[worldid % opt_timestep.shape[0]]
-    timestep = opt_timestep[opt_timestep_id]
+    impratio_invsqrt = opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
-    impratio = opt_impratio[opt_timestep_id]
     contact_efc_address_out[conid, dimid] = efcid
 
     geom = geom_in[conid]
@@ -1432,8 +1507,7 @@ def _efc_contact_elliptic(
       if solreffriction[0] or solreffriction[1]:
         ref = solreffriction
 
-      # TODO(team): precompute 1 / impratio
+      invweight = invweight * impratio_invsqrt * impratio_invsqrt
-      invweight = invweight / impratio
       friction = friction_in[conid]
 
       if dimid > 1:
@@ -1482,11 +1556,12 @@ def _num_equality(
   ne_weld_in: wp.array(dtype=int),
   ne_jnt_in: wp.array(dtype=int),
   ne_ten_in: wp.array(dtype=int),
+  ne_flex_in: wp.array(dtype=int),
   # Data out:
   ne_out: wp.array(dtype=int),
 ):
   worldid = wp.tid()
-  ne = ne_connect_in[worldid] + ne_weld_in[worldid] + ne_jnt_in[worldid] + ne_ten_in[worldid]
+  ne = ne_connect_in[worldid] + ne_weld_in[worldid] + ne_jnt_in[worldid] + ne_ten_in[worldid] + ne_flex_in[worldid]
   ne_out[worldid] = ne
 
 
@@ -1496,7 +1571,7 @@ def make_constraint(m: types.Model, d: types.Data):
   wp.launch(
     _zero_constraint_counts,
     dim=d.nworld,
-    inputs=[d.ne, d.nf, d.nl, d.nefc, d.ne_connect, d.ne_weld, d.ne_jnt, d.ne_ten],
+    inputs=[d.ne, d.nf, d.nl, d.nefc, d.ne_connect, d.ne_weld, d.ne_jnt, d.ne_ten, d.ne_flex],
   )
 
   if not (m.opt.disableflags & types.DisableBit.CONSTRAINT):
@@ -1663,10 +1738,42 @@ def make_constraint(m: types.Model, d: types.Data):
         ],
       )
 
+      wp.launch(
+        _efc_equality_flex,
+        dim=(d.nworld, m.eq_flex_adr.size, m.nflexedge),
+        inputs=[
+          m.nv,
+          m.opt.timestep,
+          m.flexedge_length0,
+          m.flexedge_invweight0,
+          m.eq_solref,
+          m.eq_solimp,
+          m.eq_flex_adr,
+          d.qvel,
+          d.flexedge_J,
+          d.flexedge_length,
+          d.njmax,
+          refsafe,
+        ],
+        outputs=[
+          d.nefc,
+          d.efc.type,
+          d.efc.id,
+          d.efc.J,
+          d.efc.pos,
+          d.efc.margin,
+          d.efc.D,
+          d.efc.vel,
+          d.efc.aref,
+          d.efc.frictionloss,
+          d.ne_flex,
+        ],
+      )
+
       wp.launch(
         _num_equality,
         dim=d.nworld,
-        inputs=[d.ne_connect, d.ne_weld, d.ne_jnt, d.ne_ten],
+        inputs=[d.ne_connect, d.ne_weld, d.ne_jnt, d.ne_ten, d.ne_flex],
         outputs=[d.ne],
       )
 
@@ -1847,7 +1954,7 @@ def make_constraint(m: types.Model, d: types.Data):
           inputs=[
             m.nv,
             m.opt.timestep,
-            m.opt.impratio,
+            m.opt.impratio_invsqrt,
             m.body_parentid,
             m.body_rootid,
             m.body_invweight0,
@@ -1892,7 +1999,7 @@ def make_constraint(m: types.Model, d: types.Data):
           inputs=[
             m.nv,
             m.opt.timestep,
-            m.opt.impratio,
+            m.opt.impratio_invsqrt,
             m.body_parentid,
             m.body_rootid,
             m.body_invweight0,

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

@@ -21,21 +21,18 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
 from mujoco.mjx.third_party.mujoco_warp._src.types import DynType
 from mujoco.mjx.third_party.mujoco_warp._src.types import GainType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
+from mujoco.mjx.third_party.mujoco_warp._src.types import TileSet
 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 nested_kernel
 
 wp.set_module_options({"enable_backward": False})
 
 
-# TODO(team): improve performance with tile operations?
 @wp.kernel
-def _qderiv_actuator_passive(
+def _qderiv_actuator_passive_vel(
   # Model:
-  nu: int,
-  opt_timestep: wp.array(dtype=float),
-  opt_disableflags: int,
-  opt_is_sparse: bool,
-  dof_damping: wp.array2d(dtype=float),
   actuator_dyntype: wp.array(dtype=int),
   actuator_gaintype: wp.array(dtype=int),
   actuator_biastype: wp.array(dtype=int),
@@ -46,9 +43,75 @@ def _qderiv_actuator_passive(
   # Data in:
   act_in: wp.array2d(dtype=float),
   ctrl_in: wp.array2d(dtype=float),
+  # Out:
+  vel_out: wp.array2d(dtype=float),
+):
+  worldid, actid = wp.tid()
+
+  actuator_gainprm_id = worldid % actuator_gainprm.shape[0]
+  actuator_biasprm_id = worldid % actuator_biasprm.shape[0]
+
+  if actuator_gaintype[actid] == GainType.AFFINE:
+    gain = actuator_gainprm[actuator_gainprm_id, actid][2]
+  else:
+    gain = 0.0
+
+  if actuator_biastype[actid] == BiasType.AFFINE:
+    bias = actuator_biasprm[actuator_biasprm_id, actid][2]
+  else:
+    bias = 0.0
+
+  if bias == 0.0 and gain == 0.0:
+    vel_out[worldid, actid] = 0.0
+    return
+
+  vel = float(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]
+
+  vel_out[worldid, actid] = vel
+
+
+@cache_kernel
+def _qderiv_actuator_passive_actuation_dense(tile: TileSet, nu: int):
+  @nested_kernel(module="unique", enable_backward=False)
+  def kernel(
+    # Data in:
+    vel_in: wp.array3d(dtype=float),
+    actuator_moment_in: wp.array3d(dtype=float),
+    # In:
+    adr: wp.array(dtype=int),
+    # Out:
+    qDeriv_out: wp.array3d(dtype=float),
+  ):
+    worldid, nodeid = wp.tid()
+    TILE_SIZE = wp.static(tile.size)
+    NU = wp.static(nu)
+
+    dofid = adr[nodeid]
+    vel_tile = wp.tile_load(vel_in[worldid], shape=(NU, 1), bounds_check=False)
+    moment_tile = wp.tile_load(actuator_moment_in[worldid], shape=(NU, TILE_SIZE), offset=(0, dofid), bounds_check=False)
+    moment_weighted = wp.tile_map(wp.mul, wp.tile_broadcast(vel_tile, shape=(NU, TILE_SIZE)), moment_tile)
+    qderiv_tile = wp.tile_matmul(wp.tile_transpose(moment_tile), moment_weighted)
+    wp.tile_store(qDeriv_out[worldid], qderiv_tile, offset=(dofid, dofid), bounds_check=False)
+
+  return kernel
+
+
+@wp.kernel
+def _qderiv_actuator_passive_actuation_sparse(
+  # Model:
+  nu: int,
+  # Data in:
   actuator_moment_in: wp.array3d(dtype=float),
-  qM_in: wp.array3d(dtype=float),
   # In:
+  vel_in: wp.array2d(dtype=float),
   qMi: wp.array(dtype=int),
   qMj: wp.array(dtype=int),
   # Out:
@@ -58,40 +121,45 @@ def _qderiv_actuator_passive(
 
   dofiid = qMi[elemid]
   dofjid = qMj[elemid]
+  qderiv_contrib = float(0.0)
+  for actid in range(nu):
+    vel = vel_in[worldid, actid]
+    if vel == 0.0:
+      continue
 
-  qderiv = float(0.0)
+    moment_i = actuator_moment_in[worldid, actid, dofiid]
-  if not opt_disableflags & DisableBit.ACTUATION:
+    moment_j = actuator_moment_in[worldid, actid, dofjid]
-    actuator_gainprm_id = worldid % actuator_gainprm.shape[0]
-    actuator_biasprm_id = worldid % actuator_biasprm.shape[0]
 
-    for actid in range(nu):
+    qderiv_contrib += moment_i * moment_j * vel
-      if actuator_gaintype[actid] == GainType.AFFINE:
-        gain = actuator_gainprm[actuator_gainprm_id, actid][2]
-      else:
-        gain = 0.0
 
-      if actuator_biastype[actid] == BiasType.AFFINE:
+  qDeriv_out[worldid, 0, elemid] = qderiv_contrib
-        bias = actuator_biasprm[actuator_biasprm_id, actid][2]
-      else:
-        bias = 0.0
 
-      if bias == 0.0 and gain == 0.0:
-        continue
 
-      vel = bias
+@wp.kernel
-      if actuator_dyntype[actid] != DynType.NONE:
+def _qderiv_actuator_passive(
-        if gain != 0.0:
+  # Model:
-          act_first = actuator_actadr[actid]
+  opt_timestep: wp.array(dtype=float),
-          act_last = act_first + actuator_actnum[actid] - 1
+  opt_disableflags: int,
-          vel += gain * act_in[worldid, act_last]
+  opt_is_sparse: bool,
-      else:
+  dof_damping: wp.array2d(dtype=float),
-        if gain != 0.0:
+  # Data in:
-          vel += gain * ctrl_in[worldid, actid]
+  qM_in: wp.array3d(dtype=float),
+  # In:
+  qMi: wp.array(dtype=int),
+  qMj: wp.array(dtype=int),
+  qDeriv_in: wp.array3d(dtype=float),
+  # Out:
+  qDeriv_out: wp.array3d(dtype=float),
+):
+  worldid, elemid = wp.tid()
 
-      if vel != 0.0:
+  dofiid = qMi[elemid]
-        qderiv += actuator_moment_in[worldid, actid, dofiid] * actuator_moment_in[worldid, actid, dofjid] * vel
+  dofjid = qMj[elemid]
 
-  # TODO(team): fluid model derivative
+  if opt_is_sparse:
+    qderiv = qDeriv_in[worldid, 0, elemid]
+  else:
+    qderiv = qDeriv_in[worldid, dofiid, dofjid]
 
   if not opt_disableflags & DisableBit.DAMPER and dofiid == dofjid:
     qderiv -= dof_damping[worldid % dof_damping.shape[0], dofiid]
@@ -143,53 +211,82 @@ def _qderiv_tendon_damping(
 
 
 @event_scope
-def deriv_smooth_vel(m: Model, d: Data, qDeriv: wp.array2d(dtype=float)):
+def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d(dtype=float)):
   """Analytical derivative of smooth forces w.r.t. velocities.
 
   Args:
     m: The model containing kinematic and dynamic information (device).
     d: The data object containing the current state and output arrays (device).
-    qDeriv: Analytical derivative of smooth forces w.r.t. velocity.
+    out: qM - dt * qDeriv (derivatives of smooth forces w.r.t velocities).
   """
-  qMi = m.qM_fullm_i if m.opt.is_sparse else m.dof_tri_row
+  qMi = m.qM_fullm_i
-  qMj = m.qM_fullm_j if m.opt.is_sparse else m.dof_tri_col
+  qMj = m.qM_fullm_j
+
+  # TODO(team): implicit requires different sparsity structure
 
   if ~(m.opt.disableflags & (DisableBit.ACTUATION | DisableBit.DAMPER)):
+    # TODO(team): only clear elements not set by _qderiv_actuator_passive
+    out.zero_()
+    if m.nu > 0 and not m.opt.disableflags & DisableBit.ACTUATION:
+      vel = wp.empty((d.nworld, m.nu), dtype=float)
+      wp.launch(
+        _qderiv_actuator_passive_vel,
+        dim=(d.nworld, m.nu),
+        inputs=[
+          m.actuator_dyntype,
+          m.actuator_gaintype,
+          m.actuator_biastype,
+          m.actuator_actadr,
+          m.actuator_actnum,
+          m.actuator_gainprm,
+          m.actuator_biasprm,
+          d.act,
+          d.ctrl,
+        ],
+        outputs=[vel],
+      )
+      if m.opt.is_sparse:
+        wp.launch(
+          _qderiv_actuator_passive_actuation_sparse,
+          dim=(d.nworld, qMi.size),
+          inputs=[m.nu, d.actuator_moment, vel, qMi, qMj],
+          outputs=[out],
+        )
+      else:
+        vel_3d = vel.reshape(vel.shape + (1,))
+        for tile in m.qM_tiles:
+          wp.launch_tiled(
+            _qderiv_actuator_passive_actuation_dense(tile, m.nu),
+            dim=(d.nworld, tile.adr.size),
+            inputs=[vel_3d, d.actuator_moment, tile.adr],
+            outputs=[out],
+            block_dim=m.block_dim.mul_m_dense,
+          )
     wp.launch(
       _qderiv_actuator_passive,
       dim=(d.nworld, qMi.size),
       inputs=[
-        m.nu,
         m.opt.timestep,
         m.opt.disableflags,
         m.opt.is_sparse,
         m.dof_damping,
-        m.actuator_dyntype,
-        m.actuator_gaintype,
-        m.actuator_biastype,
-        m.actuator_actadr,
-        m.actuator_actnum,
-        m.actuator_gainprm,
-        m.actuator_biasprm,
-        d.act,
-        d.ctrl,
-        d.actuator_moment,
         d.qM,
         qMi,
         qMj,
+        out,
       ],
-      outputs=[qDeriv],
+      outputs=[out],
     )
   else:
     # TODO(team): directly utilize qM for these settings
-    wp.copy(qDeriv, d.qM)
+    wp.copy(out, d.qM)
 
   if not m.opt.disableflags & DisableBit.DAMPER:
     wp.launch(
       _qderiv_tendon_damping,
       dim=(d.nworld, qMi.size),
       inputs=[m.ntendon, m.opt.timestep, m.opt.is_sparse, m.tendon_damping, d.ten_J, qMi, qMj],
-      outputs=[qDeriv],
+      outputs=[out],
     )
 
   # TODO(team): rne derivative

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

@@ -509,6 +509,7 @@ def fwd_position(m: Model, d: Data, factorize: bool = True):
   smooth.kinematics(m, d)
   smooth.com_pos(m, d)
   smooth.camlight(m, d)
+  smooth.flex(m, d)
   smooth.tendon(m, d)
   smooth.crb(m, d)
   smooth.tendon_armature(m, d)

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

@@ -121,7 +121,9 @@ def inv_constraint(m: Model, d: Data):
     return
 
   # update
-  solver.create_context(m, d, grad=False)
+  h = wp.empty((d.nworld, 0, 0), dtype=float)  # not used
+  hfactor = wp.empty((d.nworld, 0, 0), dtype=float)  # not used
+  solver.create_context(m, d, h, hfactor, grad=False)
 
 
 def inverse(m: Model, d: Data):

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

@@ -49,6 +49,16 @@ def _create_array(data: Any, spec: wp.array, sizes: dict[str, int]) -> Union[wp.
   return array
 
 
+def is_sparse(mjm: mujoco.MjModel) -> bool:
+  if mjm.opt.jacobian == mujoco.mjtJacobian.mjJAC_AUTO:
+    if mjm.nv > 32:
+      return True
+    else:
+      return False
+  else:
+    return bool(mujoco.mj_isSparse(mjm))
+
+
 def put_model(mjm: mujoco.MjModel) -> types.Model:
   """Creates a model on device.
 
@@ -93,9 +103,6 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     if field & ~np.bitwise_or.reduce(field_type):
       raise NotImplementedError(f"{field_type.__name__} {field} is unsupported.")
 
-  if mjm.nflex > 1:
-    raise NotImplementedError("Only one flex is unsupported.")
-
   if ((mjm.flex_contype != 0) | (mjm.flex_conaffinity != 0)).any():
     raise NotImplementedError("Flex collisions are not implemented.")
 
@@ -137,23 +144,20 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
         return True
     return False
 
-  for geoms in [
+  for objtype, objid, reftype, refid in zip(
-    (types.GeomType.BOX, types.GeomType.BOX),
+    mjm.sensor_objtype[is_collision_sensor],
-    (types.GeomType.CAPSULE, types.GeomType.BOX),
+    mjm.sensor_objid[is_collision_sensor],
-    (types.GeomType.CYLINDER, types.GeomType.BOX),
+    mjm.sensor_reftype[is_collision_sensor],
-    (types.GeomType.PLANE, types.GeomType.BOX),
+    mjm.sensor_refid[is_collision_sensor],
-  ]:
+  ):
-    for objtype, objid, reftype, refid in zip(
+    if not_implemented(objtype, objid, types.GeomType.BOX) and not_implemented(reftype, refid, types.GeomType.BOX):
-      mjm.sensor_objtype[is_collision_sensor],
+      raise NotImplementedError(f"Collision sensors with box-box collisions are not implemented.")
-      mjm.sensor_objid[is_collision_sensor],
-      mjm.sensor_reftype[is_collision_sensor],
-      mjm.sensor_refid[is_collision_sensor],
-    ):
-      if not_implemented(objtype, objid, geoms[0]) and not_implemented(reftype, refid, geoms[1]):
-        raise NotImplementedError(f"Collision sensors with {geoms[0]} and {geoms[1]} are not implemented.")
 
   # create opt
-  opt = types.Option(**{f.name: getattr(mjm.opt, f.name, None) for f in dataclasses.fields(types.Option)})
+  opt_kwargs = {f.name: getattr(mjm.opt, f.name, None) for f in dataclasses.fields(types.Option)}
+  if hasattr(mjm.opt, "impratio"):
+    opt_kwargs["impratio_invsqrt"] = 1.0 / np.sqrt(np.maximum(mjm.opt.impratio, mujoco.mjMINVAL))
+  opt = types.Option(**opt_kwargs)
 
   # C MuJoCo tolerance was chosen for float64 architecture, but we default to float32 on GPU
   # adjust the tolerance for lower precision, to avoid the solver spending iterations needlessly
@@ -161,7 +165,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   opt.tolerance = max(opt.tolerance, 1e-6)
 
   # warp only fields
-  opt.is_sparse = bool(mujoco.mj_isSparse(mjm))
+  opt.is_sparse = is_sparse(mjm)
   ls_parallel_id = mujoco.mj_name2id(mjm, mujoco.mjtObj.mjOBJ_NUMERIC, "ls_parallel")
   opt.ls_parallel = (ls_parallel_id > -1) and (mjm.numeric_data[mjm.numeric_adr[ls_parallel_id]] == 1)
   opt.ls_parallel_min_step = 1.0e-6  # TODO(team): determine good default setting
@@ -170,7 +174,11 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   opt.broadphase_filter = types.BroadphaseFilter.PLANE | types.BroadphaseFilter.SPHERE | types.BroadphaseFilter.OBB
   opt.graph_conditional = True
   opt.run_collision_detection = True
-  opt.contact_sensor_maxmatch = 64
+  contact_sensor_maxmatch_id = mujoco.mj_name2id(mjm, mujoco.mjtObj.mjOBJ_NUMERIC, "contact_sensor_maxmatch")
+  if contact_sensor_maxmatch_id > -1:
+    opt.contact_sensor_maxmatch = mjm.numeric_data[mjm.numeric_adr[contact_sensor_maxmatch_id]]
+  else:
+    opt.contact_sensor_maxmatch = 64
 
   # place opt on device
   for f in dataclasses.fields(types.Option):
@@ -188,6 +196,9 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   m.opt = opt
   m.stat = stat
 
+  m.nv_pad = _get_padded_sizes(
+    mjm.nv, 0, is_sparse(mjm), types.TILE_SIZE_JTDAJ_SPARSE if is_sparse(mjm) else types.TILE_SIZE_JTDAJ_DENSE
+  )[1]
   m.nacttrnbody = (mjm.actuator_trntype == mujoco.mjtTrn.mjTRN_BODY).sum()
   m.nsensortaxel = mjm.mesh_vertnum[mjm.sensor_objid[mjm.sensor_type == mujoco.mjtSensor.mjSENS_TACTILE]].sum()
   m.nsensorcontact = (mjm.sensor_type == mujoco.mjtSensor.mjSENS_CONTACT).sum()
@@ -249,10 +260,12 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     nxn_pairid_contact[upper_tri_index(mjm.ngeom, mjm.pair_geom1[i], mjm.pair_geom2[i])] = i
 
   sensor_collision_adr = np.nonzero(is_collision_sensor)[0]
+  collision_sensor_adr = np.full(mjm.nsensor, -1)
+  collision_sensor_adr[sensor_collision_adr] = np.arange(len(sensor_collision_adr))
+
   nxn_pairid_collision = -1 * np.ones(len(geom1), dtype=int)
   pairids = []
-  collision_geom_adr = [0]
+  sensor_collision_start_adr = []
-  m.sensor_collision_start_adr = []
   for i in range(sensor_collision_adr.size):
     sensorid = sensor_collision_adr[i]
     objtype = mjm.sensor_objtype[sensorid]
@@ -275,24 +288,20 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
       id2 = refid
 
     # collide all pairs
-    geomid = 0
     for geom1id in range(id1, id1 + n1):
       for geom2id in range(id2, id2 + n2):
         pairid = upper_tri_index(mjm.ngeom, geom1id, geom2id)
 
         if pairid in pairids:
-          m.sensor_collision_start_adr.append(nxn_pairid_collision[pairid])
+          sensor_collision_start_adr.append(nxn_pairid_collision[pairid])
         else:
+          npairids = len(pairids)
+          nxn_pairid_collision[pairid] = npairids
+          sensor_collision_start_adr.append(npairids)
           pairids.append(pairid)
-          adr = collision_geom_adr[-1] + geomid
-          nxn_pairid_collision[pairid] = adr
-          m.sensor_collision_start_adr.append(adr)
-
-        geomid += 1
-    if i < sensor_collision_adr.size - 1:
-      collision_geom_adr.append(collision_geom_adr[-1] + n1 * n2)
 
   m.nsensorcollision = (nxn_pairid_collision >= 0).sum()
+  m.sensor_collision_start_adr = np.array(sensor_collision_start_adr)
   nxn_include = (nxn_pairid_contact > -2) | (nxn_pairid_collision >= 0)
 
   if nxn_include.sum() < 250_000:
@@ -363,6 +372,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   m.eq_wld_adr = np.nonzero(mjm.eq_type == types.EqType.WELD)[0]
   m.eq_jnt_adr = np.nonzero(mjm.eq_type == types.EqType.JOINT)[0]
   m.eq_ten_adr = np.nonzero(mjm.eq_type == types.EqType.TENDON)[0]
+  m.eq_flex_adr = np.nonzero(mjm.eq_type == types.EqType.FLEX)[0]
 
   # fixed tendon
   m.tendon_jnt_adr, m.wrap_jnt_adr = [], []
@@ -545,7 +555,7 @@ def _get_padded_sizes(nv: int, njmax: int, is_sparse: bool, tile_size: int):
     return ((x + multiple - 1) // multiple) * multiple
 
   njmax_padded = round_up(njmax, tile_size)
-  nv_padded = round_up(nv, tile_size) if is_sparse else round_up(nv, 4)
+  nv_padded = round_up(nv, tile_size) if (is_sparse or nv > 32) else round_up(nv, 4)
 
   return njmax_padded, nv_padded
 
@@ -592,8 +602,8 @@ def make_data(
   sizes = dict({"*": 1}, **{f.name: getattr(mjm, f.name, None) for f in dataclasses.fields(types.Model) if f.type is int})
   sizes["nmaxcondim"] = np.concatenate(([0], mjm.geom_condim, mjm.pair_dim)).max()
   sizes["nmaxpyramid"] = np.maximum(1, 2 * (sizes["nmaxcondim"] - 1))
-  tile_size = types.TILE_SIZE_JTDAJ_SPARSE if mujoco.mj_isSparse(mjm) else types.TILE_SIZE_JTDAJ_DENSE
+  tile_size = types.TILE_SIZE_JTDAJ_SPARSE if is_sparse(mjm) else types.TILE_SIZE_JTDAJ_DENSE
-  sizes["njmax_pad"], sizes["nv_pad"] = _get_padded_sizes(mjm.nv, njmax, mujoco.mj_isSparse(mjm), tile_size)
+  sizes["njmax_pad"], sizes["nv_pad"] = _get_padded_sizes(mjm.nv, njmax, is_sparse(mjm), tile_size)
   sizes["nworld"] = nworld
   sizes["naconmax"] = naconmax
   sizes["njmax"] = njmax
@@ -601,7 +611,18 @@ def make_data(
   contact = types.Contact(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Contact)})
   efc = types.Constraint(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Constraint)})
 
+  # world body and static geom (attached to the world) poses are precomputed
+  # this speeds up scenes with many static geoms (e.g. terrains)
+  # TODO(team): remove this when we introduce dof islands + sleeping
+  mjd = mujoco.MjData(mjm)
+  mujoco.mj_kinematics(mjm, mjd)
+
+  # mocap
+  mocap_body = np.nonzero(mjm.body_mocapid >= 0)[0]
+  mocap_id = mjm.body_mocapid[mocap_body]
+
   d_kwargs = {
+    "qpos": wp.array(np.tile(mjm.qpos0, nworld), shape=(nworld, mjm.nq), dtype=float),
     "contact": contact,
     "efc": efc,
     "nworld": nworld,
@@ -609,8 +630,20 @@ def make_data(
     "njmax": njmax,
     "qM": None,
     "qLD": None,
-    "geom_xpos": None,
+    # world body
-    "geom_xmat": None,
+    "xquat": wp.array(np.tile(mjd.xquat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.quat),
+    "xmat": wp.array(np.tile(mjd.xmat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.mat33),
+    "ximat": wp.array(np.tile(mjd.ximat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.mat33),
+    # static geoms
+    "geom_xpos": wp.array(np.tile(mjd.geom_xpos, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.vec3),
+    "geom_xmat": wp.array(np.tile(mjd.geom_xmat, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.mat33),
+    # mocap
+    "mocap_pos": wp.array(np.tile(mjm.body_pos[mocap_body[mocap_id]], (nworld, 1)), shape=(nworld, mjm.nmocap), dtype=wp.vec3),
+    "mocap_quat": wp.array(
+      np.tile(mjm.body_quat[mocap_body[mocap_id]], (nworld, 1)), shape=(nworld, mjm.nmocap), dtype=wp.quat
+    ),
+    # equality constraints
+    "eq_active": wp.array(np.tile(mjm.eq_active0.astype(bool), (nworld, 1)), shape=(nworld, mjm.neq), dtype=bool),
   }
   for f in dataclasses.fields(types.Data):
     if f.name in d_kwargs:
@@ -619,21 +652,13 @@ def make_data(
 
   d = types.Data(**d_kwargs)
 
-  if mujoco.mj_isSparse(mjm):
+  if is_sparse(mjm):
     d.qM = wp.zeros((nworld, 1, mjm.nM), dtype=float)
     d.qLD = wp.zeros((nworld, 1, mjm.nC), dtype=float)
   else:
     d.qM = wp.zeros((nworld, sizes["nv_pad"], sizes["nv_pad"]), dtype=float)
     d.qLD = wp.zeros((nworld, mjm.nv, mjm.nv), dtype=float)
 
-  # static geoms (attached to the world) have their poses calculated once during make_data instead
-  # of during each physics step.  this speeds up scenes with many static geoms (e.g. terrains)
-  # TODO(team): remove this when we introduce dof islands + sleeping
-  mjd = mujoco.MjData(mjm)
-  mujoco.mj_kinematics(mjm, mjd)
-  d.geom_xpos = wp.array(np.tile(mjd.geom_xpos, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.vec3)
-  d.geom_xmat = wp.array(np.tile(mjd.geom_xmat, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.mat33)
-
   return d
 
 
@@ -691,8 +716,8 @@ def put_data(
   sizes = dict({"*": 1}, **{f.name: getattr(mjm, f.name, None) for f in dataclasses.fields(types.Model) if f.type is int})
   sizes["nmaxcondim"] = np.concatenate(([0], mjm.geom_condim, mjm.pair_dim)).max()
   sizes["nmaxpyramid"] = np.maximum(1, 2 * (sizes["nmaxcondim"] - 1))
-  tile_size = types.TILE_SIZE_JTDAJ_SPARSE if mujoco.mj_isSparse(mjm) else types.TILE_SIZE_JTDAJ_DENSE
+  tile_size = types.TILE_SIZE_JTDAJ_SPARSE if is_sparse(mjm) else types.TILE_SIZE_JTDAJ_DENSE
-  sizes["njmax_pad"], sizes["nv_pad"] = _get_padded_sizes(mjm.nv, njmax, mujoco.mj_isSparse(mjm), tile_size)
+  sizes["njmax_pad"], sizes["nv_pad"] = _get_padded_sizes(mjm.nv, njmax, is_sparse(mjm), tile_size)
   sizes["nworld"] = nworld
   sizes["naconmax"] = naconmax
   sizes["njmax"] = njmax
@@ -766,11 +791,13 @@ def put_data(
     "qLD": None,
     "ten_J": None,
     "actuator_moment": None,
+    "flexedge_J": None,
     "nacon": None,
     "ne_connect": None,
     "ne_weld": None,
     "ne_jnt": None,
     "ne_ten": None,
+    "ne_flex": None,
     "nsolving": None,
   }
   for f in dataclasses.fields(types.Data):
@@ -785,12 +812,9 @@ def put_data(
   d = types.Data(**d_kwargs)
   d.solver_niter = wp.full((nworld,), mjd.solver_niter[0], dtype=int)
 
-  if mujoco.mj_isSparse(mjm):
+  if is_sparse(mjm):
     d.qM = wp.array(np.full((nworld, 1, mjm.nM), mjd.qM), dtype=float)
     d.qLD = wp.array(np.full((nworld, 1, mjm.nC), mjd.qLD), dtype=float)
-    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)
   else:
     qM = np.zeros((mjm.nv, mjm.nv))
     mujoco.mj_fullM(mjm, qM, mjd.qM)
@@ -799,8 +823,21 @@ def put_data(
     qM_padded = np.pad(qM, ((0, padding), (0, padding)), mode="constant", constant_values=0.0)
     d.qM = wp.array(np.full((nworld, sizes["nv_pad"], sizes["nv_pad"]), qM_padded), dtype=float)
     d.qLD = wp.array(np.full((nworld, mjm.nv, mjm.nv), qLD), dtype=float)
+
+  if mujoco.mj_isSparse(mjm):
+    ten_J = np.zeros((mjm.ntendon, mjm.nv))
+    mujoco.mju_sparse2dense(ten_J, mjd.ten_J.reshape(-1), mjd.ten_J_rownnz, mjd.ten_J_rowadr, mjd.ten_J_colind.reshape(-1))
+    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), 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)
+  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)
 
   # TODO(taylorhowell): sparse actuator_moment
   actuator_moment = np.zeros((mjm.nu, mjm.nv))
@@ -812,6 +849,7 @@ def put_data(
   d.ne_weld = wp.full(nworld, 6 * np.sum((mjm.eq_type == mujoco.mjtEq.mjEQ_WELD) & mjd.eq_active), dtype=int)
   d.ne_jnt = wp.full(nworld, np.sum((mjm.eq_type == mujoco.mjtEq.mjEQ_JOINT) & mjd.eq_active), dtype=int)
   d.ne_ten = wp.full(nworld, np.sum((mjm.eq_type == mujoco.mjtEq.mjEQ_TENDON) & mjd.eq_active), dtype=int)
+  d.ne_flex = wp.full(nworld, np.sum((mjm.eq_type == mujoco.mjtEq.mjEQ_FLEX) & mjd.eq_active), dtype=int)
   d.nsolving = wp.array([nworld], dtype=int)
 
   return d
@@ -913,6 +951,7 @@ 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]
   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]
@@ -957,12 +996,9 @@ def get_data_into(
   result.contact.geom[:ncon] = d.contact.geom.numpy()[ncon_filter]
   result.contact.efc_address[:ncon] = contact_efc_address_ordered[:ncon]
 
-  if mujoco.mj_isSparse(mjm):
+  if is_sparse(mjm):
     result.qM[:] = d.qM.numpy()[world_id, 0]
     result.qLD[:] = d.qLD.numpy()[world_id, 0]
-    if nefc > 0:
-      efc_J = d.efc.J.numpy()[world_id, efc_idx, : mjm.nv]
-      mujoco.mju_dense2sparse(result.efc_J, efc_J, result.efc_J_rownnz, result.efc_J_rowadr, result.efc_J_colind)
   else:
     qM = d.qM.numpy()[world_id]
     adr = 0
@@ -973,7 +1009,12 @@ def get_data_into(
         j = mjm.dof_parentid[j]
         adr += 1
     mujoco.mj_factorM(mjm, result)
-    if nefc > 0:
+
+  if nefc > 0:
+    if mujoco.mj_isSparse(mjm):
+      efc_J = d.efc.J.numpy()[world_id, efc_idx, : mjm.nv]
+      mujoco.mju_dense2sparse(result.efc_J, efc_J, result.efc_J_rownnz, result.efc_J_rowadr, result.efc_J_colind)
+    else:
       result.efc_J[: nefc * mjm.nv] = d.efc.J.numpy()[world_id, :nefc, : mjm.nv].flatten()
 
   # efc
@@ -1072,6 +1113,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
     ne_weld_out: wp.array(dtype=int),
     ne_jnt_out: wp.array(dtype=int),
     ne_ten_out: wp.array(dtype=int),
+    ne_flex_out: wp.array(dtype=int),
     nsolving_out: wp.array(dtype=int),
   ):
     worldid = wp.tid()
@@ -1088,6 +1130,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
     ne_weld_out[worldid] = 0
     ne_jnt_out[worldid] = 0
     ne_ten_out[worldid] = 0
+    ne_flex_out[worldid] = 0
     nf_out[worldid] = 0
     nl_out[worldid] = 0
     nefc_out[worldid] = 0
@@ -1095,8 +1138,9 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
       nsolving_out[0] = nworld_in
     time_out[worldid] = 0.0
     energy_out[worldid] = wp.vec2(0.0, 0.0)
+    qpos0_id = worldid % qpos0.shape[0]
     for i in range(nq):
-      qpos_out[worldid, i] = qpos0[worldid, i]
+      qpos_out[worldid, i] = qpos0[qpos0_id, i]
       if i < nv:
         qvel_out[worldid, i] = 0.0
         qacc_warmstart_out[worldid, i] = 0.0
@@ -1254,6 +1298,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
       d.ne_weld,
       d.ne_jnt,
       d.ne_ten,
+      d.ne_flex,
       d.nsolving,
     ],
   )

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

@@ -551,11 +551,14 @@ def _qfrc_passive(
 @wp.kernel
 def _flex_elasticity(
   # Model:
+  nflex: int,
   opt_timestep: wp.array(dtype=float),
   body_dofadr: wp.array(dtype=int),
   flex_dim: wp.array(dtype=int),
   flex_vertadr: wp.array(dtype=int),
   flex_edgeadr: wp.array(dtype=int),
+  flex_elemadr: wp.array(dtype=int),
+  flex_elemnum: wp.array(dtype=int),
   flex_elemedgeadr: wp.array(dtype=int),
   flex_vertbodyid: wp.array(dtype=int),
   flex_elem: wp.array(dtype=int),
@@ -574,22 +577,27 @@ def _flex_elasticity(
 ):
   worldid, elemid = wp.tid()
   timestep = opt_timestep[worldid % opt_timestep.shape[0]]
-  f = 0  # TODO(quaglino): this should become a function of t
+
+  for i in range(nflex):
+    locid = elemid - flex_elemadr[i]
+    if locid >= 0 and locid < flex_elemnum[i]:
+      f = i
+      break
 
   dim = flex_dim[f]
   nvert = dim + 1
   nedge = nvert * (nvert - 1) / 2
   edges = wp.where(
     dim == 3,
-    wp.mat(0, 1, 1, 2, 2, 0, 2, 3, 0, 3, 1, 3, shape=(6, 2), dtype=int),
+    wp.types.matrix(0, 1, 1, 2, 2, 0, 2, 3, 0, 3, 1, 3, shape=(6, 2), dtype=int),
-    wp.mat(1, 2, 2, 0, 0, 1, 0, 0, 0, 0, 0, 0, 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),
   )
   if timestep > 0.0 and not dsbl_damper:
     kD = flex_damping[f] / timestep
   else:
     kD = 0.0
 
-  gradient = wp.mat(0.0, shape=(6, 6))
+  gradient = wp.types.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]]
@@ -601,14 +609,14 @@ def _flex_elasticity(
 
   elongation = wp.spatial_vectorf(0.0)
   for e in range(nedge):
-    idx = flex_elemedge[flex_elemedgeadr[f] + elemid * nedge + e]
+    idx = flex_elemedge[elemid * nedge + e]
     vel = flexedge_velocity_in[worldid, flex_edgeadr[f] + idx]
     deformed = flexedge_length_in[worldid, flex_edgeadr[f] + idx]
     reference = flexedge_length0[flex_edgeadr[f] + idx]
     previous = deformed - vel * timestep
     elongation[e] = deformed * deformed - reference * reference + (deformed * deformed - previous * previous) * kD
 
-  metric = wp.mat(0.0, shape=(6, 6))
+  metric = wp.types.matrix(0.0, shape=(6, 6))
   id = int(0)
   for ed1 in range(nedge):
     for ed2 in range(ed1, nedge):
@@ -616,7 +624,7 @@ def _flex_elasticity(
       metric[ed2, ed1] = flex_stiffness[elemid, id]
       id += 1
 
-  force = wp.mat(0.0, shape=(6, 3))
+  force = wp.types.matrix(0.0, shape=(6, 3))
   for ed1 in range(nedge):
     for ed2 in range(nedge):
       for i in range(2):
@@ -633,10 +641,12 @@ def _flex_elasticity(
 @wp.kernel
 def _flex_bending(
   # Model:
+  nflex: int,
   body_dofadr: wp.array(dtype=int),
   flex_dim: wp.array(dtype=int),
   flex_vertadr: wp.array(dtype=int),
   flex_edgeadr: wp.array(dtype=int),
+  flex_edgenum: wp.array(dtype=int),
   flex_vertbodyid: wp.array(dtype=int),
   flex_edge: wp.array(dtype=wp.vec2i),
   flex_edgeflap: wp.array(dtype=wp.vec2i),
@@ -648,22 +658,27 @@ def _flex_bending(
 ):
   worldid, edgeid = wp.tid()
   nvert = 4
-  f = 0  # TODO(quaglino): this should become a function of t
+
+  for i in range(nflex):
+    locid = edgeid - flex_edgeadr[i]
+    if locid >= 0 and locid < flex_edgenum[i]:
+      f = i
+      break
 
   if flex_dim[f] != 2:
     return
 
-  v = wp.vec4i(
+  if flex_edgeflap[edgeid][1] == -1:
-    flex_edge[edgeid + flex_edgeadr[f]][0],
-    flex_edge[edgeid + flex_edgeadr[f]][1],
-    flex_edgeflap[edgeid + flex_edgeadr[f]][0],
-    flex_edgeflap[edgeid + flex_edgeadr[f]][1],
-  )
-
-  if v[3] == -1:
     return
 
-  frc = wp.mat(0.0, shape=(4, 3))
+  v = wp.vec4i(
+    flex_vertadr[f] + flex_edge[edgeid][0],
+    flex_vertadr[f] + flex_edge[edgeid][1],
+    flex_vertadr[f] + flex_edgeflap[edgeid][0],
+    flex_vertadr[f] + flex_edgeflap[edgeid][1],
+  )
+
+  frc = wp.types.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]]
@@ -674,7 +689,7 @@ def _flex_bending(
     frc[3] = wp.cross(v1 - v0, v2 - v0)
     frc[0] = -(frc[1] + frc[2] + frc[3])
 
-  force = wp.mat(0.0, shape=(nvert, 3))
+  force = wp.types.matrix(0.0, shape=(nvert, 3))
   for i in range(nvert):
     for x in range(3):
       for j in range(nvert):
@@ -743,11 +758,14 @@ def passive(m: Model, d: Data):
       _flex_elasticity,
       dim=(d.nworld, m.nflexelem),
       inputs=[
+        m.nflex,
         m.opt.timestep,
         m.body_dofadr,
         m.flex_dim,
         m.flex_vertadr,
         m.flex_edgeadr,
+        m.flex_elemadr,
+        m.flex_elemnum,
         m.flex_elemedgeadr,
         m.flex_vertbodyid,
         m.flex_elem,
@@ -766,10 +784,12 @@ def passive(m: Model, d: Data):
     _flex_bending,
     dim=(d.nworld, m.nflexedge),
     inputs=[
+      m.nflex,
       m.body_dofadr,
       m.flex_dim,
       m.flex_vertadr,
       m.flex_edgeadr,
+      m.flex_edgenum,
       m.flex_vertbodyid,
       m.flex_edge,
       m.flex_edgeflap,

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

@@ -455,6 +455,7 @@ def _clock(time_in: wp.array(dtype=float), worldid: int) -> float:
 @wp.kernel
 def _sensor_pos(
   # Model:
+  ngeom: int,
   opt_magnetic: wp.array(dtype=wp.vec3),
   body_geomnum: wp.array(dtype=int),
   body_geomadr: wp.array(dtype=int),
@@ -481,10 +482,9 @@ def _sensor_pos(
   sensor_refid: wp.array(dtype=int),
   sensor_adr: wp.array(dtype=int),
   sensor_cutoff: wp.array(dtype=float),
+  nxn_pairid: wp.array(dtype=wp.vec2i),
   sensor_pos_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),
   energy_in: wp.array(dtype=wp.vec2),
@@ -605,13 +605,10 @@ def _sensor_pos(
     refid = sensor_refid[sensorid]
 
     # initialize
-    dist = float(1.0e32)
+    dist = float(sensor_cutoff[sensorid])
     pnts = vec6(0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
     flip = bool(False)
 
-    collision_sensorid = collision_sensor_adr[sensorid]
-    collision_start_adr = sensor_collision_start_adr[collision_sensorid]
-
     # check for flip direction
     if objtype == int(ObjType.BODY.value):
       n1 = body_geomnum[objid]
@@ -627,12 +624,20 @@ def _sensor_pos(
       id2 = refid
 
     for geom1 in range(n1):
+      geomid1 = id1 + geom1
       for geom2 in range(n2):
-        collisionid = collision_start_adr + geom1 * n2 + geom2
+        geomid2 = id2 + geom2
+
+        if geomid1 <= geomid2:
+          pairid = math.upper_tri_index(ngeom, geomid1, geomid2)
+        else:
+          pairid = math.upper_tri_index(ngeom, geomid2, geomid1)
+        collisionid = nxn_pairid[pairid][1]
+
         for i in range(8):
           dist_new = sensor_collision_in[worldid, collisionid, i, 0]
 
-          if dist_new <= dist:
+          if dist_new < dist:
             dist = dist_new
 
             if sensortype == SensorType.GEOMNORMAL or sensortype == SensorType.GEOMFROMTO:
@@ -645,14 +650,12 @@ def _sensor_pos(
                 sensor_collision_in[worldid, collisionid, i, 6],
               )
 
-              geomid1 = id1 + geom1
+            if geom_type[geomid1] > geom_type[geomid2]:
-              geomid2 = id2 + geom2
+              flip = True
-
+            elif geom_type[geomid1] == geom_type[geomid2]:
-              if geom_type[geomid2] < geom_type[geomid1]:
+              flip = geomid1 > geomid2
-                flip = True
+            else:
-              elif geom_type[geomid1] == geom_type[geomid2]:
+              flip = False
-                if geomid2 < geomid1:
-                  flip = True
     if sensortype == int(SensorType.GEOMDIST.value):
       _write_scalar(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, dist, out)
     elif sensortype == int(SensorType.GEOMNORMAL.value):
@@ -704,6 +707,7 @@ def _sensor_pos(
 def _sensor_collision(
   # Model:
   ngeom: int,
+  nxn_pairid: wp.array(dtype=wp.vec2i),
   # Data in:
   contact_dist_in: wp.array(dtype=float),
   contact_pos_in: wp.array(dtype=wp.vec3),
@@ -713,7 +717,6 @@ def _sensor_collision(
   contact_type_in: wp.array(dtype=int),
   contact_geomcollisionid_in: wp.array(dtype=int),
   nacon_in: wp.array(dtype=int),
-  collision_pairid_in: wp.array(dtype=wp.vec2i),
   # Out:
   sensor_collision_out: wp.array4d(dtype=float),
 ):
@@ -732,7 +735,7 @@ def _sensor_collision(
     pairid = math.upper_tri_index(ngeom, geom[1], geom[0])
 
   worldid = contact_worldid_in[conid]
-  collisionid = collision_pairid_in[pairid][1]
+  collisionid = nxn_pairid[pairid][1]
   geomcollisionid = contact_geomcollisionid_in[conid]
 
   dist = contact_dist_in[conid]
@@ -800,6 +803,7 @@ def sensor_pos(m: Model, d: Data):
       dim=d.naconmax,
       inputs=[
         m.ngeom,
+        m.nxn_pairid,
         d.contact.dist,
         d.contact.pos,
         d.contact.frame,
@@ -808,7 +812,6 @@ def sensor_pos(m: Model, d: Data):
         d.contact.type,
         d.contact.geomcollisionid,
         d.nacon,
-        d.collision_pairid,
       ],
       outputs=[sensor_collision],
     )
@@ -817,6 +820,7 @@ def sensor_pos(m: Model, d: Data):
     _sensor_pos,
     dim=(d.nworld, m.sensor_pos_adr.size),
     inputs=[
+      m.ngeom,
       m.opt.magnetic,
       m.body_geomnum,
       m.body_geomadr,
@@ -843,10 +847,9 @@ def sensor_pos(m: Model, d: Data):
       m.sensor_refid,
       m.sensor_adr,
       m.sensor_cutoff,
+      m.nxn_pairid,
       m.sensor_pos_adr,
       m.rangefinder_sensor_adr,
-      m.sensor_collision_start_adr,
-      m.collision_sensor_adr,
       d.time,
       d.energy,
       d.qpos,

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

@@ -40,28 +40,12 @@ from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
 wp.set_module_options({"enable_backward": False})
 
 
-@wp.kernel
-def _kinematics_root(
-  # Data out:
-  xpos_out: wp.array2d(dtype=wp.vec3),
-  xquat_out: wp.array2d(dtype=wp.quat),
-  xmat_out: wp.array2d(dtype=wp.mat33),
-  xipos_out: wp.array2d(dtype=wp.vec3),
-  ximat_out: wp.array2d(dtype=wp.mat33),
-):
-  worldid = wp.tid()
-  xpos_out[worldid, 0] = wp.vec3(0.0)
-  xquat_out[worldid, 0] = wp.quat(1.0, 0.0, 0.0, 0.0)
-  xipos_out[worldid, 0] = wp.vec3(0.0)
-  xmat_out[worldid, 0] = wp.identity(n=3, dtype=wp.float32)
-  ximat_out[worldid, 0] = wp.identity(n=3, dtype=wp.float32)
-
-
 @wp.kernel
 def _kinematics_level(
   # Model:
   qpos0: wp.array2d(dtype=float),
   body_parentid: wp.array(dtype=int),
+  body_mocapid: wp.array(dtype=int),
   body_jntnum: wp.array(dtype=int),
   body_jntadr: wp.array(dtype=int),
   body_pos: wp.array2d(dtype=wp.vec3),
@@ -74,6 +58,8 @@ def _kinematics_level(
   jnt_axis: wp.array2d(dtype=wp.vec3),
   # Data in:
   qpos_in: wp.array2d(dtype=float),
+  mocap_pos_in: wp.array2d(dtype=wp.vec3),
+  mocap_quat_in: wp.array2d(dtype=wp.quat),
   xpos_in: wp.array2d(dtype=wp.vec3),
   xquat_in: wp.array2d(dtype=wp.quat),
   xmat_in: wp.array2d(dtype=wp.mat33),
@@ -97,12 +83,12 @@ def _kinematics_level(
   body_quat_id = worldid % body_quat.shape[0]
   jnt_axis_id = worldid % jnt_axis.shape[0]
 
-  if jntnum == 0:
+  free_joint = False
-    # no joints - apply fixed translation and rotation relative to parent
+  if jntnum == 1:
-    pid = body_parentid[bodyid]
+    jnt_type_ = jnt_type[jntadr]
-    xpos = (xmat_in[worldid, pid] * body_pos[body_pos_id, bodyid]) + xpos_in[worldid, pid]
+    free_joint = jnt_type_ == JointType.FREE
-    xquat = math.mul_quat(xquat_in[worldid, pid], body_quat[body_quat_id, bodyid])
+
-  elif jntnum == 1 and jnt_type[jntadr] == JointType.FREE:
+  if free_joint:
     # free joint
     qadr = jnt_qposadr[jntadr]
     xpos = wp.vec3(qpos[qadr], qpos[qadr + 1], qpos[qadr + 2])
@@ -116,8 +102,19 @@ def _kinematics_level(
     qpos0_id = worldid % qpos0.shape[0]
     jnt_pos_id = worldid % jnt_pos.shape[0]
     pid = body_parentid[bodyid]
-    xpos = (xmat_in[worldid, pid] * body_pos[body_pos_id, bodyid]) + xpos_in[worldid, pid]
+
-    xquat = math.mul_quat(xquat_in[worldid, pid], body_quat[body_quat_id, bodyid])
+    # mocap bodies have world body as parent
+    mocapid = body_mocapid[bodyid]
+    if mocapid >= 0:
+      xpos = mocap_pos_in[worldid, mocapid]
+      xquat = mocap_quat_in[worldid, mocapid]
+    else:
+      xpos = body_pos[body_pos_id, bodyid]
+      xquat = body_quat[body_quat_id, bodyid]
+
+    if pid >= 0:
+      xpos = xmat_in[worldid, pid] @ xpos + xpos_in[worldid, pid]
+      xquat = math.mul_quat(xquat_in[worldid, pid], xquat)
 
     for _ in range(jntnum):
       qadr = jnt_qposadr[jntadr]
@@ -146,7 +143,8 @@ def _kinematics_level(
       jntadr += 1
 
   xpos_out[worldid, bodyid] = xpos
-  xquat_out[worldid, bodyid] = wp.normalize(xquat)
+  xquat = wp.normalize(xquat)
+  xquat_out[worldid, bodyid] = xquat
   xmat_out[worldid, bodyid] = math.quat_to_mat(xquat)
   xipos_out[worldid, bodyid] = xpos + math.rot_vec_quat(body_ipos[worldid % body_ipos.shape[0], bodyid], xquat)
   ximat_out[worldid, bodyid] = math.quat_to_mat(math.mul_quat(xquat, body_iquat[worldid % body_iquat.shape[0], bodyid]))
@@ -219,19 +217,33 @@ def _flex_vertices(
 @wp.kernel
 def _flex_edges(
   # Model:
+  nv: int,
+  nflex: int,
+  body_parentid: wp.array(dtype=int),
+  body_rootid: wp.array(dtype=int),
   body_dofadr: wp.array(dtype=int),
+  dof_bodyid: wp.array(dtype=int),
   flex_vertadr: wp.array(dtype=int),
+  flex_edgeadr: wp.array(dtype=int),
+  flex_edgenum: wp.array(dtype=int),
   flex_vertbodyid: wp.array(dtype=int),
   flex_edge: wp.array(dtype=wp.vec2i),
   # Data in:
   qvel_in: wp.array2d(dtype=float),
+  subtree_com_in: wp.array2d(dtype=wp.vec3),
+  cdof_in: wp.array2d(dtype=wp.spatial_vector),
   flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
   # Data out:
+  flexedge_J_out: wp.array3d(dtype=float),
   flexedge_length_out: wp.array2d(dtype=float),
   flexedge_velocity_out: wp.array2d(dtype=float),
 ):
   worldid, edgeid = wp.tid()
-  f = 0  # TODO(quaglino): get f from edgeid
+  for i in range(nflex):
+    locid = edgeid - flex_edgeadr[i]
+    if locid >= 0 and locid < flex_edgenum[i]:
+      f = i
+      break
   vbase = flex_vertadr[f]
   v = flex_edge[edgeid]
   pos1 = flexvert_xpos_in[worldid, vbase + v[0]]
@@ -240,38 +252,20 @@ def _flex_edges(
   vecnorm = wp.length(vec)
   flexedge_length_out[worldid, edgeid] = vecnorm
   # TODO(quaglino): use Jacobian
-  i = body_dofadr[flex_vertbodyid[vbase + v[0]]]
+  b1 = flex_vertbodyid[vbase + v[0]]
-  j = body_dofadr[flex_vertbodyid[vbase + v[1]]]
+  b2 = flex_vertbodyid[vbase + v[1]]
+  i = body_dofadr[b1]
+  j = body_dofadr[b2]
   vel1 = wp.vec3(qvel_in[worldid, i], qvel_in[worldid, i + 1], qvel_in[worldid, i + 2])
   vel2 = wp.vec3(qvel_in[worldid, j], qvel_in[worldid, j + 1], qvel_in[worldid, j + 2])
-  flexedge_velocity_out[worldid, edgeid] = math.safe_div(wp.dot(vel2 - vel1, vec), vecnorm)
+  edge = wp.normalize(vec)
-
+  flexedge_velocity_out[worldid, edgeid] = wp.dot(vel2 - vel1, edge)
-
+  # Edge jacobian
-@wp.kernel
+  for k in range(nv):
-def _mocap(
+    jacp1, _ = support.jac(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos1, b1, k, worldid)
-  # Model:
+    jacp2, _ = support.jac(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos2, b2, k, worldid)
-  body_ipos: wp.array2d(dtype=wp.vec3),
+    jacdif = jacp2 - jacp1
-  body_iquat: wp.array2d(dtype=wp.quat),
+    flexedge_J_out[worldid, edgeid, k] = wp.dot(jacdif, edge)
-  mocap_bodyid: wp.array(dtype=int),
-  # Data in:
-  mocap_pos_in: wp.array2d(dtype=wp.vec3),
-  mocap_quat_in: wp.array2d(dtype=wp.quat),
-  # Data out:
-  xpos_out: wp.array2d(dtype=wp.vec3),
-  xquat_out: wp.array2d(dtype=wp.quat),
-  xmat_out: wp.array2d(dtype=wp.mat33),
-  xipos_out: wp.array2d(dtype=wp.vec3),
-  ximat_out: wp.array2d(dtype=wp.mat33),
-):
-  worldid, mocapid = wp.tid()
-  bodyid = mocap_bodyid[mocapid]
-  mocap_quat = wp.normalize(mocap_quat_in[worldid, mocapid])
-  xpos = mocap_pos_in[worldid, mocapid]
-  xpos_out[worldid, bodyid] = xpos
-  xquat_out[worldid, bodyid] = mocap_quat
-  xmat_out[worldid, bodyid] = math.quat_to_mat(mocap_quat)
-  xipos_out[worldid, bodyid] = xpos + math.rot_vec_quat(body_ipos[worldid % body_ipos.shape[0], bodyid], mocap_quat)
-  ximat_out[worldid, bodyid] = math.quat_to_mat(math.mul_quat(mocap_quat, body_iquat[worldid % body_iquat.shape[0], bodyid]))
 
 
 @event_scope
@@ -282,8 +276,6 @@ def kinematics(m: Model, d: Data):
   derived positions and orientations of geoms, sites, and flexible elements, based on the
   current joint positions and any attached mocap bodies.
   """
-  wp.launch(_kinematics_root, dim=(d.nworld), inputs=[], outputs=[d.xpos, d.xquat, d.xmat, d.xipos, d.ximat])
-
   for i in range(1, len(m.body_tree)):
     body_tree = m.body_tree[i]
     wp.launch(
@@ -292,6 +284,7 @@ def kinematics(m: Model, d: Data):
       inputs=[
         m.qpos0,
         m.body_parentid,
+        m.body_mocapid,
         m.body_jntnum,
         m.body_jntadr,
         m.body_pos,
@@ -303,6 +296,8 @@ def kinematics(m: Model, d: Data):
         m.jnt_pos,
         m.jnt_axis,
         d.qpos,
+        d.mocap_pos,
+        d.mocap_quat,
         d.xpos,
         d.xquat,
         d.xmat,
@@ -311,13 +306,6 @@ def kinematics(m: Model, d: Data):
       outputs=[d.xpos, d.xquat, d.xmat, d.xipos, d.ximat, d.xanchor, d.xaxis],
     )
 
-  wp.launch(
-    _mocap,
-    dim=(d.nworld, m.nmocap),
-    inputs=[m.body_ipos, m.body_iquat, m.mocap_bodyid, d.mocap_pos, d.mocap_quat],
-    outputs=[d.xpos, d.xquat, d.xmat, d.xipos, d.ximat],
-  )
-
   wp.launch(
     _geom_local_to_global,
     dim=(d.nworld, m.ngeom),
@@ -332,12 +320,31 @@ def kinematics(m: Model, d: Data):
     outputs=[d.site_xpos, d.site_xmat],
   )
 
+
+@event_scope
+def flex(m: Model, d: Data):
   wp.launch(_flex_vertices, dim=(d.nworld, m.nflexvert), inputs=[m.flex_vertbodyid, d.xpos], outputs=[d.flexvert_xpos])
   wp.launch(
     _flex_edges,
     dim=(d.nworld, m.nflexedge),
-    inputs=[m.body_dofadr, m.flex_vertadr, m.flex_vertbodyid, m.flex_edge, d.qvel, d.flexvert_xpos],
+    inputs=[
-    outputs=[d.flexedge_length, d.flexedge_velocity],
+      m.nv,
+      m.nflex,
+      m.body_parentid,
+      m.body_rootid,
+      m.body_dofadr,
+      m.dof_bodyid,
+      m.flex_vertadr,
+      m.flex_edgeadr,
+      m.flex_edgenum,
+      m.flex_vertbodyid,
+      m.flex_edge,
+      d.qvel,
+      d.subtree_com,
+      d.cdof,
+      d.flexvert_xpos,
+    ],
+    outputs=[d.flexedge_J, d.flexedge_length, d.flexedge_velocity],
   )
 
 

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

@@ -31,10 +31,12 @@ from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
 
 wp.set_module_options({"enable_backward": False})
 
+_BLOCK_CHOLESKY_DIM = 32
+
 
 @wp.func
 def _rescale(nv: int, stat_meaninertia: float, value: float) -> float:
-  return value / (stat_meaninertia * float(wp.max(1, nv)))
+  return value / (stat_meaninertia * float(nv))
 
 
 @wp.func
@@ -291,10 +293,10 @@ def _eval(
 def linesearch_iterative(
   # Model:
   nv: int,
-  opt_impratio: wp.array(dtype=float),
   opt_tolerance: wp.array(dtype=float),
   opt_ls_tolerance: wp.array(dtype=float),
   opt_ls_iterations: int,
+  opt_impratio_invsqrt: wp.array(dtype=float),
   stat_meaninertia: float,
   # Data in:
   ne_in: wp.array(dtype=int),
@@ -321,7 +323,7 @@ def linesearch_iterative(
   if efc_done_in[worldid]:
     return
 
-  impratio = opt_impratio[worldid % opt_impratio.shape[0]]
+  impratio_invsqrt = opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
   efc_type = efc_type_in[worldid]
   efc_id = efc_id_in[worldid]
   efc_D = efc_D_in[worldid]
@@ -335,11 +337,10 @@ def linesearch_iterative(
   ne_clip = min(njmax_in, ne_in[worldid])
   nef_clip = min(njmax_in, ne_clip + nf_in[worldid])
   nefc_clip = min(njmax_in, nefc_in[worldid])
-  impratio_invsqrt = 1.0 / wp.sqrt(impratio)
 
   # Calculate p0
-  snorm = wp.math.sqrt(efc_search_dot_in[worldid])
+  snorm = wp.sqrt(efc_search_dot_in[worldid])
-  scale = stat_meaninertia * wp.float(wp.max(1, nv))
+  scale = stat_meaninertia * wp.float(nv)
   gtol = tolerance * ls_tolerance * snorm * scale
   p0 = wp.vec3(efc_quad_gauss[0], efc_quad_gauss[1], 2.0 * efc_quad_gauss[2])
   p0 += _eval_init(
@@ -461,10 +462,10 @@ def _linesearch_iterative(m: types.Model, d: types.Data):
     dim=d.nworld,
     inputs=[
       m.nv,
-      m.opt.impratio,
       m.opt.tolerance,
       m.opt.ls_tolerance,
       m.opt.ls_iterations,
+      m.opt.impratio_invsqrt,
       m.stat.meaninertia,
       d.ne,
       d.nf,
@@ -484,6 +485,7 @@ def _linesearch_iterative(m: types.Model, d: types.Data):
       d.njmax,
     ],
     outputs=[d.efc.alpha],
+    block_dim=m.block_dim.linesearch_iterative,
   )
 
 
@@ -496,8 +498,8 @@ def _log_scale(min_value: float, max_value: float, num_values: int, i: int) -> f
 @wp.kernel
 def linesearch_parallel_fused(
   # Model:
-  opt_impratio: wp.array(dtype=float),
   opt_ls_iterations: int,
+  opt_impratio_invsqrt: wp.array(dtype=float),
   opt_ls_parallel_min_step: float,
   # Data in:
   ne_in: wp.array(dtype=int),
@@ -569,7 +571,7 @@ def linesearch_parallel_fused(
         continue
 
       friction = contact_friction_in[conid]
-      mu = friction[0] / wp.sqrt(opt_impratio[worldid % opt_impratio.shape[0]])
+      mu = friction[0] * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
 
       # unpack quad
       efcid1 = contact_efc_address_in[conid, 1]
@@ -649,8 +651,8 @@ def _linesearch_parallel(m: types.Model, d: types.Data, cost: wp.array2d(dtype=f
     linesearch_parallel_fused,
     dim=(d.nworld, m.opt.ls_iterations),
     inputs=[
-      m.opt.impratio,
       m.opt.ls_iterations,
+      m.opt.impratio_invsqrt,
       m.opt.ls_parallel_min_step,
       d.ne,
       d.nf,
@@ -768,7 +770,7 @@ def linesearch_prepare_gauss(
 @wp.kernel
 def linesearch_prepare_quad(
   # Model:
-  opt_impratio: wp.array(dtype=float),
+  opt_impratio_invsqrt: wp.array(dtype=float),
   # Data in:
   nefc_in: wp.array(dtype=int),
   contact_friction_in: wp.array(dtype=types.vec5),
@@ -814,7 +816,7 @@ def linesearch_prepare_quad(
 
     dim = contact_dim_in[conid]
     friction = contact_friction_in[conid]
-    mu = friction[0] / wp.sqrt(opt_impratio[worldid])
+    mu = friction[0] * opt_impratio_invsqrt[worldid]
 
     u0 = Jaref * mu
     v0 = jv * mu
@@ -949,7 +951,7 @@ def _linesearch(m: types.Model, d: types.Data, cost: wp.array2d(dtype=float)):
     linesearch_prepare_quad,
     dim=(d.nworld, d.njmax),
     inputs=[
-      m.opt.impratio,
+      m.opt.impratio_invsqrt,
       d.nefc,
       d.contact.friction,
       d.contact.dim,
@@ -1061,7 +1063,7 @@ def update_constraint_init_cost(
 @wp.kernel
 def update_constraint_efc(
   # Model:
-  opt_impratio: wp.array(dtype=float),
+  opt_impratio_invsqrt: wp.array(dtype=float),
   # Data in:
   ne_in: wp.array(dtype=int),
   nf_in: wp.array(dtype=int),
@@ -1133,7 +1135,7 @@ def update_constraint_efc(
 
     dim = contact_dim_in[conid]
     friction = contact_friction_in[conid]
-    mu = friction[0] / wp.sqrt(opt_impratio[worldid])
+    mu = friction[0] * opt_impratio_invsqrt[worldid]
 
     efcid0 = contact_efc_address_in[conid, 0]
     if efcid0 < 0:
@@ -1261,7 +1263,7 @@ def _update_constraint(m: types.Model, d: types.Data):
     update_constraint_efc,
     dim=(d.nworld, d.njmax),
     inputs=[
-      m.opt.impratio,
+      m.opt.impratio_invsqrt,
       d.ne,
       d.nf,
       d.nefc,
@@ -1349,8 +1351,8 @@ def update_gradient_set_h_qM_lower_sparse(
   # Data in:
   qM_in: wp.array3d(dtype=float),
   efc_done_in: wp.array(dtype=bool),
-  # Data out:
+  # Out:
-  efc_h_out: wp.array3d(dtype=float),
+  h_out: wp.array3d(dtype=float),
 ):
   worldid, elementid = wp.tid()
 
@@ -1359,7 +1361,7 @@ def update_gradient_set_h_qM_lower_sparse(
 
   i = qM_fullm_i[elementid]
   j = qM_fullm_j[elementid]
-  efc_h_out[worldid, i, j] += qM_in[worldid, 0, elementid]
+  h_out[worldid, i, j] += qM_in[worldid, 0, elementid]
 
 
 @wp.func
@@ -1390,8 +1392,8 @@ def update_gradient_JTDAJ_sparse_tiled(tile_size: int, njmax: int):
     efc_D_in: wp.array2d(dtype=float),
     efc_state_in: wp.array2d(dtype=int),
     efc_done_in: wp.array(dtype=bool),
-    # Data out:
+    # Out:
-    efc_h_out: wp.array3d(dtype=float),
+    h_out: wp.array3d(dtype=float),
   ):
     worldid, elementid = wp.tid()
 
@@ -1442,7 +1444,7 @@ def update_gradient_JTDAJ_sparse_tiled(tile_size: int, njmax: int):
 
     # AD: setting bounds_check to True explicitly here because for some reason it was
     # slower to disable it.
-    wp.tile_store(efc_h_out[worldid], sum_val, offset=(offset_i, offset_j), bounds_check=True)
+    wp.tile_store(h_out[worldid], sum_val, offset=(offset_i, offset_j), bounds_check=True)
 
   return kernel
 
@@ -1463,8 +1465,8 @@ def update_gradient_JTDAJ_dense_tiled(nv_padded: int, tile_size: int, njmax: int
     efc_D_in: wp.array2d(dtype=float),
     efc_state_in: wp.array2d(dtype=int),
     efc_done_in: wp.array(dtype=bool),
-    # Data out:
+    # Out:
-    efc_h_out: wp.array3d(dtype=float),
+    h_out: wp.array3d(dtype=float),
   ):
     worldid = wp.tid()
 
@@ -1503,7 +1505,7 @@ def update_gradient_JTDAJ_dense_tiled(nv_padded: int, tile_size: int, njmax: int
 
       sum_val += wp.tile_matmul(J_ki, J_kj)
 
-    wp.tile_store(efc_h_out[worldid], sum_val, bounds_check=False)
+    wp.tile_store(h_out[worldid], sum_val, bounds_check=False)
 
   return kernel
 
@@ -1512,7 +1514,7 @@ def update_gradient_JTDAJ_dense_tiled(nv_padded: int, tile_size: int, njmax: int
 @wp.kernel
 def update_gradient_JTCJ(
   # Model:
-  opt_impratio: wp.array(dtype=float),
+  opt_impratio_invsqrt: wp.array(dtype=float),
   dof_tri_row: wp.array(dtype=int),
   dof_tri_col: wp.array(dtype=int),
   # Data in:
@@ -1532,8 +1534,8 @@ def update_gradient_JTCJ(
   # In:
   nblocks_perblock: int,
   dim_block: int,
-  # Data out:
+  # Out:
-  efc_h_out: wp.array3d(dtype=float),
+  h_out: wp.array3d(dtype=float),
 ):
   conid_start, elementid = wp.tid()
 
@@ -1564,7 +1566,7 @@ def update_gradient_JTCJ(
       continue
 
     fri = contact_friction_in[conid]
-    mu = math.safe_div(fri[0], wp.sqrt(opt_impratio[worldid]))
+    mu = fri[0] * opt_impratio_invsqrt[worldid]
 
     mu2 = mu * mu
     dm = math.safe_div(efc_D_in[worldid, efcid0], mu2 * (1.0 + mu2))
@@ -1589,7 +1591,7 @@ def update_gradient_JTCJ(
     t = wp.max(t, types.MJ_MINVAL)
     ttt = wp.max(t * t * t, types.MJ_MINVAL)
 
-    efc_h = float(0.0)
+    h = float(0.0)
 
     for dim1id in range(condim):
       if dim1id == 0:
@@ -1641,12 +1643,12 @@ def update_gradient_JTCJ(
         hcone *= dm * fri1 * fri2
 
         if hcone != 0.0:
-          efc_h += hcone * efc_J11 * efc_J22
+          h += hcone * efc_J11 * efc_J22
 
           if dim1id != dim2id:
-            efc_h += hcone * efc_J12 * efc_J21
+            h += hcone * efc_J12 * efc_J21
 
-    efc_h_out[worldid, dof1id, dof2id] += efc_h
+    h_out[worldid, dof1id, dof2id] += h
 
 
 @cache_kernel
@@ -1655,7 +1657,7 @@ def update_gradient_cholesky(tile_size: int):
   def kernel(
     # Data in:
     efc_grad_in: wp.array2d(dtype=float),
-    efc_h_in: wp.array3d(dtype=float),
+    h_in: wp.array3d(dtype=float),
     efc_done_in: wp.array(dtype=bool),
     # Data out:
     efc_Mgrad_out: wp.array2d(dtype=float),
@@ -1666,7 +1668,7 @@ def update_gradient_cholesky(tile_size: int):
     if efc_done_in[worldid]:
       return
 
-    mat_tile = wp.tile_load(efc_h_in[worldid], shape=(TILE_SIZE, TILE_SIZE))
+    mat_tile = wp.tile_load(h_in[worldid], shape=(TILE_SIZE, TILE_SIZE))
     fact_tile = wp.tile_cholesky(mat_tile)
     input_tile = wp.tile_load(efc_grad_in[worldid], shape=TILE_SIZE)
     output_tile = wp.tile_cholesky_solve(fact_tile, input_tile)
@@ -1676,34 +1678,48 @@ def update_gradient_cholesky(tile_size: int):
 
 
 @cache_kernel
-def update_gradient_cholesky_blocked(tile_size: int):
+def update_gradient_cholesky_blocked(tile_size: int, matrix_size: int):
   @nested_kernel(module="unique", enable_backward=False)
   def kernel(
     # Data in:
     efc_grad_in: wp.array3d(dtype=float),
-    efc_h_in: wp.array3d(dtype=float),
+    h_in: wp.array3d(dtype=float),
     efc_done_in: wp.array(dtype=bool),
-    matrix_size: int,
+    hfactor: wp.array3d(dtype=float),
-    cholesky_L_tmp: wp.array3d(dtype=float),
-    cholesky_y_tmp: wp.array3d(dtype=float),
     # Data out:
     efc_Mgrad_out: wp.array3d(dtype=float),
   ):
-    worldid, tid_block = wp.tid()
+    worldid = wp.tid()
     TILE_SIZE = wp.static(tile_size)
 
     if efc_done_in[worldid]:
       return
 
-    wp.static(create_blocked_cholesky_func(TILE_SIZE))(tid_block, efc_h_in[worldid], matrix_size, cholesky_L_tmp[worldid])
+    # We need matrix size both as a runtime input as well as a static input:
-    wp.static(create_blocked_cholesky_solve_func(TILE_SIZE))(
+    # static input is needed to specify the tile sizes for the compiler
-      tid_block, cholesky_L_tmp[worldid], efc_grad_in[worldid], cholesky_y_tmp[worldid], matrix_size, efc_Mgrad_out[worldid]
+    # runtime input is needed for the loop bounds, otherwise warp will unroll
+    # unconditionally leading to shared memory capacity issues.
+
+    wp.static(create_blocked_cholesky_func(TILE_SIZE))(h_in[worldid], matrix_size, hfactor[worldid])
+    wp.static(create_blocked_cholesky_solve_func(TILE_SIZE, matrix_size))(
+      hfactor[worldid], efc_grad_in[worldid], matrix_size, efc_Mgrad_out[worldid]
     )
 
   return kernel
 
 
-def _update_gradient(m: types.Model, d: types.Data):
+@wp.kernel
+def padding_h(nv: int, efc_done_in: wp.array(dtype=bool), h_out: wp.array3d(dtype=float)):
+  worldid, elementid = wp.tid()
+
+  if efc_done_in[worldid]:
+    return
+
+  dofid = nv + elementid
+  h_out[worldid, dofid, dofid] = 1.0
+
+
+def _update_gradient(m: types.Model, d: types.Data, h: wp.array3d(dtype=float), hfactor: wp.array3d(dtype=float)):
   # grad = Ma - qfrc_smooth - qfrc_constraint
   wp.launch(update_gradient_zero_grad_dot, dim=(d.nworld), inputs=[d.efc.done], outputs=[d.efc.grad_dot])
 
@@ -1731,7 +1747,7 @@ def _update_gradient(m: types.Model, d: types.Data):
           d.efc.state,
           d.efc.done,
         ],
-        outputs=[d.efc.h],
+        outputs=[h],
         block_dim=m.block_dim.update_gradient_JTDAJ_sparse,
       )
 
@@ -1739,7 +1755,7 @@ def _update_gradient(m: types.Model, d: types.Data):
         update_gradient_set_h_qM_lower_sparse,
         dim=(d.nworld, m.qM_fullm_i.size),
         inputs=[m.qM_fullm_i, m.qM_fullm_j, d.qM, d.efc.done],
-        outputs=[d.efc.h],
+        outputs=[h],
       )
     else:
       nv_padded = d.efc.J.shape[2]
@@ -1754,7 +1770,7 @@ def _update_gradient(m: types.Model, d: types.Data):
           d.efc.state,
           d.efc.done,
         ],
-        outputs=[d.efc.h],
+        outputs=[h],
         block_dim=m.block_dim.update_gradient_JTDAJ_dense,
       )
 
@@ -1786,7 +1802,7 @@ def _update_gradient(m: types.Model, d: types.Data):
         update_gradient_JTCJ,
         dim=(dim_block, m.dof_tri_row.size),
         inputs=[
-          m.opt.impratio,
+          m.opt.impratio_invsqrt,
           m.dof_tri_row,
           m.dof_tri_col,
           d.contact.dist,
@@ -1805,32 +1821,32 @@ def _update_gradient(m: types.Model, d: types.Data):
           nblocks_perblock,
           dim_block,
         ],
-        outputs=[d.efc.h],
+        outputs=[h],
       )
 
     # TODO(team): Define good threshold for blocked vs non-blocked cholesky
-    if m.nv < 32:
+    if m.nv <= _BLOCK_CHOLESKY_DIM:
       wp.launch_tiled(
         update_gradient_cholesky(m.nv),
         dim=d.nworld,
-        inputs=[d.efc.grad, d.efc.h, d.efc.done],
+        inputs=[d.efc.grad, h, d.efc.done],
         outputs=[d.efc.Mgrad],
         block_dim=m.block_dim.update_gradient_cholesky,
       )
     else:
+      wp.launch(
+        padding_h,
+        dim=(d.nworld, m.nv_pad - m.nv),
+        inputs=[m.nv, d.efc.done],
+        outputs=[h],
+      )
+
       wp.launch_tiled(
-        update_gradient_cholesky_blocked(16),
+        update_gradient_cholesky_blocked(types.TILE_SIZE_JTDAJ_DENSE, m.nv_pad),
         dim=d.nworld,
-        inputs=[
+        inputs=[d.efc.grad.reshape(shape=(d.nworld, d.efc.grad.shape[1], 1)), h, d.efc.done, hfactor],
-          d.efc.grad.reshape(shape=(d.nworld, m.nv, 1)),
+        outputs=[d.efc.Mgrad.reshape(shape=(d.nworld, d.efc.Mgrad.shape[1], 1))],
-          d.efc.h,
+        block_dim=m.block_dim.update_gradient_cholesky_blocked,
-          d.efc.done,
-          m.nv,
-          d.efc.cholesky_L_tmp,
-          d.efc.cholesky_y_tmp.reshape(shape=(d.nworld, m.nv, 1)),
-        ],
-        outputs=[d.efc.Mgrad.reshape(shape=(d.nworld, m.nv, 1))],
-        block_dim=m.block_dim.update_gradient_cholesky,
       )
   else:
     raise ValueError(f"Unknown solver type: {m.opt.solver}")
@@ -1951,7 +1967,7 @@ def solve_done(
   tolerance = opt_tolerance[worldid % opt_tolerance.shape[0]]
 
   improvement = _rescale(nv, stat_meaninertia, efc_prev_cost_in[worldid] - efc_cost_in[worldid])
-  gradient = _rescale(nv, stat_meaninertia, wp.math.sqrt(efc_grad_dot_in[worldid]))
+  gradient = _rescale(nv, stat_meaninertia, wp.sqrt(efc_grad_dot_in[worldid]))
   done = (improvement < tolerance) or (gradient < tolerance)
   if done or solver_niter_out[worldid] == opt_iterations:
     # if the solver has converged or the maximum number of iterations has been reached then
@@ -1964,6 +1980,8 @@ def solve_done(
 def _solver_iteration(
   m: types.Model,
   d: types.Data,
+  h: wp.array3d(dtype=float),
+  hfactor: wp.array3d(dtype=float),
   step_size_cost: wp.array2d(dtype=float),
 ):
   _linesearch(m, d, step_size_cost)
@@ -1977,7 +1995,7 @@ def _solver_iteration(
     )
 
   _update_constraint(m, d)
-  _update_gradient(m, d)
+  _update_gradient(m, d, h, hfactor)
 
   # polak-ribiere
   if m.opt.solver == types.SolverType.CG:
@@ -2014,7 +2032,9 @@ def _solver_iteration(
   )
 
 
-def create_context(m: types.Model, d: types.Data, grad: bool = True):
+def create_context(
+  m: types.Model, d: types.Data, h: wp.array3d(dtype=float), hfactor: wp.array3d(dtype=float), grad: bool = True
+):
   # initialize some efc arrays
   wp.launch(
     solve_init_efc,
@@ -2036,7 +2056,7 @@ def create_context(m: types.Model, d: types.Data, grad: bool = True):
   _update_constraint(m, d)
 
   if grad:
-    _update_gradient(m, d)
+    _update_gradient(m, d, h, hfactor)
 
 
 @event_scope
@@ -2055,8 +2075,19 @@ def _solve(m: types.Model, d: types.Data):
   else:
     wp.copy(d.qacc, d.qacc_smooth)
 
+  # Newton solver Hessian
+  if m.opt.solver == types.SolverType.NEWTON:
+    h = wp.zeros((d.nworld, m.nv_pad, m.nv_pad), dtype=float)
+    if m.nv > _BLOCK_CHOLESKY_DIM:
+      hfactor = wp.zeros((d.nworld, m.nv_pad, m.nv_pad), dtype=float)
+    else:
+      hfactor = wp.empty((d.nworld, 0, 0), dtype=float)
+  else:
+    h = wp.empty((d.nworld, 0, 0), dtype=float)
+    hfactor = wp.empty((d.nworld, 0, 0), dtype=float)
+
   # create context
-  create_context(m, d, grad=True)
+  create_context(m, d, h, hfactor, grad=True)
 
   # search = -Mgrad
   wp.launch(
@@ -2077,16 +2108,10 @@ def _solve(m: types.Model, d: types.Data):
     # becomes zero and all worlds are marked as converged to avoid an infinite loop.
     # note: we only launch the iteration kernel if everything is not done
     d.nsolving.fill_(d.nworld)
-    wp.capture_while(
+    wp.capture_while(d.nsolving, while_body=_solver_iteration, m=m, d=d, h=h, hfactor=hfactor, step_size_cost=step_size_cost)
-      d.nsolving,
-      while_body=_solver_iteration,
-      m=m,
-      d=d,
-      step_size_cost=step_size_cost,
-    )
   else:
     # This branch is mostly for when JAX is used as it is currently not compatible
     # with CUDA graph conditional.
     # It should be removed when JAX becomes compatible.
     for _ in range(m.opt.iterations):
-      _solver_iteration(m, d, step_size_cost)
+      _solver_iteration(m, d, h, hfactor, step_size_cost)

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

@@ -116,10 +116,10 @@ def mul_m_dense(tile: TileSet, check_skip: bool):
         return
 
     dofid = adr[nodeid]
-    qM_tile = wp.tile_load(qM_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
+    qM_tile = wp.tile_load(qM_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid), bounds_check=False)
-    vec_tile = wp.tile_load(vec[worldid], shape=(TILE_SIZE, 1), offset=(dofid, 0))
+    vec_tile = wp.tile_load(vec[worldid], shape=(TILE_SIZE, 1), offset=(dofid, 0), bounds_check=False)
     res_tile = wp.tile_matmul(qM_tile, vec_tile)
-    wp.tile_store(res[worldid], res_tile, offset=(dofid, 0))
+    wp.tile_store(res[worldid], res_tile, offset=(dofid, 0), bounds_check=False)
 
   return _mul_m_dense
 

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

@@ -58,8 +58,10 @@ class BlockDim:
   cholesky_factorize_solve: int = 32
   # solver
   update_gradient_cholesky: int = 64
+  update_gradient_cholesky_blocked: int = 32
   update_gradient_JTDAJ_sparse: int = 64
   update_gradient_JTDAJ_dense: int = 96
+  linesearch_iterative: int = 64
   # support
   mul_m_dense: int = 32
 
@@ -502,13 +504,15 @@ class EqType(enum.IntEnum):
     JOINT: couple the values of two scalar joints with cubic
     WELD: fix relative position and orientation of two bodies
     TENDON: couple the lengths of two tendons with cubic
+    FLEX: couple the edge lengths of a flex
   """
 
   CONNECT = mujoco.mjtEq.mjEQ_CONNECT
   WELD = mujoco.mjtEq.mjEQ_WELD
   JOINT = mujoco.mjtEq.mjEQ_JOINT
   TENDON = mujoco.mjtEq.mjEQ_TENDON
-  # unsupported: FLEX, DISTANCE
+  FLEX = mujoco.mjtEq.mjEQ_FLEX
+  # unsupported: DISTANCE
 
 
 class WrapType(enum.IntEnum):
@@ -638,7 +642,6 @@ class Option:
 
   Attributes:
     timestep: simulation timestep
-    impratio: ratio of friction-to-normal contact impedance
     tolerance: main solver tolerance
     ls_tolerance: CG/Newton linesearch tolerance
     ccd_tolerance: convex collision detection tolerance
@@ -659,6 +662,7 @@ class Option:
     sdf_iterations: max number of iterations for gradient descent
 
   warp only fields:
+    impratio_invsqrt: ratio of friction-to-normal contact impedance (stored as inverse square root)
     is_sparse: whether to use sparse representations
     ls_parallel: evaluate engine solver step sizes in parallel
     ls_parallel_min_step: minimum step size for solver linesearch
@@ -674,7 +678,6 @@ class Option:
   """
 
   timestep: array("*", float)
-  impratio: array("*", float)
   tolerance: array("*", float)
   ls_tolerance: array("*", float)
   ccd_tolerance: array("*", float)
@@ -694,6 +697,7 @@ class Option:
   sdf_initpoints: int
   sdf_iterations: int
   # warp only fields:
+  impratio_invsqrt: array("*", float)
   is_sparse: bool
   ls_parallel: bool
   ls_parallel_min_step: float
@@ -883,6 +887,9 @@ class Model:
     flex_vertadr: first vertex address                       (nflex,)
     flex_vertnum: number of vertices                         (nflex,)
     flex_edgeadr: first edge address                         (nflex,)
+    flex_edgenum: number of edges                            (nflex,)
+    flex_elemadr: first element address                      (nflex,)
+    flex_elemnum: number of elements                         (nflex,)
     flex_elemedgeadr: first element address                  (nflex,)
     flex_vertbodyid: vertex body ids                         (nflexvert,)
     flex_edge: edge vertex ids (2 per edge)                  (nflexedge, 2)
@@ -890,6 +897,7 @@ class Model:
     flex_elem: element vertex ids (dim+1 per elem)           (nflexelemdata,)
     flex_elemedge: element edge ids                          (nflexelemedge,)
     flexedge_length0: edge lengths in qpos0                  (nflexedge,)
+    flexedge_invweight0: inv. inertia for the edge           (nflexedge,)
     flex_stiffness: finite element stiffness matrix          (nflexelem, 21)
     flex_bending: bending stiffness                          (nflexedge, 17)
     flex_damping: Rayleigh's damping coefficient             (nflex,)
@@ -998,6 +1006,7 @@ class Model:
     mapM2M: index mapping from M (legacy) to M (CSR)         (nC)
 
   warp only fields:
+    nv_pad: number of degrees of freedom + padding
     nacttrnbody: number of actuators with body transmission
     nsensorcollision: number of unique collisions for
                       geom distance sensors
@@ -1031,6 +1040,7 @@ class Model:
     eq_wld_adr: eq_* addresses of type `WELD`
     eq_jnt_adr: eq_* addresses of type `JOINT`
     eq_ten_adr: eq_* addresses of type `TENDON`
+    eq_flex_adr: eq * addresses of type `FLEX
     tendon_jnt_adr: joint tendon address
     tendon_site_pair_adr: site pair tendon address
     tendon_geom_adr: geom tendon address
@@ -1225,6 +1235,9 @@ class Model:
   flex_vertadr: array("nflex", int)
   flex_vertnum: array("nflex", int)
   flex_edgeadr: array("nflex", int)
+  flex_edgenum: array("nflex", int)
+  flex_elemadr: array("nflex", int)
+  flex_elemnum: array("nflex", int)
   flex_elemedgeadr: array("nflex", int)
   flex_vertbodyid: array("nflexvert", int)
   flex_edge: array("nflexedge", wp.vec2i)
@@ -1232,6 +1245,7 @@ class Model:
   flex_elem: array("nflexelemdata", int)
   flex_elemedge: array("nflexelemedge", int)
   flexedge_length0: array("nflexedge", float)
+  flexedge_invweight0: array("nflexedge", float)
   flex_stiffness: array("nflexelem", 21, float)
   flex_bending: array("nflexedge", 17, float)
   flex_damping: array("nflex", float)
@@ -1339,6 +1353,7 @@ class Model:
   M_colind: array("nC", int)
   mapM2M: array("nC", int)
   # warp only fields:
+  nv_pad: int
   nacttrnbody: int
   nsensorcollision: int
   nsensortaxel: int
@@ -1367,6 +1382,7 @@ class Model:
   eq_wld_adr: wp.array(dtype=int)
   eq_jnt_adr: wp.array(dtype=int)
   eq_ten_adr: wp.array(dtype=int)
+  eq_flex_adr: wp.array(dtype=int)
   tendon_jnt_adr: wp.array(dtype=int)
   tendon_site_pair_adr: wp.array(dtype=int)
   tendon_geom_adr: wp.array(dtype=int)
@@ -1476,11 +1492,9 @@ class Constraint:
     force: constraint force in constraint space       (nworld, njmax)
     Jaref: Jac*qacc - aref                            (nworld, njmax)
     Ma: M*qacc                                        (nworld, nv)
-    grad: gradient of master cost                     (nworld, nv)
+    grad: gradient of master cost                     (nworld, nv_pad)
-    cholesky_L_tmp: temporary for Cholesky factor     (nworld, nv, nv)
-    cholesky_y_tmp: temporary for Cholesky solve      (nworld, nv
     grad_dot: dot(grad, grad)                         (nworld,)
-    Mgrad: M / grad                                   (nworld, nv)
+    Mgrad: M / grad                                   (nworld, nv_pad)
     search: linesearch vector                         (nworld, nv)
     search_dot: dot(search, search)                   (nworld,)
     gauss: Gauss Cost                                 (nworld,)
@@ -1491,7 +1505,6 @@ class Constraint:
     jv: efc_J @ search                                (nworld, njmax)
     quad: quadratic cost coefficients                 (nworld, njmax, 3)
     quad_gauss: quadratic cost Gauss coefficients     (nworld, 3)
-    h: Hessian                                        (nworld, nv_pad, nv_pad)
     alpha: line search step size                      (nworld,)
     prev_grad: previous grad                          (nworld, nv)
     prev_Mgrad: previous Mgrad                        (nworld, nv)
@@ -1511,11 +1524,9 @@ class Constraint:
   force: array("nworld", "njmax", float)
   Jaref: array("nworld", "njmax", float)
   Ma: array("nworld", "nv", float)
-  grad: array("nworld", "nv", float)
+  grad: array("nworld", "nv_pad", float)
-  cholesky_L_tmp: array("nworld", "nv", "nv", float)
-  cholesky_y_tmp: array("nworld", "nv", float)
   grad_dot: array("nworld", float)
-  Mgrad: array("nworld", "nv", float)
+  Mgrad: array("nworld", "nv_pad", float)
   search: array("nworld", "nv", float)
   search_dot: array("nworld", float)
   gauss: array("nworld", float)
@@ -1526,7 +1537,6 @@ class Constraint:
   jv: array("nworld", "njmax", float)
   quad: array("nworld", "njmax", wp.vec3)
   quad_gauss: array("nworld", wp.vec3)
-  h: array("nworld", "nv_pad", "nv_pad", float)
   alpha: array("nworld", float)
   prev_grad: array("nworld", "nv", float)
   prev_Mgrad: array("nworld", "nv", float)
@@ -1578,6 +1588,7 @@ class Data:
     cdof: com-based motion axis of each dof (rot:lin)           (nworld, nv, 6)
     cinert: com-based body inertia and mass                     (nworld, nbody, 10)
     flexvert_xpos: cartesian flex vertex positions              (nworld, nflexvert, 3)
+    flexedge_J: edge length Jacobian                            (nworld, nflexedge, nv)
     flexedge_length: flex edge lengths                          (nworld, nflexedge, 1)
     ten_wrapadr: start address of tendon's path                 (nworld, ntendon)
     ten_wrapnum: number of wrap points in path                  (nworld, ntendon)
@@ -1629,6 +1640,7 @@ class Data:
     ne_weld: number of equality weld constraints                (nworld,)
     ne_jnt: number of equality joint constraints                (nworld,)
     ne_ten: number of equality tendon constraints               (nworld,)
+    ne_flex: number of flex edge equality constraints           (nworld,)
     nsolving: number of unconverged worlds                      (1,)
     subtree_bodyvel: subtree body velocity (ang, vel)           (nworld, nbody, 6)
     collision_pair: collision pairs from broadphase             (naconmax, 2)
@@ -1676,6 +1688,7 @@ class Data:
   cdof: array("nworld", "nv", wp.spatial_vector)
   cinert: array("nworld", "nbody", vec10)
   flexvert_xpos: array("nworld", "nflexvert", wp.vec3)
+  flexedge_J: array("nworld", "nflexedge", "nv", float)
   flexedge_length: array("nworld", "nflexedge", float)
   ten_wrapadr: array("nworld", "ntendon", int)
   ten_wrapnum: array("nworld", "ntendon", int)
@@ -1723,6 +1736,7 @@ class Data:
   ne_weld: array("nworld", int)
   ne_jnt: array("nworld", int)
   ne_ten: array("nworld", int)
+  ne_flex: array("nworld", int)
   nsolving: array(1, int)
   subtree_bodyvel: array("nworld", "nbody", wp.spatial_vector)
 

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

@@ -18,8 +18,8 @@ import inspect
 from typing import Callable, Optional
 
 import warp as wp
-from warp.context import Module
+from warp._src.context import Module
-from warp.context import get_module
+from warp._src.context import get_module
 
 _STACK = None
 
@@ -220,8 +220,8 @@ def cache_kernel(func):
 
 
 def check_toolkit_driver():
-  if wp.context.runtime is None:
+  if wp._src.context.runtime is None:
-    wp.context.init()
+    wp._src.context.init()
   if wp.get_device().is_cuda:
-    if wp.context.runtime.toolkit_version < (12, 4) or wp.context.runtime.driver_version < (12, 4):
+    if wp._src.context.runtime.toolkit_version < (12, 4) or wp._src.context.runtime.driver_version < (12, 4):
       RuntimeError("Minimum supported CUDA version: 12.4.")

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

@@ -18,13 +18,12 @@ classifiers = [
     "Intended Audience :: Science/Research",
     "Natural Language :: English",
     "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Topic :: Scientific/Engineering",
 ]
-requires-python = ">=3.9"
+requires-python = ">=3.10"
 dependencies = [
     "absl-py",
     "etils[epath]",
@@ -35,7 +34,7 @@ dependencies = [
 
 [[tool.uv.index]]
 name = "nvidia"
-url = "https://pypi.nvidia.com"
+url = "https://pypi.nvidia.com/"
 explicit = true
 
 [[tool.uv.index]]
@@ -56,6 +55,8 @@ dev = [
     "ruff",
     "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",
 ]
 # TODO(team): cpu and cuda JAX optional dependencies are temporary, remove after we land MJX:Warp
 cpu = [

mjx/mujoco/mjx/warp/forward_test.py

@@ -101,6 +101,7 @@ class ForwardTest(parameterized.TestCase):
     m = test_util.load_test_file(xml)
     m.opt.iterations = 10
     m.opt.ls_iterations = 10
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE
     mx = mjx.put_model(m, impl='warp')
 
     d = mujoco.MjData(m)
@@ -150,11 +151,8 @@ class ForwardTest(parameterized.TestCase):
       tu.assert_attr_eq(dx._impl, d, 'wrap_obj')
       tu.assert_attr_eq(dx._impl, d, 'crb')
 
-      if not mx.opt._impl.is_sparse:
+      qm = np.zeros((m.nv, m.nv))
-        qm = np.zeros((m.nv, m.nv))
+      mujoco.mj_fullM(m, qm, d.qM)
-        mujoco.mj_fullM(m, qm, d.qM)
-      else:
-        qm = d.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.

mjx/mujoco/mjx/warp/smooth.py

@@ -46,7 +46,6 @@ _e = mjwarp.Constraint(
 def _kinematics_shim(
     # Model
     nworld: int,
-    body_dofadr: wp.array(dtype=int),
     body_ipos: wp.array2d(dtype=wp.vec3),
     body_iquat: wp.array2d(dtype=wp.quat),
     body_jntadr: wp.array(dtype=int),
@@ -58,9 +57,6 @@ def _kinematics_shim(
     body_rootid: wp.array(dtype=int),
     body_tree: tuple[wp.array(dtype=int), ...],
     body_weldid: wp.array(dtype=int),
-    flex_edge: wp.array(dtype=wp.vec2i),
-    flex_vertadr: wp.array(dtype=int),
-    flex_vertbodyid: wp.array(dtype=int),
     geom_bodyid: wp.array(dtype=int),
     geom_pos: wp.array2d(dtype=wp.vec3),
     geom_quat: wp.array2d(dtype=wp.quat),
@@ -68,26 +64,18 @@ def _kinematics_shim(
     jnt_pos: wp.array2d(dtype=wp.vec3),
     jnt_qposadr: wp.array(dtype=int),
     jnt_type: wp.array(dtype=int),
-    mocap_bodyid: wp.array(dtype=int),
-    nflexedge: int,
-    nflexvert: int,
     ngeom: int,
-    nmocap: int,
     nsite: int,
     qpos0: wp.array2d(dtype=float),
     site_bodyid: wp.array(dtype=int),
     site_pos: wp.array2d(dtype=wp.vec3),
     site_quat: wp.array2d(dtype=wp.quat),
     # Data
-    flexedge_length: wp.array2d(dtype=float),
-    flexedge_velocity: wp.array2d(dtype=float),
-    flexvert_xpos: wp.array2d(dtype=wp.vec3),
     geom_xmat: wp.array2d(dtype=wp.mat33),
     geom_xpos: wp.array2d(dtype=wp.vec3),
     mocap_pos: wp.array2d(dtype=wp.vec3),
     mocap_quat: wp.array2d(dtype=wp.quat),
     qpos: wp.array2d(dtype=float),
-    qvel: wp.array2d(dtype=float),
     site_xmat: wp.array2d(dtype=wp.mat33),
     site_xpos: wp.array2d(dtype=wp.vec3),
     xanchor: wp.array2d(dtype=wp.vec3),
@@ -102,7 +90,6 @@ def _kinematics_shim(
   _m.opt = _o
   _d.efc = _e
   _d.contact = _c
-  _m.body_dofadr = body_dofadr
   _m.body_ipos = body_ipos
   _m.body_iquat = body_iquat
   _m.body_jntadr = body_jntadr
@@ -114,9 +101,6 @@ def _kinematics_shim(
   _m.body_rootid = body_rootid
   _m.body_tree = body_tree
   _m.body_weldid = body_weldid
-  _m.flex_edge = flex_edge
-  _m.flex_vertadr = flex_vertadr
-  _m.flex_vertbodyid = flex_vertbodyid
   _m.geom_bodyid = geom_bodyid
   _m.geom_pos = geom_pos
   _m.geom_quat = geom_quat
@@ -124,25 +108,17 @@ def _kinematics_shim(
   _m.jnt_pos = jnt_pos
   _m.jnt_qposadr = jnt_qposadr
   _m.jnt_type = jnt_type
-  _m.mocap_bodyid = mocap_bodyid
-  _m.nflexedge = nflexedge
-  _m.nflexvert = nflexvert
   _m.ngeom = ngeom
-  _m.nmocap = nmocap
   _m.nsite = nsite
   _m.qpos0 = qpos0
   _m.site_bodyid = site_bodyid
   _m.site_pos = site_pos
   _m.site_quat = site_quat
-  _d.flexedge_length = flexedge_length
-  _d.flexedge_velocity = flexedge_velocity
-  _d.flexvert_xpos = flexvert_xpos
   _d.geom_xmat = geom_xmat
   _d.geom_xpos = geom_xpos
   _d.mocap_pos = mocap_pos
   _d.mocap_quat = mocap_quat
   _d.qpos = qpos
-  _d.qvel = qvel
   _d.site_xmat = site_xmat
   _d.site_xpos = site_xpos
   _d.xanchor = xanchor
@@ -158,15 +134,11 @@ def _kinematics_shim(
 
 def _kinematics_jax_impl(m: types.Model, d: types.Data):
   output_dims = {
-      'flexedge_length': d._impl.flexedge_length.shape,
-      'flexedge_velocity': d._impl.flexedge_velocity.shape,
-      'flexvert_xpos': d._impl.flexvert_xpos.shape,
       'geom_xmat': d.geom_xmat.shape,
       'geom_xpos': d.geom_xpos.shape,
       'mocap_pos': d.mocap_pos.shape,
       'mocap_quat': d.mocap_quat.shape,
       'qpos': d.qpos.shape,
-      'qvel': d.qvel.shape,
       'site_xmat': d.site_xmat.shape,
       'site_xpos': d.site_xpos.shape,
       'xanchor': d.xanchor.shape,
@@ -179,19 +151,15 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
   }
   jf = ffi.jax_callable_variadic_tuple(
       _kinematics_shim,
-      num_outputs=18,
+      num_outputs=14,
       output_dims=output_dims,
       vmap_method=None,
       in_out_argnames={
-          'flexedge_length',
-          'flexedge_velocity',
-          'flexvert_xpos',
           'geom_xmat',
           'geom_xpos',
           'mocap_pos',
           'mocap_quat',
           'qpos',
-          'qvel',
           'site_xmat',
           'site_xpos',
           'xanchor',
@@ -205,7 +173,6 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
   )
   out = jf(
       d.qpos.shape[0],
-      m.body_dofadr,
       m.body_ipos,
       m.body_iquat,
       m.body_jntadr,
@@ -217,9 +184,6 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
       m.body_rootid,
       m._impl.body_tree,
       m.body_weldid,
-      m._impl.flex_edge,
-      m._impl.flex_vertadr,
-      m._impl.flex_vertbodyid,
       m.geom_bodyid,
       m.geom_pos,
       m.geom_quat,
@@ -227,25 +191,17 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
       m.jnt_pos,
       m.jnt_qposadr,
       m.jnt_type,
-      m._impl.mocap_bodyid,
-      m._impl.nflexedge,
-      m._impl.nflexvert,
       m.ngeom,
-      m.nmocap,
       m.nsite,
       m.qpos0,
       m.site_bodyid,
       m.site_pos,
       m.site_quat,
-      d._impl.flexedge_length,
-      d._impl.flexedge_velocity,
-      d._impl.flexvert_xpos,
       d.geom_xmat,
       d.geom_xpos,
       d.mocap_pos,
       d.mocap_quat,
       d.qpos,
-      d.qvel,
       d.site_xmat,
       d.site_xpos,
       d.xanchor,
@@ -257,24 +213,20 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
       d.xquat,
   )
   d = d.tree_replace({
-      '_impl.flexedge_length': out[0],
+      'geom_xmat': out[0],
-      '_impl.flexedge_velocity': out[1],
+      'geom_xpos': out[1],
-      '_impl.flexvert_xpos': out[2],
+      'mocap_pos': out[2],
-      'geom_xmat': out[3],
+      'mocap_quat': out[3],
-      'geom_xpos': out[4],
+      'qpos': out[4],
-      'mocap_pos': out[5],
+      'site_xmat': out[5],
-      'mocap_quat': out[6],
+      'site_xpos': out[6],
-      'qpos': out[7],
+      'xanchor': out[7],
-      'qvel': out[8],
+      'xaxis': out[8],
-      'site_xmat': out[9],
+      'ximat': out[9],
-      'site_xpos': out[10],
+      'xipos': out[10],
-      'xanchor': out[11],
+      'xmat': out[11],
-      'xaxis': out[12],
+      'xpos': out[12],
-      'ximat': out[13],
+      'xquat': out[13],
-      'xipos': out[14],
-      'xmat': out[15],
-      'xpos': out[16],
-      'xquat': out[17],
   })
   return d
 

mjx/mujoco/mjx/warp/smooth_test.py

@@ -73,7 +73,7 @@ class SmoothTest(parameterized.TestCase):
     _, key1, key2 = jax.random.split(rng, 3)
     mocap_pos = jax.random.normal(key1, (m.nmocap, 3))
     mocap_quat = jax.random.normal(key2, (m.nmocap, 4))
-    mocap_quat = math.normalize(mocap_quat)
+    mocap_quat = math.normalize(mocap_quat, axis=0)
     dx = dx.replace(qpos=qpos, mocap_pos=mocap_pos, mocap_quat=mocap_quat)
 
     dx = jax.jit(smooth.kinematics)(mx, dx)
@@ -81,7 +81,7 @@ class SmoothTest(parameterized.TestCase):
     d.qpos[:] = qpos
     d.mocap_pos[:] = mocap_pos
     d.mocap_quat[:] = mocap_quat
-    mujoco.mj_forward(m, d)
+    mujoco.mj_kinematics(m, d)
 
     tu.assert_attr_eq(d, dx, 'xanchor')
     tu.assert_attr_eq(d, dx, 'xaxis')
@@ -110,9 +110,9 @@ class SmoothTest(parameterized.TestCase):
 
     worldids = jp.arange(batch_size)
     dx_batch = jax.vmap(functools.partial(tu.make_data, m))(worldids)
-    fields = ('xanchor', 'xaxis', 'xpos', 'xquat', 'xmat', 'xipos', 'ximat',
+    # don't zero xmat, ximat, xquat, geom_xpos, or geom_xmat
-              'geom_xpos', 'geom_xmat', 'site_xpos', 'site_xmat')  # fmt: skip
+    # these fields are precomputed in make_data
-    for f in fields:
+    for f in ('xanchor', 'xaxis', 'xpos', 'xipos', 'site_xpos', 'site_xmat'):
       dx_batch = dx_batch.replace(**{f: jp.zeros_like(getattr(dx_batch, f))})
 
     dx_batch = jax.jit(jax.vmap(smooth.kinematics, in_axes=(None, 0)))(
@@ -125,7 +125,7 @@ class SmoothTest(parameterized.TestCase):
       d.qpos[:] = dx.qpos
       d.mocap_pos[:] = dx.mocap_pos
       d.mocap_quat[:] = dx.mocap_quat
-      mujoco.mj_forward(m, d)
+      mujoco.mj_kinematics(m, d)
 
       tu.assert_attr_eq(d, dx, 'xanchor')
       tu.assert_attr_eq(d, dx, 'xaxis')
@@ -154,9 +154,9 @@ class SmoothTest(parameterized.TestCase):
 
     worldids = jp.arange(16).reshape((4, 4))
     dx_batch = jax.vmap(jax.vmap(functools.partial(tu.make_data, m)))(worldids)
-    fields = ('xanchor', 'xaxis', 'xpos', 'xquat', 'xmat', 'xipos', 'ximat',
+    # don't zero xmat, ximat, xquat, geom_xpos, or geom_xmat
-              'geom_xpos', 'geom_xmat', 'site_xpos', 'site_xmat')  # fmt: skip
+    # these fields are precomputed in make_data
-    for f in fields:
+    for f in ('xanchor', 'xaxis', 'xpos', 'xipos', 'site_xpos', 'site_xmat'):
       dx_batch = dx_batch.replace(**{f: jp.zeros_like(getattr(dx_batch, f))})
 
     dx_batch = jax.jit(
@@ -208,9 +208,9 @@ class SmoothTest(parameterized.TestCase):
 
     worldids = jp.arange(batch_size)
     dx_batch = jax.vmap(functools.partial(tu.make_data, m))(worldids)
-    fields = ('xanchor', 'xaxis', 'xpos', 'xquat', 'xmat', 'xipos', 'ximat',
+    # don't zero xmat, ximat, xquat, geom_xpos, or geom_xmat
-              'geom_xpos', 'geom_xmat', 'site_xpos', 'site_xmat')  # fmt: skip
+    # these fields are precomputed in make_data
-    for f in fields:
+    for f in ('xanchor', 'xaxis', 'xpos', 'xipos', 'site_xpos', 'site_xmat'):
       dx_batch = dx_batch.replace(**{f: jp.zeros_like(getattr(dx_batch, f))})
 
     dx_batch = jax.jit(jax.vmap(smooth.kinematics, in_axes=(None, 0)))(
@@ -224,7 +224,7 @@ class SmoothTest(parameterized.TestCase):
       d.mocap_pos[:] = dx.mocap_pos
       d.mocap_quat[:] = dx.mocap_quat
       m.geom_pos[:] = mx.geom_pos[i]
-      mujoco.mj_forward(m, d)
+      mujoco.mj_kinematics(m, d)
 
       tu.assert_attr_eq(d, dx, 'xanchor')
       tu.assert_attr_eq(d, dx, 'xaxis')

mjx/mujoco/mjx/warp/types.py

@@ -62,6 +62,7 @@ class BlockDim:
   contact_sort: int
   energy_vel_kinetic: int
   euler_dense: int
+  linesearch_iterative: int
   mul_m_dense: int
   ray: int
   segmented_sort: int
@@ -69,6 +70,7 @@ class BlockDim:
   update_gradient_JTDAJ_dense: int
   update_gradient_JTDAJ_sparse: int
   update_gradient_cholesky: int
+  update_gradient_cholesky_blocked: int
 
   def tree_flatten(self):
     children = list((getattr(self, k) for k in self.__dataclass_fields__))
@@ -93,6 +95,7 @@ class OptionWarp(PyTreeNode):
   contact_sensor_maxmatch: int
   graph_conditional: bool
   has_fluid: bool
+  impratio_invsqrt: jax.Array
   is_sparse: bool
   ls_parallel: bool
   ls_parallel_min_step: float
@@ -113,6 +116,7 @@ class ModelWarp(PyTreeNode):
   dof_tri_col: np.ndarray
   dof_tri_row: np.ndarray
   eq_connect_adr: np.ndarray
+  eq_flex_adr: np.ndarray
   eq_jnt_adr: np.ndarray
   eq_ten_adr: np.ndarray
   eq_wld_adr: np.ndarray
@@ -122,13 +126,17 @@ class ModelWarp(PyTreeNode):
   flex_edge: np.ndarray
   flex_edgeadr: np.ndarray
   flex_edgeflap: np.ndarray
+  flex_edgenum: np.ndarray
   flex_elem: np.ndarray
+  flex_elemadr: np.ndarray
   flex_elemedge: np.ndarray
   flex_elemedgeadr: np.ndarray
+  flex_elemnum: np.ndarray
   flex_stiffness: np.ndarray
   flex_vertadr: np.ndarray
   flex_vertbodyid: np.ndarray
   flex_vertnum: np.ndarray
+  flexedge_invweight0: np.ndarray
   flexedge_length0: np.ndarray
   geom_pair_type_count: Tuple[int, ...]
   geom_plugin_index: np.ndarray
@@ -167,6 +175,7 @@ class ModelWarp(PyTreeNode):
   nsensorcollision: int
   nsensorcontact: int
   nsensortaxel: int
+  nv_pad: int
   nxn_geom_pair: np.ndarray
   nxn_geom_pair_filtered: np.ndarray
   nxn_pairid: np.ndarray
@@ -249,8 +258,6 @@ class DataWarp(PyTreeNode):
   efc__alpha: jax.Array
   efc__aref: jax.Array
   efc__beta: jax.Array
-  efc__cholesky_L_tmp: jax.Array
-  efc__cholesky_y_tmp: jax.Array
   efc__cost: jax.Array
   efc__done: jax.Array
   efc__force: jax.Array
@@ -258,7 +265,6 @@ class DataWarp(PyTreeNode):
   efc__gauss: jax.Array
   efc__grad: jax.Array
   efc__grad_dot: jax.Array
-  efc__h: jax.Array
   efc__id: jax.Array
   efc__jv: jax.Array
   efc__margin: jax.Array
@@ -275,6 +281,7 @@ class DataWarp(PyTreeNode):
   efc__type: jax.Array
   efc__vel: jax.Array
   energy: jax.Array
+  flexedge_J: jax.Array
   flexedge_length: jax.Array
   flexedge_velocity: jax.Array
   flexvert_xpos: jax.Array
@@ -285,6 +292,7 @@ class DataWarp(PyTreeNode):
   ncollision: jax.Array
   ne: jax.Array
   ne_connect: jax.Array
+  ne_flex: jax.Array
   ne_jnt: jax.Array
   ne_ten: jax.Array
   ne_weld: jax.Array
@@ -404,8 +412,6 @@ _NDIM = {
         'efc__alpha': 1,
         'efc__aref': 2,
         'efc__beta': 1,
-        'efc__cholesky_L_tmp': 3,
-        'efc__cholesky_y_tmp': 2,
         'efc__cost': 1,
         'efc__done': 1,
         'efc__force': 2,
@@ -413,7 +419,6 @@ _NDIM = {
         'efc__gauss': 1,
         'efc__grad': 2,
         'efc__grad_dot': 1,
-        'efc__h': 3,
         'efc__id': 2,
         'efc__jv': 2,
         'efc__margin': 2,
@@ -431,6 +436,7 @@ _NDIM = {
         'efc__vel': 2,
         'energy': 2,
         'eq_active': 2,
+        'flexedge_J': 3,
         'flexedge_length': 2,
         'flexedge_velocity': 2,
         'flexvert_xpos': 3,
@@ -445,6 +451,7 @@ _NDIM = {
         'ncollision': 1,
         'ne': 1,
         'ne_connect': 1,
+        'ne_flex': 1,
         'ne_jnt': 1,
         'ne_ten': 1,
         'ne_weld': 1,
@@ -531,6 +538,7 @@ _NDIM = {
         'block_dim__contact_sort': 0,
         'block_dim__energy_vel_kinetic': 0,
         'block_dim__euler_dense': 0,
+        'block_dim__linesearch_iterative': 0,
         'block_dim__mul_m_dense': 0,
         'block_dim__ray': 0,
         'block_dim__segmented_sort': 0,
@@ -538,6 +546,7 @@ _NDIM = {
         'block_dim__update_gradient_JTDAJ_dense': 0,
         'block_dim__update_gradient_JTDAJ_sparse': 0,
         'block_dim__update_gradient_cholesky': 0,
+        'block_dim__update_gradient_cholesky_blocked': 0,
         'body_conaffinity': 1,
         'body_contype': 1,
         'body_dofadr': 1,
@@ -589,6 +598,7 @@ _NDIM = {
         'eq_active0': 1,
         'eq_connect_adr': 1,
         'eq_data': 3,
+        'eq_flex_adr': 1,
         'eq_jnt_adr': 1,
         'eq_obj1id': 1,
         'eq_obj2id': 1,
@@ -605,13 +615,17 @@ _NDIM = {
         'flex_edge': 2,
         'flex_edgeadr': 1,
         'flex_edgeflap': 2,
+        'flex_edgenum': 1,
         'flex_elem': 1,
+        'flex_elemadr': 1,
         'flex_elemedge': 1,
         'flex_elemedgeadr': 1,
+        'flex_elemnum': 1,
         'flex_stiffness': 2,
         'flex_vertadr': 1,
         'flex_vertbodyid': 1,
         'flex_vertnum': 1,
+        'flexedge_invweight0': 1,
         'flexedge_length0': 1,
         'geom_aabb': 4,
         'geom_bodyid': 1,
@@ -743,6 +757,7 @@ _NDIM = {
         'ntendon': 0,
         'nu': 0,
         'nv': 0,
+        'nv_pad': 0,
         'nwrap': 0,
         'nxn_geom_pair': 2,
         'nxn_geom_pair_filtered': 2,
@@ -761,7 +776,7 @@ _NDIM = {
         'opt__graph_conditional': 0,
         'opt__gravity': 2,
         'opt__has_fluid': 0,
-        'opt__impratio': 1,
+        'opt__impratio_invsqrt': 1,
         'opt__integrator': 0,
         'opt__is_sparse': 0,
         'opt__iterations': 0,
@@ -878,7 +893,7 @@ _NDIM = {
         'graph_conditional': 0,
         'gravity': 2,
         'has_fluid': 0,
-        'impratio': 1,
+        'impratio_invsqrt': 1,
         'integrator': 0,
         'is_sparse': 0,
         'iterations': 0,
@@ -942,8 +957,6 @@ _BATCH_DIM = {
         'efc__alpha': True,
         'efc__aref': True,
         'efc__beta': True,
-        'efc__cholesky_L_tmp': True,
-        'efc__cholesky_y_tmp': True,
         'efc__cost': True,
         'efc__done': True,
         'efc__force': True,
@@ -951,7 +964,6 @@ _BATCH_DIM = {
         'efc__gauss': True,
         'efc__grad': True,
         'efc__grad_dot': True,
-        'efc__h': True,
         'efc__id': True,
         'efc__jv': True,
         'efc__margin': True,
@@ -969,6 +981,7 @@ _BATCH_DIM = {
         'efc__vel': True,
         'energy': True,
         'eq_active': True,
+        'flexedge_J': True,
         'flexedge_length': True,
         'flexedge_velocity': True,
         'flexvert_xpos': True,
@@ -983,6 +996,7 @@ _BATCH_DIM = {
         'ncollision': False,
         'ne': True,
         'ne_connect': True,
+        'ne_flex': True,
         'ne_jnt': True,
         'ne_ten': True,
         'ne_weld': True,
@@ -1069,6 +1083,7 @@ _BATCH_DIM = {
         'block_dim__contact_sort': False,
         'block_dim__energy_vel_kinetic': False,
         'block_dim__euler_dense': False,
+        'block_dim__linesearch_iterative': False,
         'block_dim__mul_m_dense': False,
         'block_dim__ray': False,
         'block_dim__segmented_sort': False,
@@ -1076,6 +1091,7 @@ _BATCH_DIM = {
         'block_dim__update_gradient_JTDAJ_dense': False,
         'block_dim__update_gradient_JTDAJ_sparse': False,
         'block_dim__update_gradient_cholesky': False,
+        'block_dim__update_gradient_cholesky_blocked': False,
         'body_conaffinity': False,
         'body_contype': False,
         'body_dofadr': False,
@@ -1127,6 +1143,7 @@ _BATCH_DIM = {
         'eq_active0': False,
         'eq_connect_adr': False,
         'eq_data': True,
+        'eq_flex_adr': False,
         'eq_jnt_adr': False,
         'eq_obj1id': False,
         'eq_obj2id': False,
@@ -1143,13 +1160,17 @@ _BATCH_DIM = {
         'flex_edge': False,
         'flex_edgeadr': False,
         'flex_edgeflap': False,
+        'flex_edgenum': False,
         'flex_elem': False,
+        'flex_elemadr': False,
         'flex_elemedge': False,
         'flex_elemedgeadr': False,
+        'flex_elemnum': False,
         'flex_stiffness': False,
         'flex_vertadr': False,
         'flex_vertbodyid': False,
         'flex_vertnum': False,
+        'flexedge_invweight0': False,
         'flexedge_length0': False,
         'geom_aabb': True,
         'geom_bodyid': False,
@@ -1281,6 +1302,7 @@ _BATCH_DIM = {
         'ntendon': False,
         'nu': False,
         'nv': False,
+        'nv_pad': False,
         'nwrap': False,
         'nxn_geom_pair': False,
         'nxn_geom_pair_filtered': False,
@@ -1299,7 +1321,7 @@ _BATCH_DIM = {
         'opt__graph_conditional': False,
         'opt__gravity': True,
         'opt__has_fluid': False,
-        'opt__impratio': True,
+        'opt__impratio_invsqrt': True,
         'opt__integrator': False,
         'opt__is_sparse': False,
         'opt__iterations': False,
@@ -1416,7 +1438,7 @@ _BATCH_DIM = {
         'graph_conditional': False,
         'gravity': True,
         'has_fluid': False,
-        'impratio': True,
+        'impratio_invsqrt': True,
         'integrator': False,
         'is_sparse': False,
         'iterations': False,