Optimization: replace mat33 with wp.quat, move to quaternion math

mujoco_warp/_src/broadphase_test.py

@@ -15,6 +15,7 @@
 
 """Tests for broadphase functions."""
 
+import mujoco
 import numpy as np
 import warp as wp
 from absl.testing import absltest
@@ -145,10 +146,11 @@ def test_broadphase(self, broadphase, filter):
       np.vstack([np.expand_dims(mjd1.geom_xpos, axis=0), np.expand_dims(mjd2.geom_xpos, axis=0)]),
       dtype=wp.vec3,
     )
-    d3.geom_xmat = wp.array(
-      np.vstack([np.expand_dims(mjd1.geom_xmat, axis=0), np.expand_dims(mjd2.geom_xmat, axis=0)]),
-      dtype=wp.mat33,
-    )
+    geom_xquat = np.zeros((2, mjm.ngeom, 4))
+    for j, mjd in enumerate([mjd1, mjd2]):
+      for i in range(mjm.ngeom):
+        mujoco.mju_mat2Quat(geom_xquat[j, i], mjd.geom_xmat[i])
+    d3.geom_xquat = wp.array(geom_xquat, dtype=wp.quat)
     broadphase_caller(m, d3)
 
     ncollision = d3.ncollision.numpy()[0]

mujoco_warp/_src/collision_convex.py

@@ -25,6 +25,9 @@
 from .collision_primitive import geom_collision_pair
 from .collision_primitive import write_contact
 from .math import make_frame
+from .math import mul_quat
+from .math import quat_inv
+from .math import rot_vec_quat
 from .math import upper_trid_index
 from .types import MJ_MAX_EPAFACES
 from .types import MJ_MAX_EPAHORIZON
@@ -93,7 +96,7 @@ def _hfield_filter(
   hfield_size: wp.array(dtype=wp.vec4),
   # Data in:
   geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  geom_xquat_in: wp.array2d(dtype=wp.quat),
   # In:
   worldid: int,
   g1: int,
@@ -113,10 +116,10 @@ def _hfield_filter(
   size_id = worldid % geom_size.shape[0]
 
   pos1 = geom_xpos_in[worldid, g1]
-  mat1 = geom_xmat_in[worldid, g1]
-  mat1T = wp.transpose(mat1)
   pos2 = geom_xpos_in[worldid, g2]
-  pos = mat1T @ (pos2 - pos1)
+  quat1 = geom_xquat_in[worldid, g1]
+  rot1_inv = quat_inv(quat1)
+  pos = rot_vec_quat(pos2 - pos1, rot1_inv)
   r2 = geom_rbound[rbound_id, g2]
 
   # TODO(team): margin?
@@ -134,13 +137,13 @@ def _hfield_filter(
   if -size1[3] > pos[2] + r2 + margin:  # down
     return True, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf
 
-  mat2 = geom_xmat_in[worldid, g2]
-  mat = mat1T @ mat2
+  quat2 = geom_xquat_in[worldid, g2]
+  rot = mul_quat(quat1, quat2)
 
   # create geom in height field frame for support function queries
   geom2 = Geom()
   geom2.pos = pos
-  geom2.rot = mat
+  geom2.rot = rot
   geom2.size = geom_size[size_id, g2]
   geom2.margin = 0.0  # margin handled separately
   geom2.index = -1
@@ -224,7 +227,7 @@ def ccd_hfield_kernel(
     # Data in:
     naconmax_in: int,
     geom_xpos_in: wp.array2d(dtype=wp.vec3),
-    geom_xmat_in: wp.array2d(dtype=wp.mat33),
+    geom_xquat_in: wp.array2d(dtype=wp.quat),
     collision_pair_in: wp.array(dtype=wp.vec2i),
     collision_pairid_in: wp.array(dtype=wp.vec2i),
     collision_worldid_in: wp.array(dtype=int),
@@ -270,7 +273,7 @@ def ccd_hfield_kernel(
 
     # height field filter
     no_hf_collision, xmin, xmax, ymin, ymax, zmin, zmax = _hfield_filter(
-      geom_type, geom_dataid, geom_size, geom_rbound, geom_margin, hfield_size, geom_xpos_in, geom_xmat_in, worldid, g1, g2
+      geom_type, geom_dataid, geom_size, geom_rbound, geom_margin, hfield_size, geom_xpos_in, geom_xquat_in, worldid, g1, g2
     )
     if no_hf_collision:
       return
@@ -316,22 +319,22 @@ def ccd_hfield_kernel(
       mesh_polymapnum,
       mesh_polymap,
       geom_xpos_in,
-      geom_xmat_in,
+      geom_xquat_in,
       geoms,
       worldid,
     )
 
     # transform geom2 into heightfield frame
     hf_pos = geom_xpos_in[worldid, g1]
-    hf_mat = geom_xmat_in[worldid, g1]
-    hf_matT = wp.transpose(hf_mat)
+    hf_quat = geom_xquat_in[worldid, g1]
+    hf_quat_inv = quat_inv(hf_quat)
 
-    geom2.pos = hf_matT @ (geom2.pos - hf_pos)
-    geom2.rot = hf_matT @ geom2.rot
+    geom2.pos = rot_vec_quat(geom2.pos - hf_pos, hf_quat_inv)
+    geom2.rot = mul_quat(hf_quat_inv, geom2.rot)
 
     # geom1 has identity pose
     geom1.pos = wp.vec3(0.0, 0.0, 0.0)
-    geom1.rot = wp.identity(n=3, dtype=float)
+    geom1.rot = wp.quat(1.0, 0.0, 0.0, 0.0)
 
     # see MuJoCo mjc_ConvexHField
     geom1_dataid = geom_dataid[g1]
@@ -444,7 +447,7 @@ def ccd_hfield_kernel(
           x1_ = wp.vec3(0.0, 0.0, 0.0)
           for i in range(6):
             x1_ += prism[i]
-          x1 += geom1.rot @ (x1_ / 6.0)
+          x1 += rot_vec_quat(x1_ / 6.0, geom1.rot)
 
           dist, ncontact, w1, w2, idx = ccd(
             opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
@@ -476,14 +479,14 @@ def ccd_hfield_kernel(
 
           # transform contact to global frame
           pos_local = 0.5 * (w1 + w2)
-          pos = hf_mat @ pos_local + hf_pos
+          pos = rot_vec_quat(pos_local, hf_quat) + hf_pos
           hfield_contact_pos[count, 0] = pos[0]
           hfield_contact_pos[count, 1] = pos[1]
           hfield_contact_pos[count, 2] = pos[2]
 
           frame_local = make_frame(w1 - w2)
           normal_local = wp.vec3(frame_local[0, 0], frame_local[0, 1], frame_local[0, 2])
-          normal = hf_mat @ normal_local
+          normal = rot_vec_quat(normal_local, hf_quat)
           hfield_contact_normal[count, 0] = normal[0]
           hfield_contact_normal[count, 1] = normal[1]
           hfield_contact_normal[count, 2] = normal[2]
@@ -926,7 +929,7 @@ def ccd_kernel(
     # Data in:
     naconmax_in: int,
     geom_xpos_in: wp.array2d(dtype=wp.vec3),
-    geom_xmat_in: wp.array2d(dtype=wp.mat33),
+    geom_xquat_in: wp.array2d(dtype=wp.quat),
     collision_pair_in: wp.array(dtype=wp.vec2i),
     collision_pairid_in: wp.array(dtype=wp.vec2i),
     collision_worldid_in: wp.array(dtype=int),
@@ -1022,7 +1025,7 @@ def ccd_kernel(
       mesh_polymapnum,
       mesh_polymap,
       geom_xpos_in,
-      geom_xmat_in,
+      geom_xquat_in,
       geoms,
       worldid,
     )
@@ -1233,7 +1236,7 @@ def _pair_count(p1: int, p2: int) -> int:
           m.pair_friction,
           d.naconmax,
           d.geom_xpos,
-          d.geom_xmat,
+          d.geom_xquat,
           d.collision_pair,
           d.collision_pairid,
           d.collision_worldid,
@@ -1319,7 +1322,7 @@ def _pair_count(p1: int, p2: int) -> int:
           m.pair_friction,
           d.naconmax,
           d.geom_xpos,
-          d.geom_xmat,
+          d.geom_xquat,
           d.collision_pair,
           d.collision_pairid,
           d.collision_worldid,

mujoco_warp/_src/collision_driver.py

@@ -20,6 +20,8 @@
 from .collision_convex import convex_narrowphase
 from .collision_primitive import primitive_narrowphase
 from .collision_sdf import sdf_narrowphase
+from .math import quat_to_mat
+from .math import rot_vec_quat
 from .math import upper_tri_index
 from .types import MJ_MAXVAL
 from .types import BroadphaseFilter
@@ -48,15 +50,17 @@ def _zero_nacon_ncollision(
 
 @wp.func
 def _plane_filter(
-  size1: float, size2: float, margin1: float, margin2: float, xpos1: wp.vec3, xpos2: wp.vec3, xmat1: wp.mat33, xmat2: wp.mat33
+  size1: float, size2: float, margin1: float, margin2: float, xpos1: wp.vec3, xpos2: wp.vec3, xquat1: wp.quat, xquat2: wp.quat
 ) -> bool:
   if size1 == 0.0:
     # geom1 is a plane
-    dist = wp.dot(xpos2 - xpos1, wp.vec3(xmat1[0, 2], xmat1[1, 2], xmat1[2, 2]))
+    normal = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), xquat1)
+    dist = wp.dot(xpos2 - xpos1, normal)
     return dist <= size2 + wp.max(margin1, margin2)
   elif size2 == 0.0:
     # geom2 is a plane
-    dist = wp.dot(xpos1 - xpos2, wp.vec3(xmat2[0, 2], xmat2[1, 2], xmat2[2, 2]))
+    normal = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), xquat2)
+    dist = wp.dot(xpos1 - xpos2, normal)
     return dist <= size1 + wp.max(margin1, margin2)
 
   return True
@@ -82,16 +86,16 @@ def _aabb_filter(
   margin2: float,
   xpos1: wp.vec3,
   xpos2: wp.vec3,
-  xmat1: wp.mat33,
-  xmat2: wp.mat33,
+  xquat1: wp.quat,
+  xquat2: wp.quat,
 ) -> bool:
   """Axis aligned boxes collision.
 
   references: see Ericson, Real-time Collision Detection section 4.2.
               filterBox: filter contact based on global AABBs.
   """
-  center1 = xmat1 @ center1 + xpos1
-  center2 = xmat2 @ center2 + xpos2
+  center1 = rot_vec_quat(center1, xquat1) + xpos1
+  center2 = rot_vec_quat(center2, xquat2) + xpos2
 
   margin = wp.max(margin1, margin2)
 
@@ -115,10 +119,10 @@ def _aabb_filter(
     for j in range(2):
       for k in range(2):
         corner1 = wp.vec3(sign[i] * size1[0], sign[j] * size1[1], sign[k] * size1[2])
-        pos1 = xmat1 @ corner1
+        pos1 = rot_vec_quat(corner1, xquat1)
 
         corner2 = wp.vec3(sign[i] * size2[0], sign[j] * size2[1], sign[k] * size2[2])
-        pos2 = xmat2 @ corner2
+        pos2 = rot_vec_quat(corner2, xquat2)
 
         if pos1[0] > max_x1:
           max_x1 = pos1[0]
@@ -184,8 +188,8 @@ def _obb_filter(
   margin2: float,
   xpos1: wp.vec3,
   xpos2: wp.vec3,
-  xmat1: wp.mat33,
-  xmat2: wp.mat33,
+  xquat1: wp.quat,
+  xquat2: wp.quat,
 ) -> bool:
   """Oriented bounding boxes collision (see Gottschalk et al.), see mj_collideOBB."""
   margin = wp.max(margin1, margin2)
@@ -196,8 +200,11 @@ def _obb_filter(
   radius = wp.vec2()
 
   # compute centers in local coordinates
-  xcenter[0] = xmat1 @ center1 + xpos1
-  xcenter[1] = xmat2 @ center2 + xpos2
+  xcenter[0] = rot_vec_quat(center1, xquat1) + xpos1
+  xcenter[1] = rot_vec_quat(center2, xquat2) + xpos2
+
+  xmat1 = quat_to_mat(xquat1)
+  xmat2 = quat_to_mat(xquat2)
 
   # compute normals in global coordinates
   normal[0] = wp.vec3(xmat1[0, 0], xmat1[1, 0], xmat1[2, 0])
@@ -239,7 +246,7 @@ def func(
     geom_margin: wp.array2d(dtype=float),
     # Data in:
     geom_xpos_in: wp.array2d(dtype=wp.vec3),
-    geom_xmat_in: wp.array2d(dtype=wp.mat33),
+    geom_xquat_in: wp.array2d(dtype=wp.quat),
     # In:
     geom1: int,
     geom2: int,
@@ -259,20 +266,20 @@ def func(
     margin_id = worldid % ngeom_margin if wp.static(ngeom_margin > 1) else 0
     margin1, margin2 = geom_margin[margin_id, geom1], geom_margin[margin_id, geom2]
     xpos1, xpos2 = geom_xpos_in[worldid, geom1], geom_xpos_in[worldid, geom2]
-    xmat1, xmat2 = geom_xmat_in[worldid, geom1], geom_xmat_in[worldid, geom2]
+    xquat1, xquat2 = geom_xquat_in[worldid, geom1], geom_xquat_in[worldid, geom2]
 
     if rbound1 == 0.0 or rbound2 == 0.0:
       if wp.static(opt_broadphase_filter & BroadphaseFilter.PLANE):
-        return _plane_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2, xmat1, xmat2)
+        return _plane_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2, xquat1, xquat2)
     else:
       if wp.static(opt_broadphase_filter & BroadphaseFilter.SPHERE):
         if not _sphere_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2):
           return False
       if wp.static(opt_broadphase_filter & BroadphaseFilter.AABB):
-        if not _aabb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2):
+        if not _aabb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xquat1, xquat2):
           return False
       if wp.static(opt_broadphase_filter & BroadphaseFilter.OBB):
-        if not _obb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2):
+        if not _obb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xquat1, xquat2):
           return False
 
     return True
@@ -415,7 +422,7 @@ def kernel(
     nworld_in: int,
     naconmax_in: int,
     geom_xpos_in: wp.array2d(dtype=wp.vec3),
-    geom_xmat_in: wp.array2d(dtype=wp.mat33),
+    geom_xquat_in: wp.array2d(dtype=wp.quat),
     # In:
     sort_index_in: wp.array2d(dtype=int),
     cumulative_sum_in: wp.array(dtype=int),
@@ -460,7 +467,7 @@ def kernel(
 
       if (
         wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin))(
-          geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid
+          geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xquat_in, geom1, geom2, worldid
         )
         or pairid[1] >= 0
       ):
@@ -591,7 +598,7 @@ def sap_broadphase(m: Model, d: Data):
       d.nworld,
       d.naconmax,
       d.geom_xpos,
-      d.geom_xmat,
+      d.geom_xquat,
       sort_index.reshape((-1, m.ngeom)),
       cumulative_sum.reshape(-1),
       nsweep,
@@ -614,7 +621,7 @@ def kernel(
     # Data in:
     naconmax_in: int,
     geom_xpos_in: wp.array2d(dtype=wp.vec3),
-    geom_xmat_in: wp.array2d(dtype=wp.mat33),
+    geom_xquat_in: wp.array2d(dtype=wp.quat),
     # Data out:
     collision_pair_out: wp.array(dtype=wp.vec2i),
     collision_pairid_out: wp.array(dtype=wp.vec2i),
@@ -629,7 +636,7 @@ def kernel(
 
     if (
       wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin))(
-        geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid
+        geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xquat_in, geom1, geom2, worldid
       )
       or nxn_pairid[elementid][1] >= 0
     ):
@@ -676,7 +683,7 @@ def nxn_broadphase(m: Model, d: Data):
       m.nxn_pairid_filtered,
       d.naconmax,
       d.geom_xpos,
-      d.geom_xmat,
+      d.geom_xquat,
     ],
     outputs=[
       d.collision_pair,

mujoco_warp/_src/collision_driver_test.py

@@ -629,7 +629,7 @@ def test_plane_meshtet(self):
     # tetrahedron, separated in z by 0.1
     convex = Geom()
     convex.pos = wp.vec3(0.0)
-    convex.rot = wp.mat33(np.eye(3))
+    convex.rot = wp.quat(1.0, 0.0, 0.0, 0.0)
     convex.graphadr = -1
     convex.vertnum = 4
     convex.vertadr = 0