diff --git a/mujoco_warp/_src/broadphase_test.py b/mujoco_warp/_src/broadphase_test.py
index 1e1e35c63..30939c29d 100644
--- a/mujoco_warp/_src/broadphase_test.py
+++ b/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]
diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py
index 81deb82b8..e6c14ad4d 100644
--- a/mujoco_warp/_src/collision_convex.py
+++ b/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,
diff --git a/mujoco_warp/_src/collision_driver.py b/mujoco_warp/_src/collision_driver.py
index a9cd39c97..4101147fa 100644
--- a/mujoco_warp/_src/collision_driver.py
+++ b/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,
diff --git a/mujoco_warp/_src/collision_driver_test.py b/mujoco_warp/_src/collision_driver_test.py
index 585ca0e5d..ed9251f4c 100644
--- a/mujoco_warp/_src/collision_driver_test.py
+++ b/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
diff --git a/mujoco_warp/_src/collision_gjk.py b/mujoco_warp/_src/collision_gjk.py
index 491aa786b..6d4f8e7e9 100644
--- a/mujoco_warp/_src/collision_gjk.py
+++ b/mujoco_warp/_src/collision_gjk.py
@@ -18,6 +18,8 @@
 import warp as wp
 
 from .collision_primitive import Geom
+from .math import quat_inv
+from .math import rot_vec_quat
 from .types import GeomType
 from .types import mat43
 from .types import mat63
@@ -97,11 +99,11 @@ def support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint:
     sp.point = geom.pos + (0.5 * geom.margin) * geom.size[0] * dir
     return sp
 
-  local_dir = wp.transpose(geom.rot) @ dir
+  local_dir = rot_vec_quat(dir, quat_inv(geom.rot))
   if geomtype == GeomType.BOX:
     tmp = wp.sign(local_dir)
     res = wp.cw_mul(tmp, geom.size)
-    sp.point = geom.rot @ res + geom.pos
+    sp.point = rot_vec_quat(res, geom.rot) + geom.pos
     sp.vertex_index = wp.where(tmp[0] > 0, 1, 0)
     sp.vertex_index += wp.where(tmp[1] > 0, 2, 0)
     sp.vertex_index += wp.where(tmp[2] > 0, 4, 0)
@@ -109,13 +111,13 @@ def support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint:
     res = local_dir * geom.size[0]
     # add cylinder contribution
     res[2] += wp.sign(local_dir[2]) * geom.size[1]
-    sp.point = geom.rot @ res + geom.pos
+    sp.point = rot_vec_quat(res, geom.rot) + geom.pos
   elif geomtype == GeomType.ELLIPSOID:
     res = wp.cw_mul(local_dir, geom.size)
     res = wp.normalize(res)
     # transform to ellipsoid
     res = wp.cw_mul(res, geom.size)
-    sp.point = geom.rot @ res + geom.pos
+    sp.point = rot_vec_quat(res, geom.rot) + geom.pos
   elif geomtype == GeomType.CYLINDER:
     res = wp.vec3(0.0, 0.0, 0.0)
     # set result in XY plane: support on circle
@@ -126,7 +128,7 @@ def support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint:
       res[1] = local_dir[1] * scl
     # set result in Z direction
     res[2] = wp.sign(local_dir[2]) * geom.size[1]
-    sp.point = geom.rot @ res + geom.pos
+    sp.point = rot_vec_quat(res, geom.rot) + geom.pos
   elif geomtype == GeomType.MESH:
     max_dist = float(FLOAT_MIN)
     if geom.graphadr == -1 or geom.vertnum < 10:
@@ -168,7 +170,7 @@ def support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint:
       sp.vertex_index = geom.graph[vert_globalid + imax]
       sp.point = geom.vert[geom.vertadr + sp.vertex_index]
 
-    sp.point = geom.rot @ sp.point + geom.pos
+    sp.point = rot_vec_quat(sp.point, geom.rot) + geom.pos
   elif geomtype == GeomType.HFIELD:
     max_dist = float(FLOAT_MIN)
     # TODO(kbayes): Support edge prisms
@@ -797,25 +799,14 @@ def _replace_simplex3(pt: Polytope, v1: int, v2: int, v3: int) -> GJKResult:
 
 
 @wp.func
-def _rotmat(axis: wp.vec3) -> wp.mat33:
-  n = wp.norm_l2(axis)
-  u1 = axis[0] / n
-  u2 = axis[1] / n
-  u3 = axis[2] / n
-
-  sin = 0.86602540378  # sin(120 deg)
-  cos = -0.5  # cos(120 deg)
-  R = wp.mat33()
-  R[0, 0] = cos + u1 * u1 * (1.0 - cos)
-  R[0, 1] = u1 * u2 * (1.0 - cos) - u3 * sin
-  R[0, 2] = u1 * u3 * (1.0 - cos) + u2 * sin
-  R[1, 0] = u2 * u1 * (1.0 - cos) + u3 * sin
-  R[1, 1] = cos + u2 * u2 * (1.0 - cos)
-  R[1, 2] = u2 * u3 * (1.0 - cos) - u1 * sin
-  R[2, 0] = u1 * u3 * (1.0 - cos) - u2 * sin
-  R[2, 1] = u2 * u3 * (1.0 - cos) + u1 * sin
-  R[2, 2] = cos + u3 * u3 * (1.0 - cos)
-  return R
+def _rotquat(axis: wp.vec3) -> wp.quat:
+  """Returns quaternion for 120 degree rotation around axis."""
+  u = wp.normalize(axis)
+  # θ = 120°, so θ/2 = 60°
+  # cos(60°) = 0.5, sin(60°) = √3/2
+  half_sin = 0.86602540378  # sin(60°)
+  half_cos = 0.5  # cos(60°)
+  return wp.quat(half_cos, half_sin * u[0], half_sin * u[1], half_sin * u[2])
 
 
 @wp.func
@@ -961,9 +952,9 @@ def _polytope2(
   d1 = wp.cross(e, diff)
 
   # rotate around the line segment to get three more points spaced 120 degrees apart
-  R = _rotmat(diff)
-  d2 = R @ d1
-  d3 = R @ d2
+  rot_qaut = _rotquat(diff)
+  d2 = rot_vec_quat(d1, rot_qaut)
+  d3 = rot_vec_quat(d2, rot_qaut)
 
   # save vertices and get indices for each one
   pt.vert[0] = simplex[0]
@@ -1492,7 +1483,7 @@ def _mesh_normals(
   # In:
   feature_dim: int,
   feature_index: wp.vec3i,
-  mat: wp.mat33,
+  rot: wp.quat,
   vertadr: int,
   polyadr: int,
   polynormal: wp.array(dtype=wp.vec3),
@@ -1525,7 +1516,7 @@ def _mesh_normals(
       return 0
 
     # three vertices on mesh define a unique face
-    normals_out[0] = mat @ polynormal[polyadr + faceset[0]]
+    normals_out[0] = rot_vec_quat(polynormal[polyadr + faceset[0]], rot)
     indices_out[0] = faceset[0]
     return 1
 
@@ -1541,7 +1532,7 @@ def _mesh_normals(
     if n == 0:
       return 0
     for i in range(n):
-      normals_out[i] = mat @ polynormal[polyadr + edgeset[i]]
+      normals_out[i] = rot_vec_quat(polynormal[polyadr + edgeset[i]], rot)
       indices_out[i] = edgeset[i]
     return n
 
@@ -1550,7 +1541,7 @@ def _mesh_normals(
     v1_num = polymapnum[vertadr + v1]
     for i in range(v1_num):
       index = polymap[v1_adr + i]
-      normals_out[i] = mat @ polynormal[polyadr + index]
+      normals_out[i] = rot_vec_quat(polynormal[polyadr + index], rot)
       indices_out[i] = index
     return v1_num
   return 0
@@ -1561,7 +1552,7 @@ def _mesh_normals(
 def _mesh_edge_normals(
   # In:
   dim: int,
-  mat: wp.mat33,
+  quat: wp.quat,
   pos: wp.vec3,
   vertadr: int,
   polyadr: int,
@@ -1598,7 +1589,7 @@ def _mesh_edge_normals(
       for j in range(nvert):
         if polyvert[adr + j] == v1i:
           k = wp.where(j == 0, nvert - 1, j - 1)
-          endverts_out[i] = mat @ vert[vertadr + polyvert[adr + k]] + pos
+          endverts_out[i] = rot_vec_quat(vert[vertadr + polyvert[adr + k]], quat) + pos
           normals_out[i] = wp.normalize(endverts_out[i] - v1)
     return v1_num
   return 0
@@ -1608,7 +1599,7 @@ def _mesh_edge_normals(
 @wp.func
 def _box_normals2(
   # In:
-  mat: wp.mat33,
+  quat: wp.quat,
   n: wp.vec3,
   # Out:
   normal_out: wp.array(dtype=wp.vec3),
@@ -1618,19 +1609,12 @@ def _box_normals2(
   face_normals = mat63(1.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, -1.0)
 
   # get local coordinates of the normal
-
-  local_n = wp.normalize(
-    wp.vec3(
-      mat[0][0] * n[0] + mat[1][0] * n[1] + mat[2][0] * n[2],
-      mat[0][1] * n[0] + mat[1][1] * n[1] + mat[2][1] * n[2],
-      mat[0][2] * n[0] + mat[1][2] * n[1] + mat[2][2] * n[2],
-    )
-  )
+  local_n = wp.normalize(rot_vec_quat(n, quat_inv(quat)))
 
   # determine if there is a side close to the normal
   for i in range(6):
     if wp.dot(local_n, face_normals[i]) > FACE_TOL:
-      normal_out[0] = mat @ face_normals[i]
+      normal_out[0] = rot_vec_quat(face_normals[i], quat)
       index_out[0] = i
       return 1
 
@@ -1643,7 +1627,7 @@ def _box_normals(
   # In:
   feature_dim: int,
   feature_index: wp.vec3i,
-  mat: wp.mat33,
+  quat: wp.quat,
   dir: wp.vec3,
   # Out:
   normal_out: wp.array(dtype=wp.vec3),
@@ -1658,7 +1642,7 @@ def _box_normals(
     x = float((v1 & 1) and (v2 & 1) and (v3 & 1)) - float(not (v1 & 1) and not (v2 & 1) and not (v3 & 1))
     y = float((v1 & 2) and (v2 & 2) and (v3 & 2)) - float(not (v1 & 2) and not (v2 & 2) and not (v3 & 2))
     z = float((v1 & 4) and (v2 & 4) and (v3 & 4)) - float(not (v1 & 4) and not (v2 & 4) and not (v3 & 4))
-    normal_out[0] = mat @ wp.vec3(x, y, z)
+    normal_out[0] = rot_vec_quat(wp.vec3(x, y, z), quat)
     sgn = x + y + z
     if x != 0.0:
       index_out[c] = 0
@@ -1673,35 +1657,35 @@ def _box_normals(
       index_out[0] = index_out[0] + 1
     if c == 1:
       return 1
-    return _box_normals2(mat, dir, normal_out, index_out)
+    return _box_normals2(quat, dir, normal_out, index_out)
   if feature_dim == 2:
     c = 0
     x = float((v1 & 1) and (v2 & 1)) - float(not (v1 & 1) and not (v2 & 1))
     y = float((v1 & 2) and (v2 & 2)) - float(not (v1 & 2) and not (v2 & 2))
     z = float((v1 & 4) and (v2 & 4)) - float(not (v1 & 4) and not (v2 & 4))
     if x != 0.0:
-      normal_out[c] = mat @ wp.vec3(float(x), 0.0, 0.0)
+      normal_out[c] = rot_vec_quat(wp.vec3(float(x), 0.0, 0.0), quat)
       index_out[c] = wp.where(x > 0.0, 0, 1)
       c += 1
     if y != 0.0:
-      normal_out[c] = mat @ wp.vec3(0.0, y, 0.0)
+      normal_out[c] = rot_vec_quat(wp.vec3(0.0, y, 0.0), quat)
       index_out[c] = wp.where(y > 0.0, 2, 3)
       c += 1
     if z != 0.0:
-      normal_out[c] = mat @ wp.vec3(0.0, 0.0, z)
+      normal_out[c] = rot_vec_quat(wp.vec3(0.0, 0.0, z), quat)
       index_out[c] = wp.where(z > 0.0, 4, 5)
       c += 1
     if c == 2:
       return 2
-    return _box_normals2(mat, dir, normal_out, index_out)
+    return _box_normals2(quat, dir, normal_out, index_out)
 
   if feature_dim == 1:
     x = wp.where(v1 & 1, 1.0, -1.0)
     y = wp.where(v1 & 2, 1.0, -1.0)
     z = wp.where(v1 & 4, 1.0, -1.0)
-    normal_out[0] = mat @ wp.vec3(x, 0.0, 0.0)
-    normal_out[1] = mat @ wp.vec3(0.0, y, 0.0)
-    normal_out[2] = mat @ wp.vec3(0.0, 0.0, z)
+    normal_out[0] = rot_vec_quat(wp.vec3(x, 0.0, 0.0), quat)
+    normal_out[1] = rot_vec_quat(wp.vec3(0.0, y, 0.0), quat)
+    normal_out[2] = rot_vec_quat(wp.vec3(0.0, 0.0, z), quat)
     index_out[0] = wp.where(x > 0.0, 0, 1)
     index_out[1] = wp.where(y > 0.0, 2, 3)
     index_out[2] = wp.where(z > 0.0, 4, 5)
@@ -1714,7 +1698,7 @@ def _box_normals(
 def _box_edge_normals(
   # In:
   dim: int,
-  mat: wp.mat33,
+  quat: wp.quat,
   pos: wp.vec3,
   size: wp.vec3,
   v1: wp.vec3,
@@ -1735,13 +1719,13 @@ def _box_edge_normals(
     y = wp.where(v1i & 2, size[1], -size[1])
     z = wp.where(v1i & 4, size[2], -size[2])
 
-    endvert_out[0] = mat @ wp.vec3(-x, y, z) + pos
+    endvert_out[0] = rot_vec_quat(wp.vec3(-x, y, z), quat) + pos
     normal_out[0] = wp.normalize(endvert_out[0] - v1)
 
-    endvert_out[1] = mat @ wp.vec3(x, -y, z) + pos
+    endvert_out[1] = rot_vec_quat(wp.vec3(x, -y, z), quat) + pos
     normal_out[1] = wp.normalize(endvert_out[1] - v1)
 
-    endvert_out[2] = mat @ wp.vec3(x, y, -z) + pos
+    endvert_out[2] = rot_vec_quat(wp.vec3(x, y, -z), quat) + pos
     normal_out[2] = wp.normalize(endvert_out[2] - v1)
     return 3
   return 0
@@ -1749,43 +1733,43 @@ def _box_edge_normals(
 
 # recover face of a box from its index
 @wp.func
-def _box_face(mat: wp.mat33, pos: wp.vec3, size: wp.vec3, idx: int, face_out: wp.array(dtype=wp.vec3)) -> int:
+def _box_face(quat: wp.quat, 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.vec3(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.vec3(size[0], -size[1], -size[2]) + pos
-    face_out[3] = mat @ wp.vec3(size[0], -size[1], size[2]) + pos
+    face_out[0] = rot_vec_quat(wp.vec3(size[0], size[1], size[2]), quat) + pos
+    face_out[1] = rot_vec_quat(wp.vec3(size[0], size[1], -size[2]), quat) + pos
+    face_out[2] = rot_vec_quat(wp.vec3(size[0], -size[1], -size[2]), quat) + pos
+    face_out[3] = rot_vec_quat(wp.vec3(size[0], -size[1], size[2]), quat) + pos
     return 4
   if idx == 1:  # left
-    face_out[0] = mat @ wp.vec3(-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.vec3(-size[0], -size[1], size[2]) + pos
-    face_out[3] = mat @ wp.vec3(-size[0], -size[1], -size[2]) + pos
+    face_out[0] = rot_vec_quat(wp.vec3(-size[0], size[1], -size[2]), quat) + pos
+    face_out[1] = rot_vec_quat(wp.vec3(-size[0], size[1], size[2]), quat) + pos
+    face_out[2] = rot_vec_quat(wp.vec3(-size[0], -size[1], size[2]), quat) + pos
+    face_out[3] = rot_vec_quat(wp.vec3(-size[0], -size[1], -size[2]), quat) + pos
     return 4
   if idx == 2:  # top
-    face_out[0] = mat @ wp.vec3(-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.vec3(size[0], size[1], size[2]) + pos
-    face_out[3] = mat @ wp.vec3(-size[0], size[1], size[2]) + pos
+    face_out[0] = rot_vec_quat(wp.vec3(-size[0], size[1], -size[2]), quat) + pos
+    face_out[1] = rot_vec_quat(wp.vec3(size[0], size[1], -size[2]), quat) + pos
+    face_out[2] = rot_vec_quat(wp.vec3(size[0], size[1], size[2]), quat) + pos
+    face_out[3] = rot_vec_quat(wp.vec3(-size[0], size[1], size[2]), quat) + pos
     return 4
   if idx == 3:  # bottom
-    face_out[0] = mat @ wp.vec3(-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.vec3(size[0], -size[1], -size[2]) + pos
-    face_out[3] = mat @ wp.vec3(-size[0], -size[1], -size[2]) + pos
+    face_out[0] = rot_vec_quat(wp.vec3(-size[0], -size[1], size[2]), quat) + pos
+    face_out[1] = rot_vec_quat(wp.vec3(size[0], -size[1], size[2]), quat) + pos
+    face_out[2] = rot_vec_quat(wp.vec3(size[0], -size[1], -size[2]), quat) + pos
+    face_out[3] = rot_vec_quat(wp.vec3(-size[0], -size[1], -size[2]), quat) + pos
     return 4
   if idx == 4:  # front
-    face_out[0] = mat @ wp.vec3(-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.vec3(size[0], -size[1], size[2]) + pos
-    face_out[3] = mat @ wp.vec3(-size[0], -size[1], size[2]) + pos
+    face_out[0] = rot_vec_quat(wp.vec3(-size[0], size[1], size[2]), quat) + pos
+    face_out[1] = rot_vec_quat(wp.vec3(size[0], size[1], size[2]), quat) + pos
+    face_out[2] = rot_vec_quat(wp.vec3(size[0], -size[1], size[2]), quat) + pos
+    face_out[3] = rot_vec_quat(wp.vec3(-size[0], -size[1], size[2]), quat) + pos
     return 4
   if idx == 5:  # back
-    face_out[0] = mat @ wp.vec3(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.vec3(-size[0], -size[1], -size[2]) + pos
-    face_out[3] = mat @ wp.vec3(size[0], -size[1], -size[2]) + pos
+    face_out[0] = rot_vec_quat(wp.vec3(size[0], size[1], -size[2]), quat) + pos
+    face_out[1] = rot_vec_quat(wp.vec3(-size[0], size[1], -size[2]), quat) + pos
+    face_out[2] = rot_vec_quat(wp.vec3(-size[0], -size[1], -size[2]), quat) + pos
+    face_out[3] = rot_vec_quat(wp.vec3(size[0], -size[1], -size[2]), quat) + pos
     return 4
   return 0
 
@@ -1794,7 +1778,7 @@ def _box_face(mat: wp.mat33, pos: wp.vec3, size: wp.vec3, idx: int, face_out: wp
 @wp.func
 def _mesh_face(
   # In:
-  mat: wp.mat33,
+  quat: wp.quat,
   pos: wp.vec3,
   vertadr: int,
   polyadr: int,
@@ -1811,7 +1795,7 @@ def _mesh_face(
   nvert = polyvertnum[polyadr + idx]
   for i in range(nvert - 1, -1, -1):
     v = vert[vertadr + polyvert[adr + i]]
-    face_out[j] = mat @ v + pos
+    face_out[j] = rot_vec_quat(v, quat) + pos
     j += 1
   return nvert
 
@@ -1825,7 +1809,7 @@ def _plane_normal(v1: wp.vec3, v2: wp.vec3, n: wp.vec3) -> Tuple[float, wp.vec3]
 
 @wp.func
 def _halfspace(a: wp.vec3, n: wp.vec3, p: wp.vec3) -> bool:
-  return wp.dot(p - a, n) > -1e-10
+  return wp.dot(p - a, n) > -1e-6
 
 
 @wp.func
diff --git a/mujoco_warp/_src/collision_gjk_test.py b/mujoco_warp/_src/collision_gjk_test.py
index f57703415..74505bbcb 100644
--- a/mujoco_warp/_src/collision_gjk_test.py
+++ b/mujoco_warp/_src/collision_gjk_test.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
+import mujoco
 import numpy as np
 import warp as wp
 from absl.testing import absltest
@@ -39,8 +40,8 @@ def _geom_dist(
   margin=0.0,
   pos1: wp.vec3 | None = None,
   pos2: wp.vec3 | None = None,
-  mat1: wp.mat33 | None = None,
-  mat2: wp.mat33 | None = None,
+  quat1: wp.quat | None = None,
+  quat2: wp.quat | None = None,
 ):
   # we run multiccd on static scenes so these need to be initialized
   nmaxpolygon = 10 if multiccd else 0
@@ -86,7 +87,7 @@ def _ccd_kernel(
     mesh_polymap: wp.array(dtype=int),
     # 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:
     gid1: int,
     gid2: int,
@@ -124,10 +125,10 @@ def _ccd_kernel(
       geom1.pos = geom_xpos_in[0, gid1]
     else:
       geom1.pos = pos1
-    if wp.static(mat1 == None):
-      geom1.rot = geom_xmat_in[0, gid1]
+    if wp.static(quat1 == None):
+      geom1.rot = geom_xquat_in[0, gid1]
     else:
-      geom1.rot = mat1
+      geom1.rot = quat1
     geom1.size = geom_size[0, gid1]
     geom1.margin = margin
     geom1.graphadr = -1
@@ -155,10 +156,10 @@ def _ccd_kernel(
       geom2.pos = geom_xpos_in[0, gid2]
     else:
       geom2.pos = pos2
-    if wp.static(mat2 == None):
-      geom2.rot = geom_xmat_in[0, gid2]
+    if wp.static(quat2 == None):
+      geom2.rot = geom_xquat_in[0, gid2]
     else:
-      geom2.rot = mat2
+      geom2.rot = quat2
     geom2.size = geom_size[0, gid2]
     geom2.margin = margin
     geom2.graphadr = -1
@@ -261,7 +262,7 @@ def _ccd_kernel(
       m.mesh_polymapnum,
       m.mesh_polymap,
       d.geom_xpos,
-      d.geom_xmat,
+      d.geom_xquat,
       gid1,
       gid2,
       m.opt.ccd_iterations,
@@ -549,19 +550,24 @@ def test_cylinder_box(self):
     )
 
     pos = wp.vec3(0.00015228791744448245, -0.00074981129728257656, 0.29839199781417846680)
-    rot = wp.mat33(
-      0.99996972084045410156,
-      0.00776371126994490623,
-      -0.00043433305108919740,
-      -0.00776385562494397163,
-      0.99996984004974365234,
-      -0.00033095158869400620,
-      0.00043175052269361913,
-      0.00033431366318836808,
-      0.99999988079071044922,
+    rot_mat = np.array(
+      [
+        0.99996972084045410156,
+        0.00776371126994490623,
+        -0.00043433305108919740,
+        -0.00776385562494397163,
+        0.99996984004974365234,
+        -0.00033095158869400620,
+        0.00043175052269361913,
+        0.00033431366318836808,
+        0.99999988079071044922,
+      ]
     )
+    rot_quat = np.zeros(4)
+    mujoco.mju_mat2Quat(rot_quat, rot_mat)
+    rot = wp.quat(rot_quat[0], rot_quat[1], rot_quat[2], rot_quat[3])
 
-    dist, _, _, _ = _geom_dist(m, d, 0, 1, pos2=pos, mat2=rot)
+    dist, _, _, _ = _geom_dist(m, d, 0, 1, pos2=pos, quat2=rot)
     self.assertAlmostEqual(dist, -0.0016624178339902445)
 
   def test_box_box_float(self):
@@ -578,32 +584,42 @@ def test_box_box_float(self):
     )
 
     pos1 = wp.vec3(-0.17624500393867492676, -0.12375499308109283447, 0.12499777972698211670)
-    rot1 = wp.mat33(
-      1.00000000000000000000,
-      -0.00000000184385418045,
-      -0.00000025833372774287,
-      0.00000000184391857339,
-      1.00000000000000000000,
-      0.00000024928382913458,
-      0.00000025833372774287,
-      -0.00000024928382913458,
-      1.0000000000000000000,
+    rot1_mat = np.array(
+      [
+        1.00000000000000000000,
+        -0.00000000184385418045,
+        -0.00000025833372774287,
+        0.00000000184391857339,
+        1.00000000000000000000,
+        0.00000024928382913458,
+        0.00000025833372774287,
+        -0.00000024928382913458,
+        1.0000000000000000000,
+      ]
     )
+    rot1_quat = np.zeros(4)
+    mujoco.mju_mat2Quat(rot1_quat, rot1_mat)
+    rot1 = wp.quat(rot1_quat[0], rot1_quat[1], rot1_quat[2], rot1_quat[3])
 
     pos2 = wp.vec3(-0.17624500393867492676, -0.12375499308109283447, 0.17499557137489318848)
-    rot2 = wp.mat33(
-      1.00000000000000000000,
-      -0.00000000184292525685,
-      0.00000012980596864054,
-      0.00000000184294413064,
-      1.00000000000000000000,
-      -0.00000014602545661546,
-      -0.00000012980596864054,
-      0.00000014602545661546,
-      1.00000000000000000000,
+    rot2_mat = np.array(
+      [
+        1.00000000000000000000,
+        -0.00000000184292525685,
+        0.00000012980596864054,
+        0.00000000184294413064,
+        1.00000000000000000000,
+        -0.00000014602545661546,
+        -0.00000012980596864054,
+        0.00000014602545661546,
+        1.00000000000000000000,
+      ]
     )
+    rot2_quat = np.zeros(4)
+    mujoco.mju_mat2Quat(rot2_quat, rot2_mat)
+    rot2 = wp.quat(rot2_quat[0], rot2_quat[1], rot2_quat[2], rot2_quat[3])
 
-    dist, ncon, _, _ = _geom_dist(m, d, 0, 1, multiccd=False, pos1=pos1, mat1=rot1, pos2=pos2, mat2=rot2)
+    dist, ncon, _, _ = _geom_dist(m, d, 0, 1, multiccd=False, pos1=pos1, quat1=rot1, pos2=pos2, quat2=rot2)
     self.assertEqual(ncon, 1)
     self.assertLess(dist, 0.0001)  # real depth is ~ 2E-6
 
diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py
index 602a3de89..157e1e20c 100644
--- a/mujoco_warp/_src/collision_primitive.py
+++ b/mujoco_warp/_src/collision_primitive.py
@@ -30,6 +30,9 @@
 from .collision_primitive_core import sphere_cylinder
 from .collision_primitive_core import sphere_sphere
 from .math import make_frame
+from .math import quat_inv
+from .math import quat_to_mat
+from .math import rot_vec_quat
 from .math import safe_div
 from .math import upper_trid_index
 from .types import MJ_MINMU
@@ -51,7 +54,7 @@
 @wp.struct
 class Geom:
   pos: wp.vec3
-  rot: wp.mat33
+  rot: wp.quat
   normal: wp.vec3
   size: wp.vec3
   margin: float
@@ -95,7 +98,7 @@ def geom_collision_pair(
   mesh_polymap: wp.array(dtype=int),
   # 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:
   geoms: wp.vec2i,
   worldid: int,
@@ -109,14 +112,16 @@ def geom_collision_pair(
   geom_type2 = geom_type[g2]
 
   geom1.pos = geom_xpos_in[worldid, g1]
-  geom1.rot = geom_xmat_in[worldid, g1]
+  geom1.rot = geom_xquat_in[worldid, g1]
   geom1.size = geom_size[worldid % geom_size.shape[0], g1]
-  geom1.normal = wp.vec3(geom1.rot[0, 2], geom1.rot[1, 2], geom1.rot[2, 2])  # plane
+  if geom_type1 == GeomType.PLANE:
+    geom1.normal = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), geom1.rot)
 
   geom2.pos = geom_xpos_in[worldid, g2]
-  geom2.rot = geom_xmat_in[worldid, g2]
+  geom2.rot = geom_xquat_in[worldid, g2]
   geom2.size = geom_size[worldid % geom_size.shape[0], g2]
-  geom2.normal = wp.vec3(geom2.rot[0, 2], geom2.rot[1, 2], geom2.rot[2, 2])  # plane
+  if geom_type2 == GeomType.PLANE:
+    geom2.normal = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), geom2.rot)
 
   if geom_type1 == GeomType.MESH:
     dataid = geom_dataid[g1]
@@ -184,8 +189,9 @@ def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tup
   contact_count = int(0)
 
   # get points in the convex frame
-  plane_pos_local = wp.transpose(convex.rot) @ (plane_pos - convex.pos)
-  n = wp.transpose(convex.rot) @ plane_normal
+  rot_inv = quat_inv(convex.rot)
+  plane_pos_local = rot_vec_quat(plane_pos - convex.pos, rot_inv)
+  n = rot_vec_quat(plane_normal, rot_inv)
 
   # Store indices in vec4
   indices = wp.vec4i(-1, -1, -1, -1)
@@ -380,7 +386,7 @@ def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tup
     # Check if the index is unique (appears exactly once)
     if count == 1:
       pos = convex.vert[convex.vertadr + idx]
-      pos = convex.pos + convex.rot @ pos
+      pos = convex.pos + rot_vec_quat(pos, convex.rot)
       support = wp.dot(plane_pos_local - convex.vert[convex.vertadr + idx], n)
       dist = -support
       pos = pos - 0.5 * dist * plane_normal
@@ -738,7 +744,7 @@ def sphere_capsule_wrapper(
 ):
   """Calculates one contact between a sphere and a capsule."""
   # capsule axis
-  axis = wp.vec3(cap.rot[0, 2], cap.rot[1, 2], cap.rot[2, 2])
+  axis = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), cap.rot)
 
   dist, pos, normal = sphere_capsule(sphere.pos, sphere.size[0], cap.pos, axis, cap.size[0], cap.size[1])
 
@@ -810,8 +816,8 @@ def capsule_capsule_wrapper(
 ):
   """Calculates contacts between two capsules."""
   # capsule axes
-  cap1_axis = wp.vec3(cap1.rot[0, 2], cap1.rot[1, 2], cap1.rot[2, 2])
-  cap2_axis = wp.vec3(cap2.rot[0, 2], cap2.rot[1, 2], cap2.rot[2, 2])
+  cap1_axis = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), cap1.rot)
+  cap2_axis = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), cap2.rot)
 
   dist, pos, normal = capsule_capsule(
     cap1.pos,
@@ -892,7 +898,7 @@ def plane_capsule_wrapper(
 ):
   """Calculates contacts between a capsule and a plane."""
   # capsule axis
-  capsule_axis = wp.vec3(cap.rot[0, 2], cap.rot[1, 2], cap.rot[2, 2])
+  capsule_axis = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), cap.rot)
 
   dist, pos, frame = plane_capsule(
     plane.normal,
@@ -971,7 +977,7 @@ def plane_ellipsoid_wrapper(
   nacon_out: wp.array(dtype=int),
 ):
   """Calculates contacts between an ellipsoid and a plane."""
-  dist, pos, normal = plane_ellipsoid(plane.normal, plane.pos, ellipsoid.pos, ellipsoid.rot, ellipsoid.size)
+  dist, pos, normal = plane_ellipsoid(plane.normal, plane.pos, ellipsoid.pos, quat_to_mat(ellipsoid.rot), ellipsoid.size)
 
   write_contact(
     naconmax_in,
@@ -1040,7 +1046,7 @@ def plane_box_wrapper(
   nacon_out: wp.array(dtype=int),
 ):
   """Calculates contacts between a box and a plane."""
-  dist, pos, normal = plane_box(plane.normal, plane.pos, box.pos, box.rot, box.size)
+  dist, pos, normal = plane_box(plane.normal, plane.pos, box.pos, quat_to_mat(box.rot), box.size)
   frame = make_frame(normal)
 
   for i in range(8):
@@ -1183,7 +1189,7 @@ def sphere_cylinder_wrapper(
 ):
   """Calculates contacts between a sphere and a cylinder."""
   # cylinder axis
-  cylinder_axis = wp.vec3(cylinder.rot[0, 2], cylinder.rot[1, 2], cylinder.rot[2, 2])
+  cylinder_axis = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), cylinder.rot)
 
   dist, pos, normal = sphere_cylinder(
     sphere.pos,
@@ -1262,7 +1268,7 @@ def plane_cylinder_wrapper(
 ):
   """Calculates contacts between a cylinder and a plane."""
   # cylinder axis
-  cylinder_axis = wp.vec3(cylinder.rot[0, 2], cylinder.rot[1, 2], cylinder.rot[2, 2])
+  cylinder_axis = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), cylinder.rot)
 
   dist, pos, normal = plane_cylinder(
     plane.normal,
@@ -1341,7 +1347,7 @@ def sphere_box_wrapper(
   contact_geomcollisionid_out: wp.array(dtype=int),
   nacon_out: wp.array(dtype=int),
 ):
-  dist, pos, normal = sphere_box(sphere.pos, sphere.size[0], box.pos, box.rot, box.size)
+  dist, pos, normal = sphere_box(sphere.pos, sphere.size[0], box.pos, quat_to_mat(box.rot), box.size)
 
   write_contact(
     naconmax_in,
@@ -1411,7 +1417,7 @@ def capsule_box_wrapper(
 ):
   """Calculates contacts between a capsule and a box."""
   # Extract capsule axis
-  axis = wp.vec3(cap.rot[0, 2], cap.rot[1, 2], cap.rot[2, 2])
+  axis = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), cap.rot)
 
   # Call the core function to get contact geometry
   dist, pos, normal = capsule_box(
@@ -1420,7 +1426,7 @@ def capsule_box_wrapper(
     cap.size[0],  # capsule radius
     cap.size[1],  # capsule half length
     box.pos,
-    box.rot,
+    quat_to_mat(box.rot),
     box.size,
   )
 
@@ -1496,10 +1502,10 @@ def box_box_wrapper(
   # Call the core function to get contact geometry
   dist, pos, normal = box_box(
     box1.pos,
-    box1.rot,
+    quat_to_mat(box1.rot),
     box1.size,
     box2.pos,
-    box2.rot,
+    quat_to_mat(box2.rot),
     box2.size,
     margin,
   )
@@ -1608,7 +1614,7 @@ def primitive_narrowphase(
     pair_friction: wp.array2d(dtype=vec5),
     # 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),
     naconmax_in: int,
     collision_pair_in: wp.array(dtype=wp.vec2i),
     collision_pairid_in: wp.array(dtype=wp.vec2i),
@@ -1679,7 +1685,7 @@ def primitive_narrowphase(
       mesh_polymapnum,
       mesh_polymap,
       geom_xpos_in,
-      geom_xmat_in,
+      geom_xquat_in,
       geoms,
       worldid,
     )
@@ -1789,7 +1795,7 @@ def primitive_narrowphase(m: Model, d: Data):
       m.pair_gap,
       m.pair_friction,
       d.geom_xpos,
-      d.geom_xmat,
+      d.geom_xquat,
       d.naconmax,
       d.collision_pair,
       d.collision_pairid,
diff --git a/mujoco_warp/_src/collision_sdf.py b/mujoco_warp/_src/collision_sdf.py
index 0842421a0..1e6c0728a 100644
--- a/mujoco_warp/_src/collision_sdf.py
+++ b/mujoco_warp/_src/collision_sdf.py
@@ -21,6 +21,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 .ray import ray_mesh
 from .types import Data
 from .types import GeomType
@@ -36,7 +39,7 @@
 
 @wp.struct
 class OptimizationParams:
-  rel_mat: wp.mat33
+  rel_mat: wp.quat
   rel_pos: wp.vec3
   attr1: wp.vec3
   attr2: wp.vec3
@@ -68,7 +71,7 @@ class MeshData:
   data_id: int
   data_id: int
   pos: wp.vec3
-  mat: wp.mat33
+  quat: wp.quat
   pnt: wp.vec3
   vec: wp.vec3
   valid: bool = False
@@ -108,7 +111,7 @@ def get_sdf_params(
 
 
 @wp.func
-def transform_aabb(aabb_pos: wp.vec3, aabb_size: wp.vec3, pos: wp.vec3, ori: wp.mat33) -> AABB:
+def transform_aabb(aabb_pos: wp.vec3, aabb_size: wp.vec3, pos: wp.vec3, ori: wp.quat) -> AABB:
   aabb = AABB()
   aabb.max = wp.vec3(-1000000000.0, -1000000000.0, -1000000000.0)
   aabb.min = wp.vec3(1000000000.0, 1000000000.0, 1000000000.0)
@@ -118,7 +121,7 @@ def transform_aabb(aabb_pos: wp.vec3, aabb_size: wp.vec3, pos: wp.vec3, ori: wp.
       aabb_size.y * (1.0 if (i & 2) else -1.0),
       aabb_size.z * (1.0 if (i & 4) else -1.0),
     )
-    frame_vec = ori * (vec + aabb_pos) + pos
+    frame_vec = rot_vec_quat(vec + aabb_pos, ori) + pos
     aabb.min = wp.min(aabb.min, frame_vec)
     aabb.max = wp.max(aabb.max, frame_vec)
   return aabb
@@ -375,7 +378,7 @@ def sdf(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: Volume
       mesh_data.mesh_face,
       mesh_data.data_id,
       mesh_data.pos,
-      mesh_data.mat,
+      mesh_data.quat,
       mesh_data.pnt,
       mesh_data.vec,
     )
@@ -388,7 +391,7 @@ def sdf(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: Volume
         mesh_data.mesh_face,
         mesh_data.data_id,
         mesh_data.pos,
-        mesh_data.mat,
+        mesh_data.quat,
         mesh_data.pnt,
         -mesh_data.vec,
       )
@@ -425,7 +428,7 @@ def sdf_grad(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: V
       mesh_data.mesh_face,
       mesh_data.data_id,
       mesh_data.pos,
-      mesh_data.mat,
+      mesh_data.quat,
       mesh_data.pnt,
       mesh_data.vec,
     )
@@ -486,7 +489,7 @@ def compute_grad(
   B = sdf(GeomType.SDF, p2, params.attr2, sdf_type2, volume_data2, mesh_data2)
   grad1 = sdf_grad(type1, p1, params.attr1, sdf_type1, volume_data1, mesh_data1)
   grad2 = sdf_grad(GeomType.SDF, p2, params.attr2, sdf_type2, volume_data2, mesh_data2)
-  grad1_transformed = wp.transpose(params.rel_mat) * grad1
+  grad1_transformed = rot_vec_quat(grad1, quat_inv(params.rel_mat))
   if sfd_intersection:
     if A > B:
       return grad1_transformed
@@ -526,7 +529,7 @@ def gradient_step(
   for i in range(niter):
     alpha = float(2.0)
     x2 = wp.vec3(x[0], x[1], x[2])
-    x1 = params.rel_mat * x2 + params.rel_pos
+    x1 = rot_vec_quat(x2, params.rel_mat) + params.rel_pos
     grad = compute_grad(
       type1, x1, x2, params, sdf_type1, sdf_type2, sfd_intersection, volume_data1, volume_data2, mesh_data1, mesh_data2
     )
@@ -552,7 +555,7 @@ def gradient_step(
       alpha *= rho
       wolfe *= rho
       x = x2 - grad * alpha
-      x1 = params.rel_mat * x + params.rel_pos
+      x1 = rot_vec_quat(x, params.rel_mat) + params.rel_pos
       dist = clearance(
         type1,
         x1,
@@ -582,9 +585,9 @@ def gradient_descent(
   attr1: wp.vec3,
   attr2: wp.vec3,
   pos1: wp.vec3,
-  rot1: wp.mat33,
+  rot1: wp.quat,
   pos2: wp.vec3,
-  rot2: wp.mat33,
+  rot2: wp.quat,
   sdf_type1: int,
   sdf_type2: int,
   sdf_iterations: int,
@@ -594,24 +597,24 @@ def gradient_descent(
   mesh_data2: MeshData,
 ) -> Tuple[float, wp.vec3, wp.vec3]:
   params = OptimizationParams()
-  params.rel_mat = wp.transpose(rot1) * rot2
-  params.rel_pos = wp.transpose(rot1) * (pos2 - pos1)
+  params.rel_mat = mul_quat(quat_inv(rot1), rot2)
+  params.rel_pos = rot_vec_quat(pos2 - pos1, quat_inv(rot1))
   params.attr1 = attr1
   params.attr2 = attr2
   dist, x = gradient_step(
     type1, x0_initial, params, sdf_type1, sdf_type2, sdf_iterations, False, volume_data1, volume_data2, mesh_data1, mesh_data2
   )
   dist, x = gradient_step(type1, x, params, sdf_type1, sdf_type2, 1, True, volume_data1, volume_data2, mesh_data1, mesh_data2)
-  x_1 = params.rel_mat * x + params.rel_pos
+  x_1 = rot_vec_quat(x, params.rel_mat) + params.rel_pos
   grad1 = sdf_grad(type1, x_1, params.attr1, sdf_type1, volume_data1, mesh_data1)
-  grad1 = wp.transpose(params.rel_mat) * grad1
+  grad1 = rot_vec_quat(grad1, quat_inv(params.rel_mat))
   grad1 = wp.normalize(grad1)
   grad2 = sdf_grad(GeomType.SDF, x, params.attr2, sdf_type2, volume_data2, mesh_data2)
   grad2 = wp.normalize(grad2)
   n = grad1 - grad2
   n = wp.normalize(n)
-  pos = rot2 * x + pos2
-  n = rot2 * n
+  pos = rot_vec_quat(x, rot2) + pos2
+  n = rot_vec_quat(n, rot2)
   pos3 = pos - n * dist / 2.0
   return dist, pos3, n
 
@@ -663,12 +666,12 @@ def _sdf_narrowphase(
   geom_plugin_index: wp.array(dtype=int),
   # Data in:
   geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
   naconmax_in: int,
   collision_pair_in: wp.array(dtype=wp.vec2i),
   collision_pairid_in: wp.array(dtype=wp.vec2i),
   collision_worldid_in: wp.array(dtype=int),
   ncollision_in: wp.array(dtype=int),
+  geom_xquat_in: wp.array2d(dtype=wp.quat),
   # In:
   sdf_initpoints: int,
   sdf_iterations: int,
@@ -740,7 +743,7 @@ def _sdf_narrowphase(
     mesh_polymapnum,
     mesh_polymap,
     geom_xpos_in,
-    geom_xmat_in,
+    geom_xquat_in,
     geoms,
     worldid,
   )
@@ -752,15 +755,13 @@ def _sdf_narrowphase(
   g1_plugin = geom_plugin_index[g1]
   g2_plugin = geom_plugin_index[g2]
 
-  g1_to_g2_rot = wp.transpose(geom1.rot) * geom2.rot
-  g1_to_g2_pos = wp.transpose(geom1.rot) * (geom2.pos - geom1.pos)
+  g1_to_g2_pos = rot_vec_quat(geom2.pos - geom1.pos, quat_inv(geom1.rot))
   aabb_pos = geom_aabb[aabb_id, g1, 0]
   aabb_size = geom_aabb[aabb_id, g1, 1]
-  identity = wp.identity(3, dtype=float)
-  aabb1 = transform_aabb(aabb_pos, aabb_size, wp.vec3(0.0), identity)
+  aabb1 = transform_aabb(aabb_pos, aabb_size, wp.vec3(0.0), wp.quat(1.0, 0.0, 0.0, 0.0))
   aabb_pos = geom_aabb[aabb_id, g2, 0]
   aabb_size = geom_aabb[aabb_id, g2, 1]
-  aabb2 = transform_aabb(aabb_pos, aabb_size, g1_to_g2_pos, g1_to_g2_rot)
+  aabb2 = transform_aabb(aabb_pos, aabb_size, g1_to_g2_pos, mul_quat(quat_inv(geom1.rot), geom2.rot))
   aabb_intersection = AABB()
   aabb_intersection.min = wp.max(aabb1.min, aabb2.min)
   aabb_intersection.max = wp.min(aabb1.max, aabb2.max)
@@ -785,7 +786,7 @@ def _sdf_narrowphase(
   mesh_data1.mesh_face = mesh_face
   mesh_data1.data_id = geom_dataid[g1]
   mesh_data1.pos = geom1.pos
-  mesh_data1.mat = geom1.rot
+  mesh_data1.quat = geom1.rot
   mesh_data1.pnt = wp.vec3(-1.0)
   mesh_data1.vec = wp.vec3(0.0)
   mesh_data1.valid = True
@@ -797,7 +798,7 @@ def _sdf_narrowphase(
   mesh_data2.mesh_face = mesh_face
   mesh_data2.data_id = geom_dataid[g2]
   mesh_data2.pos = geom2.pos
-  mesh_data2.mat = geom2.rot
+  mesh_data2.quat = geom2.rot
   mesh_data2.pnt = wp.vec3(-1.0)
   mesh_data2.vec = wp.vec3(0.0)
   mesh_data2.valid = True
@@ -807,8 +808,8 @@ def _sdf_narrowphase(
     aabb_intersection.min[1] + (aabb_intersection.max[1] - aabb_intersection.min[1]) * halton(i, 3),
     aabb_intersection.min[2] + (aabb_intersection.max[2] - aabb_intersection.min[2]) * halton(i, 5),
   )
-  x = geom1.rot * x_g2 + geom1.pos
-  x0_initial = wp.transpose(rot2) * (x - pos2)
+  x = rot_vec_quat(x_g2, geom1.rot) + geom1.pos
+  x0_initial = rot_vec_quat(x - pos2, quat_inv(rot2))
   dist, pos, n = gradient_descent(
     type1,
     x0_initial,
@@ -908,12 +909,12 @@ def sdf_narrowphase(m: Model, d: Data):
       m.plugin_attr,
       m.geom_plugin_index,
       d.geom_xpos,
-      d.geom_xmat,
       d.naconmax,
       d.collision_pair,
       d.collision_pairid,
       d.collision_worldid,
       d.ncollision,
+      d.geom_xquat,
       m.opt.sdf_initpoints,
       m.opt.sdf_iterations,
     ],
diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py
index 9f04d2655..dacf55135 100644
--- a/mujoco_warp/_src/constraint.py
+++ b/mujoco_warp/_src/constraint.py
@@ -147,7 +147,7 @@ def _efc_equality_connect(
   qvel_in: wp.array2d(dtype=float),
   eq_active_in: wp.array2d(dtype=bool),
   xpos_in: wp.array2d(dtype=wp.vec3),
-  xmat_in: wp.array2d(dtype=wp.mat33),
+  xquat_in: wp.array2d(dtype=wp.quat),
   site_xpos_in: wp.array2d(dtype=wp.vec3),
   subtree_com_in: wp.array2d(dtype=wp.vec3),
   cdof_in: wp.array2d(dtype=wp.spatial_vector),
@@ -196,8 +196,8 @@ def _efc_equality_connect(
   else:
     body1id = obj1id
     body2id = obj2id
-    pos1 = xpos_in[worldid, body1id] + xmat_in[worldid, body1id] @ anchor1
-    pos2 = xpos_in[worldid, body2id] + xmat_in[worldid, body2id] @ anchor2
+    pos1 = xpos_in[worldid, body1id] + math.rot_vec_quat(anchor1, xquat_in[worldid, body1id])
+    pos2 = xpos_in[worldid, body2id] + math.rot_vec_quat(anchor2, xquat_in[worldid, body2id])
 
   # error is difference in global positions
   pos = pos1 - pos2
@@ -740,7 +740,6 @@ def _efc_equality_weld(
   eq_active_in: wp.array2d(dtype=bool),
   xpos_in: wp.array2d(dtype=wp.vec3),
   xquat_in: wp.array2d(dtype=wp.quat),
-  xmat_in: wp.array2d(dtype=wp.mat33),
   site_xpos_in: wp.array2d(dtype=wp.vec3),
   subtree_com_in: wp.array2d(dtype=wp.vec3),
   cdof_in: wp.array2d(dtype=wp.spatial_vector),
@@ -797,8 +796,8 @@ def _efc_equality_weld(
   else:
     body1id = obj1id
     body2id = obj2id
-    pos1 = xpos_in[worldid, body1id] + xmat_in[worldid, body1id] @ anchor2
-    pos2 = xpos_in[worldid, body2id] + xmat_in[worldid, body2id] @ anchor1
+    pos1 = xpos_in[worldid, body1id] + math.rot_vec_quat(anchor2, xquat_in[worldid, body1id])
+    pos2 = xpos_in[worldid, body2id] + math.rot_vec_quat(anchor1, xquat_in[worldid, body2id])
 
     quat = math.mul_quat(xquat_in[worldid, body1id], relpose)
     quat1 = math.quat_inv(xquat_in[worldid, body2id])
@@ -1600,7 +1599,7 @@ def make_constraint(m: types.Model, d: types.Data):
           d.qvel,
           d.eq_active,
           d.xpos,
-          d.xmat,
+          d.xquat,
           d.site_xpos,
           d.subtree_com,
           d.cdof,
@@ -1645,7 +1644,6 @@ def make_constraint(m: types.Model, d: types.Data):
           d.eq_active,
           d.xpos,
           d.xquat,
-          d.xmat,
           d.site_xpos,
           d.subtree_com,
           d.cdof,
diff --git a/mujoco_warp/_src/forward_test.py b/mujoco_warp/_src/forward_test.py
index 1880e2433..e47e3a182 100644
--- a/mujoco_warp/_src/forward_test.py
+++ b/mujoco_warp/_src/forward_test.py
@@ -315,14 +315,10 @@ def test_step1(self, xml, energy):
     step1_field = [
       "xpos",
       "xquat",
-      "xmat",
       "xipos",
-      "ximat",
       "xanchor",
       "xaxis",
       "geom_xpos",
-      "geom_xmat",
-      "site_xmat",
       "subtree_com",
       "cinert",
       "cdof",
@@ -373,7 +369,7 @@ def _getattr(arr):
       return getattr(d, arr), False
 
     for arr in step1_field:
-      if arr in ("geom_xpos", "geom_xmat"):
+      if arr in ("geom_xpos", "geom_xmat", "geom_xquat"):
         # leave geom_xpos and geom_xmat untouched because they have static data
         continue
       attr, _ = _getattr(arr)
@@ -394,7 +390,7 @@ def _getattr(arr):
       d_arr, is_nefc = _getattr(arr)
       d_arr = d_arr.numpy()[0]
       mjd_arr = getattr(mjd, arr)
-      if arr in ["xmat", "ximat", "geom_xmat", "site_xmat", "cam_xmat"]:
+      if arr in ["cam_xmat"]:
         mjd_arr = mjd_arr.reshape(-1)
         d_arr = d_arr.reshape(-1)
       elif arr == "qM":
@@ -434,6 +430,27 @@ def _getattr(arr):
 
       _assert_eq(d_arr, mjd_arr, arr)
 
+    # compare xiquat (warp) to ximat (mujoco) by converting quaternion to matrix
+    xiquat = d.xiquat.numpy()[0]
+    ximat = np.zeros((xiquat.shape[0], 9))
+    for i in range(xiquat.shape[0]):
+      mujoco.mju_quat2Mat(ximat[i], xiquat[i])
+    _assert_eq(ximat.reshape(-1), mjd.ximat.reshape(-1), "xiquat->ximat")
+
+    # compare geom_xquat (warp) to geom_xmat (mujoco) by converting quaternion to matrix
+    geom_xquat = d.geom_xquat.numpy()[0]
+    geom_xmat = np.zeros((geom_xquat.shape[0], 9))
+    for i in range(geom_xquat.shape[0]):
+      mujoco.mju_quat2Mat(geom_xmat[i], geom_xquat[i])
+    _assert_eq(geom_xmat.reshape(-1), mjd.geom_xmat.reshape(-1), "geom_xquat->geom_xmat")
+
+    # compare site_xquat (warp) to site_xmat (mujoco) by converting quaternion to matrix
+    site_xquat = d.site_xquat.numpy()[0]
+    site_xmat = np.zeros((site_xquat.shape[0], 9))
+    for i in range(site_xquat.shape[0]):
+      mujoco.mju_quat2Mat(site_xmat[i], site_xquat[i])
+    _assert_eq(site_xmat.reshape(-1), mjd.site_xmat.reshape(-1), "site_xquat->site_xmat")
+
     # TODO(team): sensor_pos
     # TODO(team): sensor_vel
 
diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py
index e2ed2d238..508adf0bf 100644
--- a/mujoco_warp/_src/io.py
+++ b/mujoco_warp/_src/io.py
@@ -649,6 +649,14 @@ def make_data(
   mocap_body = np.nonzero(mjm.body_mocapid >= 0)[0]
   mocap_id = mjm.body_mocapid[mocap_body]
 
+  xiquat = np.zeros((mjm.nbody, 4))
+  for i in range(mjm.nbody):
+    mujoco.mju_mat2Quat(xiquat[i], mjd.ximat[i])
+
+  geom_xquat = np.zeros((mjm.ngeom, 4))
+  for i in range(mjm.ngeom):
+    mujoco.mju_mat2Quat(geom_xquat[i], mjd.geom_xmat[i])
+
   d_kwargs = {
     "qpos": wp.array(np.tile(mjm.qpos0, nworld), shape=(nworld, mjm.nq), dtype=float),
     "contact": contact,
@@ -660,11 +668,10 @@ def make_data(
     "qLD": None,
     # world body
     "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),
+    "xiquat": wp.array(np.tile(xiquat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.quat),
     # 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),
+    "geom_xquat": wp.array(np.tile(geom_xquat, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.quat),
     # 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(
@@ -827,6 +834,9 @@ def put_data(
     "ne_ten": None,
     "ne_flex": None,
     "nsolving": None,
+    "xiquat": None,
+    "geom_xquat": None,
+    "site_xquat": None,
   }
   for f in dataclasses.fields(types.Data):
     if f.name in d_kwargs:
@@ -872,6 +882,17 @@ def put_data(
   mujoco.mju_sparse2dense(actuator_moment, mjd.actuator_moment, mjd.moment_rownnz, mjd.moment_rowadr, mjd.moment_colind)
   d.actuator_moment = wp.array(np.full((nworld, mjm.nu, mjm.nv), actuator_moment), dtype=float)
 
+  # convert xmat fields to xquat
+  for xmat, attr, n in [
+    (mjd.ximat, "xiquat", mjm.nbody),
+    (mjd.geom_xmat, "geom_xquat", mjm.ngeom),
+    (mjd.site_xmat, "site_xquat", mjm.nsite),
+  ]:
+    xquat = np.zeros((n, 4))
+    for i in range(n):
+      mujoco.mju_mat2Quat(xquat[i], xmat[i])
+    setattr(d, attr, wp.array(np.full((nworld, n, 4), xquat), dtype=wp.quat))
+
   d.nacon = wp.array([mjd.ncon * nworld], dtype=int)
 
   d.ne_connect = wp.full(nworld, 3 * int(np.sum((mjm.eq_type == mujoco.mjtEq.mjEQ_CONNECT) & mjd.eq_active)), dtype=int)
@@ -942,6 +963,18 @@ def get_data_into(
 
   efc_idx = efc_idx[:nefc]  # dont emit indices for overflow constraints
 
+  # convert quaternions to rotation matrices
+  for n, quat, xmat in [
+    (mjm.nbody, d.xquat, result.xmat),
+    (mjm.nbody, d.xiquat, result.ximat),
+    (mjm.ngeom, d.geom_xquat, result.geom_xmat),
+    (mjm.nsite, d.site_xquat, result.site_xmat),
+  ]:
+    mat = np.zeros((n, 9))
+    for i in range(n):
+      mujoco.mju_quat2Mat(mat[i], quat.numpy()[world_id, i])
+    xmat[:] = mat
+
   result.solver_niter[0] = d.solver_niter.numpy()[world_id]
   result.ncon = ncon
   result.ne = ne
@@ -963,15 +996,11 @@ def get_data_into(
   result.act_dot[:] = d.act_dot.numpy()[world_id]
   result.xpos[:] = d.xpos.numpy()[world_id]
   result.xquat[:] = d.xquat.numpy()[world_id]
-  result.xmat[:] = d.xmat.numpy()[world_id].reshape((-1, 9))
   result.xipos[:] = d.xipos.numpy()[world_id]
-  result.ximat[:] = d.ximat.numpy()[world_id].reshape((-1, 9))
   result.xanchor[:] = d.xanchor.numpy()[world_id]
   result.xaxis[:] = d.xaxis.numpy()[world_id]
   result.geom_xpos[:] = d.geom_xpos.numpy()[world_id]
-  result.geom_xmat[:] = d.geom_xmat.numpy()[world_id].reshape((-1, 9))
   result.site_xpos[:] = d.site_xpos.numpy()[world_id]
-  result.site_xmat[:] = d.site_xmat.numpy()[world_id].reshape((-1, 9))
   result.cam_xpos[:] = d.cam_xpos.numpy()[world_id]
   result.cam_xmat[:] = d.cam_xmat.numpy()[world_id].reshape((-1, 9))
   result.light_xpos[:] = d.light_xpos.numpy()[world_id]
diff --git a/mujoco_warp/_src/io_test.py b/mujoco_warp/_src/io_test.py
index a7b39ad7b..7d6cb1f2c 100644
--- a/mujoco_warp/_src/io_test.py
+++ b/mujoco_warp/_src/io_test.py
@@ -141,15 +141,14 @@ def test_get_data_into(self, nworld, world_id):
       "act_dot",
       "xpos",
       "xquat",
-      "xmat",
       "xipos",
-      "ximat",
+      # "ximat",
       "xanchor",
       "xaxis",
       "geom_xpos",
-      "geom_xmat",
+      # "geom_xmat",
       "site_xpos",
-      "site_xmat",
+      # "site_xmat",
       "cam_xpos",
       "cam_xmat",
       "light_xpos",
diff --git a/mujoco_warp/_src/passive.py b/mujoco_warp/_src/passive.py
index 3b1ddd9a9..dc0583757 100644
--- a/mujoco_warp/_src/passive.py
+++ b/mujoco_warp/_src/passive.py
@@ -280,11 +280,11 @@ def _fluid_force(
   body_fluid_ellipsoid: wp.array(dtype=bool),
   # Data in:
   xipos_in: wp.array2d(dtype=wp.vec3),
-  ximat_in: wp.array2d(dtype=wp.mat33),
   geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
   subtree_com_in: wp.array2d(dtype=wp.vec3),
   cvel_in: wp.array2d(dtype=wp.spatial_vector),
+  xiquat_in: wp.array2d(dtype=wp.quat),
+  geom_xquat_in: wp.array2d(dtype=wp.quat),
   # Out:
   fluid_applied_out: wp.array2d(dtype=wp.spatial_vector),
 ):
@@ -308,8 +308,6 @@ def _fluid_force(
 
   # Body kinematics
   xipos = xipos_in[worldid, bodyid]
-  rot = ximat_in[worldid, bodyid]
-  rotT = wp.transpose(rot)
   cvel = cvel_in[worldid, bodyid]
   ang_global = wp.spatial_top(cvel)
   lin_global = wp.spatial_bottom(cvel)
@@ -331,17 +329,17 @@ def _fluid_force(
 
       size = geom_size[worldid % geom_size.shape[0], geomid]
       semiaxes = _geom_semiaxes(size, geom_type[geomid])
-      geom_rot = geom_xmat_in[worldid, geomid]
-      geom_rotT = wp.transpose(geom_rot)
       geom_pos = geom_xpos_in[worldid, geomid]
+      geom_rot = geom_xquat_in[worldid, geomid]
+      geom_rot_inv = math.quat_inv(geom_rot)
 
       lin_point = lin_com + wp.cross(ang_global, geom_pos - xipos)
 
-      l_ang = geom_rotT @ ang_global
-      l_lin = geom_rotT @ lin_point
+      l_ang = math.rot_vec_quat(ang_global, geom_rot_inv)
+      l_lin = math.rot_vec_quat(lin_point, geom_rot_inv)
 
       if wind[0] or wind[1] or wind[2]:
-        l_lin -= geom_rotT @ wind
+        l_lin -= math.rot_vec_quat(wind, geom_rot_inv)
 
       lfrc_torque = wp.vec3(0.0)
       lfrc_force = wp.vec3(0.0)
@@ -437,17 +435,20 @@ def _fluid_force(
       lfrc_force *= coef
 
       # map force/torque from local to world frame: lfrc -> bfrc
-      torque_global += geom_rot @ lfrc_torque
-      force_global += geom_rot @ lfrc_force
+      torque_global += math.rot_vec_quat(lfrc_torque, geom_rot)
+      force_global += math.rot_vec_quat(lfrc_force, geom_rot)
 
     fluid_applied_out[worldid, bodyid] = wp.spatial_vector(force_global, torque_global)
     return
 
-  l_ang = rotT @ ang_global
-  l_lin = rotT @ lin_com
+  rot = xiquat_in[worldid, bodyid]
+  rot_inv = math.quat_inv(rot)
+
+  l_ang = math.rot_vec_quat(ang_global, rot_inv)
+  l_lin = math.rot_vec_quat(lin_com, rot_inv)
 
   if wind[0] or wind[1] or wind[2]:
-    l_lin -= rotT @ wind
+    l_lin -= math.rot_vec_quat(wind, rot_inv)
 
   lfrc_torque = wp.vec3(0.0)
   lfrc_force = wp.vec3(0.0)
@@ -485,8 +486,8 @@ def _fluid_force(
       box2 * (box0_pow4 + box1_pow4) * wp.abs(l_ang[2]) * l_ang[2] * scl,
     )
 
-  torque_global = rot @ lfrc_torque
-  force_global = rot @ lfrc_force
+  torque_global = math.rot_vec_quat(lfrc_torque, rot)
+  force_global = math.rot_vec_quat(lfrc_force, rot)
 
   fluid_applied_out[worldid, bodyid] = wp.spatial_vector(force_global, torque_global)
 
@@ -511,11 +512,11 @@ def _fluid(m: Model, d: Data):
       m.geom_fluid,
       m.body_fluid_ellipsoid,
       d.xipos,
-      d.ximat,
       d.geom_xpos,
-      d.geom_xmat,
       d.subtree_com,
       d.cvel,
+      d.xiquat,
+      d.geom_xquat,
     ],
     outputs=[fluid_applied],
   )
diff --git a/mujoco_warp/_src/ray.py b/mujoco_warp/_src/ray.py
index 76ca3646d..c497abb61 100644
--- a/mujoco_warp/_src/ray.py
+++ b/mujoco_warp/_src/ray.py
@@ -17,6 +17,8 @@
 
 import warp as wp
 
+from .math import quat_inv
+from .math import rot_vec_quat
 from .math import safe_div
 from .types import MJ_MINVAL
 from .types import Data
@@ -28,22 +30,21 @@
 
 
 @wp.func
-def _ray_map(pos: wp.vec3, mat: wp.mat33, pnt: wp.vec3, vec: wp.vec3) -> Tuple[wp.vec3, wp.vec3]:
+def _ray_map(pos: wp.vec3, quat: wp.quat, pnt: wp.vec3, vec: wp.vec3) -> Tuple[wp.vec3, wp.vec3]:
   """Maps ray to local geom frame coordinates.
 
   Args:
       pos: position of geom frame
-      mat: orientation of geom frame
+      quat: orientation of geom frame
       pnt: starting point of ray in world coordinates
       vec: direction of ray in world coordinates
 
   Returns:
       3D point and 3D direction in local geom frame
   """
-  matT = wp.transpose(mat)
-  lpnt = matT @ (pnt - pos)
-  lvec = matT @ vec
-
+  quat_inv = quat_inv(quat)
+  lpnt = rot_vec_quat(pnt - pos, quat_inv)
+  lvec = rot_vec_quat(vec, quat_inv)
   return lpnt, lvec
 
 
@@ -183,10 +184,10 @@ def _ray_triangle(
 
 
 @wp.func
-def _ray_plane(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
+def _ray_plane(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
   """Returns the distance and normal at which a ray intersects with a plane."""
   # map to local frame
-  lpnt, lvec = _ray_map(pos, mat, pnt, vec)
+  lpnt, lvec = _ray_map(pos, quat, pnt, vec)
 
   # z-vec not pointing towards front face: reject
   if lvec[2] > -MJ_MINVAL:
@@ -201,7 +202,7 @@ def _ray_plane(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp
 
   # accept only within rendered rectangle
   if (size[0] <= 0.0 or wp.abs(p[0]) <= size[0]) and (size[1] <= 0.0 or wp.abs(p[1]) <= size[1]):
-    return x, wp.vec3(mat[0, 2], mat[1, 2], mat[2, 2])
+    return x, rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), quat)
   else:
     return -1.0, wp.vec3()
 
@@ -224,7 +225,7 @@ def _ray_sphere(pos: wp.vec3, dist_sqr: float, pnt: wp.vec3, vec: wp.vec3) -> Tu
 
 
 @wp.func
-def _ray_capsule(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
+def _ray_capsule(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
   """Returns the distance and normal at which a ray intersects with a capsule."""
   # bounding sphere test
   ssz = size[0] + size[1]
@@ -233,7 +234,7 @@ def _ray_capsule(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec:
     return -1.0, wp.vec3()
 
   # map to local frame
-  lpnt, lvec = _ray_map(pos, mat, pnt, vec)
+  lpnt, lvec = _ray_map(pos, quat, pnt, vec)
 
   # init solution
   x = -1.0
@@ -291,16 +292,16 @@ def _ray_capsule(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec:
 
     # normalize, rotate into global frame
     normal = wp.normalize(normal)
-    normal = mat @ normal
+    normal = rot_vec_quat(normal, quat)
 
   return x, normal
 
 
 @wp.func
-def _ray_ellipsoid(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
+def _ray_ellipsoid(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
   """Returns the distance and normal at which a ray intersects with an ellipsoid."""
   # map to local frame
-  lpnt, lvec = _ray_map(pos, mat, pnt, vec)
+  lpnt, lvec = _ray_map(pos, quat, pnt, vec)
 
   # invert size^2
   s = wp.vec3(safe_div(1.0, size[0] * size[0]), safe_div(1.0, size[1] * size[1]), safe_div(1.0, size[2] * size[2]))
@@ -322,13 +323,13 @@ def _ray_ellipsoid(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec
     # gradient of ellipsoid function
     normal = wp.cw_mul(s, l)
     normal = wp.normalize(normal)
-    normal = mat @ normal
+    normal = rot_vec_quat(normal, quat)
 
   return sol, normal
 
 
 @wp.func
-def _ray_cylinder(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
+def _ray_cylinder(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
   """Returns the distance and normal at which a ray intersects with a cylinder."""
   # bounding sphere test
   ssz = size[0] * size[0] + size[1] * size[1]
@@ -337,7 +338,7 @@ def _ray_cylinder(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec:
     return -1.0, wp.vec3()
 
   # map to local frame
-  lpnt, lvec = _ray_map(pos, mat, pnt, vec)
+  lpnt, lvec = _ray_map(pos, quat, pnt, vec)
 
   # init solution
   x = -1.0
@@ -383,7 +384,7 @@ def _ray_cylinder(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec:
     else:
       normal = wp.vec3(0.0, 0.0, float(part))
 
-    normal = mat @ normal
+    normal = rot_vec_quat(normal, quat)
 
   return x, normal
 
@@ -392,8 +393,8 @@ def _ray_cylinder(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec:
 
 
 @wp.func
-def _ray_box(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, vec6, wp.vec3]:
-  """Returns distance, per side information, and normal at which a ray intersects with a box."""
+def _ray_box(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, vec6, wp.vec3]:
+  """Returns the distance, per side information, and normal at which a ray intersects with a box."""
   all = vec6(-1.0, -1.0, -1.0, -1.0, -1.0, -1.0)
 
   # bounding sphere test
@@ -403,7 +404,7 @@ def _ray_box(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.v
     return -1.0, all, wp.vec3()
 
   # map to local frame
-  lpnt, lvec = _ray_map(pos, mat, pnt, vec)
+  lpnt, lvec = _ray_map(pos, quat, pnt, vec)
 
   # init solution
   x = float(-1.0)
@@ -440,7 +441,7 @@ def _ray_box(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.v
   normal = wp.vec3()
   if x >= 0:
     normal[face_axis] = float(face_side)
-    normal = mat @ normal
+    normal = rot_vec_quat(normal, quat)
 
   return x, all, normal
 
@@ -457,7 +458,7 @@ def _ray_hfield(
   hfield_data: wp.array(dtype=float),
   # In:
   pos: wp.vec3,
-  mat: wp.mat33,
+  quat: wp.quat,
   pnt: wp.vec3,
   vec: wp.vec3,
   id: int,
@@ -474,12 +475,12 @@ def _ray_hfield(
   size = hfield_size[hid]
   adr = hfield_adr[hid]
 
-  mat_col = wp.vec3(mat[0, 2], mat[1, 2], mat[2, 2])
+  mat_col = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), quat)
 
   # compute size and pos of base box
   base_scale = size[3] * 0.5
   base_size = wp.vec3(size[0], size[1], base_scale)
-  base_pos = pos - mat_col * base_scale
+  base_pos = pos + mat_col * base_scale
 
   # compute size and pos of top box
   top_scale = size[2] * 0.5
@@ -487,16 +488,16 @@ def _ray_hfield(
   top_pos = pos + mat_col * top_scale
 
   # init: intersection with base box
-  x, _, normal_base = _ray_box(base_pos, mat, base_size, pnt, vec)
+  x, _, normal_base = _ray_box(base_pos, quat, base_size, pnt, vec)
 
   # check top box: done if no intersection
-  top_intersect, all, normal_top = _ray_box(top_pos, mat, top_size, pnt, vec)
+  top_intersect, all, normal_top = _ray_box(top_pos, quat, top_size, pnt, vec)
 
   if top_intersect < 0.0:
     return x, normal_base
 
   # map to local frame
-  lpnt, lvec = _ray_map(pos, mat, pnt, vec)
+  lpnt, lvec = _ray_map(pos, quat, pnt, vec)
 
   # construct basis vectors of normal plane
   b0 = wp.vec3(1.0, 1.0, 1.0)
@@ -540,7 +541,7 @@ def _ray_hfield(
 
   normal_local = wp.vec3()
   if x >= 0:
-    normal_local = wp.transpose(mat) @ normal_base
+    normal_local = rot_vec_quat(normal_top, quat_inv(quat))
 
   # check triangles within bounds
   for r in range(rmin, rmax):
@@ -611,7 +612,7 @@ def _ray_hfield(
         normal_local = wp.vec3(float(i == 1) - float(i == 0), float(i == 3) - float(i == 2), 0.0)
 
   if x >= 0:
-    normal_local = mat @ normal_local
+    normal_local = rot_vec_quat(normal_local, quat)
   return x, normal_local
 
 
@@ -626,12 +627,12 @@ def ray_mesh(
   # In:
   data_id: int,
   pos: wp.vec3,
-  mat: wp.mat33,
+  quat: wp.quat,
   pnt: wp.vec3,
   vec: wp.vec3,
 ) -> Tuple[float, wp.vec3]:
   """Returns the distance and normal for ray mesh intersections."""
-  pnt, vec = _ray_map(pos, mat, pnt, vec)
+  pnt, vec = _ray_map(pos, quat, pnt, vec)
 
   # compute orthogonal basis vectors
   if wp.abs(vec[0]) < wp.abs(vec[1]):
@@ -682,30 +683,30 @@ def ray_mesh(
       x = dist
       normal = normal_tri
 
-  normal = mat @ normal
+  normal = rot_vec_quat(normal, quat)
 
   return x, normal
 
 
 @wp.func
-def ray_geom(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3, geomtype: int) -> Tuple[float, wp.vec3]:
+def ray_geom(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3, geomtype: int) -> Tuple[float, wp.vec3]:
   """Returns distance along ray to intersection with geom and normal at intersection point.
 
   If no intersection is found, returns -1 and zero vector.
   """
   # TODO(team): static loop unrolling to remove unnecessary branching
   if geomtype == GeomType.PLANE:
-    return _ray_plane(pos, mat, size, pnt, vec)
+    return _ray_plane(pos, quat, size, pnt, vec)
   elif geomtype == GeomType.SPHERE:
     return _ray_sphere(pos, size[0] * size[0], pnt, vec)
   elif geomtype == GeomType.CAPSULE:
-    return _ray_capsule(pos, mat, size, pnt, vec)
+    return _ray_capsule(pos, quat, size, pnt, vec)
   elif geomtype == GeomType.ELLIPSOID:
-    return _ray_ellipsoid(pos, mat, size, pnt, vec)
+    return _ray_ellipsoid(pos, quat, size, pnt, vec)
   elif geomtype == GeomType.CYLINDER:
-    return _ray_cylinder(pos, mat, size, pnt, vec)
+    return _ray_cylinder(pos, quat, size, pnt, vec)
   elif geomtype == GeomType.BOX:
-    dist, _, normal = _ray_box(pos, mat, size, pnt, vec)
+    dist, _, normal = _ray_box(pos, quat, size, pnt, vec)
     return dist, normal
   else:
     return -1.0, wp.vec3()
@@ -735,7 +736,7 @@ def _ray_geom_mesh(
   mat_rgba: wp.array2d(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,
   pnt: wp.vec3,
@@ -758,7 +759,7 @@ def _ray_geom_mesh(
     bodyexclude,
   ):
     pos = geom_xpos_in[worldid, geomid]
-    mat = geom_xmat_in[worldid, geomid]
+    quat = geom_xquat_in[worldid, geomid]
     type = geom_type[geomid]
 
     if type == GeomType.MESH:
@@ -770,7 +771,7 @@ def _ray_geom_mesh(
         mesh_face,
         geom_dataid[geomid],
         pos,
-        mat,
+        quat,
         pnt,
         vec,
       )
@@ -784,13 +785,13 @@ def _ray_geom_mesh(
         hfield_adr,
         hfield_data,
         pos,
-        mat,
+        quat,
         pnt,
         vec,
         geomid,
       )
     else:
-      return ray_geom(pos, mat, geom_size[worldid % geom_size.shape[0], geomid], pnt, vec, type)
+      return ray_geom(pos, quat, geom_size[worldid % geom_size.shape[0], geomid], pnt, vec, type)
   else:
     return -1.0, wp.vec3()
 
@@ -820,7 +821,7 @@ def _ray(
   mat_rgba: wp.array2d(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:
   pnt: wp.array2d(dtype=wp.vec3),
   vec: wp.array2d(dtype=wp.vec3),
@@ -864,7 +865,7 @@ def _ray(
         hfield_data,
         mat_rgba,
         geom_xpos_in,
-        geom_xmat_in,
+        geom_xquat_in,
         worldid,
         pnt[worldid, rayid],
         vec[worldid, rayid],
@@ -994,7 +995,7 @@ def rays(
       m.hfield_data,
       m.mat_rgba,
       d.geom_xpos,
-      d.geom_xmat,
+      d.geom_xquat,
       pnt,
       vec,
       geomgroup,
diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py
index 431ad9d51..4237c993a 100644
--- a/mujoco_warp/_src/sensor.py
+++ b/mujoco_warp/_src/sensor.py
@@ -112,13 +112,13 @@ def _magnetometer(
   # Model:
   opt_magnetic: wp.array(dtype=wp.vec3),
   # Data in:
-  site_xmat_in: wp.array2d(dtype=wp.mat33),
+  site_xquat_in: wp.array2d(dtype=wp.quat),
   # In:
   worldid: int,
   objid: int,
 ) -> wp.vec3:
   magnetic = opt_magnetic[worldid % opt_magnetic.shape[0]]
-  return wp.transpose(site_xmat_in[worldid, objid]) @ magnetic
+  return math.rot_vec_quat(magnetic, math.quat_inv(site_xquat_in[worldid, objid]))
 
 
 @wp.func
@@ -195,7 +195,7 @@ def _sensor_rangefinder_init(
   sensor_rangefinder_adr: wp.array(dtype=int),
   # Data in:
   site_xpos_in: wp.array2d(dtype=wp.vec3),
-  site_xmat_in: wp.array2d(dtype=wp.mat33),
+  site_xquat_in: wp.array2d(dtype=wp.quat),
   # Out:
   pnt_out: wp.array2d(dtype=wp.vec3),
   vec_out: wp.array2d(dtype=wp.vec3),
@@ -204,10 +204,10 @@ def _sensor_rangefinder_init(
   sensorid = sensor_rangefinder_adr[rfid]
   objid = sensor_objid[sensorid]
   site_xpos = site_xpos_in[worldid, objid]
-  site_xmat = site_xmat_in[worldid, objid]
+  site_normal = math.rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), site_xquat_in[worldid, objid])
 
   pnt_out[worldid, rfid] = site_xpos
-  vec_out[worldid, rfid] = wp.vec3(site_xmat[0, 2], site_xmat[1, 2], site_xmat[2, 2])
+  vec_out[worldid, rfid] = site_normal
 
 
 @wp.func
@@ -279,15 +279,15 @@ def _limit_pos(
 def _frame_pos(
   # Data in:
   xpos_in: wp.array2d(dtype=wp.vec3),
-  xmat_in: wp.array2d(dtype=wp.mat33),
+  xquat_in: wp.array2d(dtype=wp.quat),
   xipos_in: wp.array2d(dtype=wp.vec3),
-  ximat_in: wp.array2d(dtype=wp.mat33),
   geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
   site_xpos_in: wp.array2d(dtype=wp.vec3),
-  site_xmat_in: wp.array2d(dtype=wp.mat33),
   cam_xpos_in: wp.array2d(dtype=wp.vec3),
   cam_xmat_in: wp.array2d(dtype=wp.mat33),
+  xiquat_in: wp.array2d(dtype=wp.quat),
+  geom_xquat_in: wp.array2d(dtype=wp.quat),
+  site_xquat_in: wp.array2d(dtype=wp.quat),
   # In:
   worldid: int,
   objid: int,
@@ -313,78 +313,75 @@ def _frame_pos(
 
   if reftype == ObjType.BODY:
     xpos_ref = xipos_in[worldid, refid]
-    xmat_ref = ximat_in[worldid, refid]
+    xquat_ref = xiquat_in[worldid, refid]
   elif objtype == ObjType.XBODY:
     xpos_ref = xpos_in[worldid, refid]
-    xmat_ref = xmat_in[worldid, refid]
+    xquat_ref = xquat_in[worldid, refid]
   elif reftype == ObjType.GEOM:
     xpos_ref = geom_xpos_in[worldid, refid]
-    xmat_ref = geom_xmat_in[worldid, refid]
+    xquat_ref = geom_xquat_in[worldid, refid]
   elif reftype == ObjType.SITE:
     xpos_ref = site_xpos_in[worldid, refid]
-    xmat_ref = site_xmat_in[worldid, refid]
+    xquat_ref = site_xquat_in[worldid, refid]
   elif reftype == ObjType.CAMERA:
     xpos_ref = cam_xpos_in[worldid, refid]
     xmat_ref = cam_xmat_in[worldid, refid]
+    return wp.transpose(xmat_ref) @ (xpos - xpos_ref)
 
   else:  # UNKNOWN
     xpos_ref = wp.vec3(0.0)
-    xmat_ref = wp.identity(3, wp.float32)
+    xquat_ref = wp.quat(1.0, 0.0, 0.0, 0.0)
 
-  return wp.transpose(xmat_ref) @ (xpos - xpos_ref)
+  return math.rot_vec_quat(xpos - xpos_ref, math.quat_inv(xquat_ref))
 
 
 @wp.func
 def _frame_axis(
   # Data in:
-  xmat_in: wp.array2d(dtype=wp.mat33),
-  ximat_in: wp.array2d(dtype=wp.mat33),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
-  site_xmat_in: wp.array2d(dtype=wp.mat33),
+  xquat_in: wp.array2d(dtype=wp.quat),
   cam_xmat_in: wp.array2d(dtype=wp.mat33),
+  xiquat_in: wp.array2d(dtype=wp.quat),
+  geom_xquat_in: wp.array2d(dtype=wp.quat),
+  site_xquat_in: wp.array2d(dtype=wp.quat),
   # In:
   worldid: int,
   objid: int,
   objtype: int,
   refid: int,
   reftype: int,
-  frame_axis: int,
+  frame_axis_index: int,
+  frame_axis_vec: wp.vec3,
 ) -> wp.vec3:
   if objtype == ObjType.BODY:
-    xmat = ximat_in[worldid, objid]
-    axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis])
+    axis = math.rot_vec_quat(frame_axis_vec, xiquat_in[worldid, objid])
   elif objtype == ObjType.XBODY:
-    xmat = xmat_in[worldid, objid]
-    axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis])
+    axis = math.rot_vec_quat(frame_axis_vec, xquat_in[worldid, objid])
   elif objtype == ObjType.GEOM:
-    xmat = geom_xmat_in[worldid, objid]
-    axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis])
+    axis = math.rot_vec_quat(frame_axis_vec, geom_xquat_in[worldid, objid])
   elif objtype == ObjType.SITE:
-    xmat = site_xmat_in[worldid, objid]
-    axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis])
+    axis = math.rot_vec_quat(frame_axis_vec, site_xquat_in[worldid, objid])
   elif objtype == ObjType.CAMERA:
     xmat = cam_xmat_in[worldid, objid]
-    axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis])
-  else:  # UNKNOWN
-    axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis])
+    axis = wp.vec3(xmat[0, frame_axis_index], xmat[1, frame_axis_index], xmat[2, frame_axis_index])
 
   if refid == -1:
     return axis
 
   if reftype == ObjType.BODY:
-    xmat_ref = ximat_in[worldid, refid]
+    xquat_ref = xiquat_in[worldid, refid]
   elif reftype == ObjType.XBODY:
-    xmat_ref = xmat_in[worldid, refid]
+    xquat_ref = xquat_in[worldid, refid]
   elif reftype == ObjType.GEOM:
-    xmat_ref = geom_xmat_in[worldid, refid]
+    xquat_ref = geom_xquat_in[worldid, refid]
   elif reftype == ObjType.SITE:
-    xmat_ref = site_xmat_in[worldid, refid]
+    xquat_ref = site_xquat_in[worldid, refid]
   elif reftype == ObjType.CAMERA:
     xmat_ref = cam_xmat_in[worldid, refid]
+    return wp.transpose(xmat_ref) @ axis
   else:  # UNKNOWN
     xmat_ref = wp.identity(3, dtype=wp.float32)
 
-  return wp.transpose(xmat_ref) @ axis
+  return math.rot_vec_quat(axis, math.quat_inv(xquat_ref))
 
 
 @wp.func
@@ -491,18 +488,17 @@ def _sensor_pos(
   qpos_in: wp.array2d(dtype=float),
   xpos_in: wp.array2d(dtype=wp.vec3),
   xquat_in: wp.array2d(dtype=wp.quat),
-  xmat_in: wp.array2d(dtype=wp.mat33),
   xipos_in: wp.array2d(dtype=wp.vec3),
-  ximat_in: wp.array2d(dtype=wp.mat33),
   geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
   site_xpos_in: wp.array2d(dtype=wp.vec3),
-  site_xmat_in: wp.array2d(dtype=wp.mat33),
   cam_xpos_in: wp.array2d(dtype=wp.vec3),
   cam_xmat_in: wp.array2d(dtype=wp.mat33),
   subtree_com_in: wp.array2d(dtype=wp.vec3),
   ten_length_in: wp.array2d(dtype=float),
   actuator_length_in: wp.array2d(dtype=float),
+  xiquat_in: wp.array2d(dtype=wp.quat),
+  geom_xquat_in: wp.array2d(dtype=wp.quat),
+  site_xquat_in: wp.array2d(dtype=wp.quat),
   # In:
   rangefinder_dist_in: wp.array2d(dtype=float),
   sensor_collision_in: wp.array4d(dtype=float),
@@ -516,7 +512,7 @@ def _sensor_pos(
   out = sensordata_out[worldid]
 
   if sensortype == SensorType.MAGNETOMETER:
-    vec3 = _magnetometer(opt_magnetic, site_xmat_in, worldid, objid)
+    vec3 = _magnetometer(opt_magnetic, site_xquat_in, worldid, objid)
     _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
   elif sensortype == SensorType.CAMPROJECTION:
     refid = sensor_refid[sensorid]
@@ -545,15 +541,15 @@ def _sensor_pos(
     reftype = sensor_reftype[sensorid]
     vec3 = _frame_pos(
       xpos_in,
-      xmat_in,
+      xquat_in,
       xipos_in,
-      ximat_in,
       geom_xpos_in,
-      geom_xmat_in,
       site_xpos_in,
-      site_xmat_in,
       cam_xpos_in,
       cam_xmat_in,
+      xiquat_in,
+      geom_xquat_in,
+      site_xquat_in,
       worldid,
       objid,
       objtype,
@@ -566,13 +562,27 @@ def _sensor_pos(
     refid = sensor_refid[sensorid]
     reftype = sensor_reftype[sensorid]
     if sensortype == SensorType.FRAMEXAXIS:
-      axis = 0
+      frame_axis_vec = wp.vec3(1.0, 0.0, 0.0)
+      frame_axis_index = 0
     elif sensortype == SensorType.FRAMEYAXIS:
-      axis = 1
+      frame_axis_vec = wp.vec3(0.0, 1.0, 0.0)
+      frame_axis_index = 1
     elif sensortype == SensorType.FRAMEZAXIS:
-      axis = 2
+      frame_axis_vec = wp.vec3(0.0, 0.0, 1.0)
+      frame_axis_index = 2
     vec3 = _frame_axis(
-      ximat_in, xmat_in, geom_xmat_in, site_xmat_in, cam_xmat_in, worldid, objid, objtype, refid, reftype, axis
+      xquat_in,
+      cam_xmat_in,
+      xiquat_in,
+      geom_xquat_in,
+      site_xquat_in,
+      worldid,
+      objid,
+      objtype,
+      refid,
+      reftype,
+      frame_axis_index,
+      frame_axis_vec,
     )
     _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
   elif sensortype == SensorType.FRAMEQUAT:
@@ -690,7 +700,9 @@ def _sensor_pos(
     else:
       return  # should not occur
     refid = sensor_refid[sensorid]
-    val_bool = inside_geom(site_xpos_in[worldid, refid], site_xmat_in[worldid, refid], site_size[refid], site_type[refid], xpos)
+    val_bool = inside_geom(
+      site_xpos_in[worldid, refid], site_xquat_in[worldid, refid], site_size[refid], site_type[refid], xpos
+    )
     _write_scalar(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, float(val_bool), out)
   elif sensortype == SensorType.E_POTENTIAL:
     val = energy_in[worldid][0]
@@ -772,7 +784,7 @@ def sensor_pos(m: Model, d: Data):
     wp.launch(
       _sensor_rangefinder_init,
       dim=(d.nworld, m.sensor_rangefinder_adr.size),
-      inputs=[m.sensor_objid, m.sensor_rangefinder_adr, d.site_xpos, d.site_xmat],
+      inputs=[m.sensor_objid, m.sensor_rangefinder_adr, d.site_xpos, d.site_xquat],
       outputs=[rangefinder_pnt, rangefinder_vec],
     )
 
@@ -857,18 +869,17 @@ def sensor_pos(m: Model, d: Data):
       d.qpos,
       d.xpos,
       d.xquat,
-      d.xmat,
       d.xipos,
-      d.ximat,
       d.geom_xpos,
-      d.geom_xmat,
       d.site_xpos,
-      d.site_xmat,
       d.cam_xpos,
       d.cam_xmat,
       d.subtree_com,
       d.ten_length,
       d.actuator_length,
+      d.xiquat,
+      d.geom_xquat,
+      d.site_xquat,
       rangefinder_dist,
       sensor_collision,
     ],
@@ -907,22 +918,22 @@ def _velocimeter(
   site_bodyid: wp.array(dtype=int),
   # Data in:
   site_xpos_in: wp.array2d(dtype=wp.vec3),
-  site_xmat_in: wp.array2d(dtype=wp.mat33),
   subtree_com_in: wp.array2d(dtype=wp.vec3),
   cvel_in: wp.array2d(dtype=wp.spatial_vector),
+  site_xquat_in: wp.array2d(dtype=wp.quat),
   # In:
   worldid: int,
   objid: int,
 ) -> wp.vec3:
   bodyid = site_bodyid[objid]
   pos = site_xpos_in[worldid, objid]
-  rot = site_xmat_in[worldid, objid]
+  rot = site_xquat_in[worldid, objid]
   cvel = cvel_in[worldid, bodyid]
   ang = wp.spatial_top(cvel)
   lin = wp.spatial_bottom(cvel)
   subtree_com = subtree_com_in[worldid, body_rootid[bodyid]]
   dif = pos - subtree_com
-  return wp.transpose(rot) @ (lin - wp.cross(dif, ang))
+  return math.rot_vec_quat(lin - wp.cross(dif, ang), math.quat_inv(rot))
 
 
 @wp.func
@@ -930,17 +941,17 @@ def _gyro(
   # Model:
   site_bodyid: wp.array(dtype=int),
   # Data in:
-  site_xmat_in: wp.array2d(dtype=wp.mat33),
   cvel_in: wp.array2d(dtype=wp.spatial_vector),
+  site_xquat_in: wp.array2d(dtype=wp.quat),
   # In:
   worldid: int,
   objid: int,
 ) -> wp.vec3:
   bodyid = site_bodyid[objid]
-  rot = site_xmat_in[worldid, objid]
+  rot = site_xquat_in[worldid, objid]
   cvel = cvel_in[worldid, bodyid]
   ang = wp.spatial_top(cvel)
-  return wp.transpose(rot) @ ang
+  return math.rot_vec_quat(ang, math.quat_inv(rot))
 
 
 @wp.func
@@ -1053,17 +1064,17 @@ def _frame_linvel(
   cam_bodyid: wp.array(dtype=int),
   # Data in:
   xpos_in: wp.array2d(dtype=wp.vec3),
-  xmat_in: wp.array2d(dtype=wp.mat33),
+  xquat_in: wp.array2d(dtype=wp.quat),
   xipos_in: wp.array2d(dtype=wp.vec3),
-  ximat_in: wp.array2d(dtype=wp.mat33),
   geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
   site_xpos_in: wp.array2d(dtype=wp.vec3),
-  site_xmat_in: wp.array2d(dtype=wp.mat33),
   cam_xpos_in: wp.array2d(dtype=wp.vec3),
   cam_xmat_in: wp.array2d(dtype=wp.mat33),
   subtree_com_in: wp.array2d(dtype=wp.vec3),
   cvel_in: wp.array2d(dtype=wp.spatial_vector),
+  xiquat_in: wp.array2d(dtype=wp.quat),
+  geom_xquat_in: wp.array2d(dtype=wp.quat),
+  site_xquat_in: wp.array2d(dtype=wp.quat),
   # In:
   worldid: int,
   objid: int,
@@ -1086,22 +1097,22 @@ def _frame_linvel(
 
   if reftype == ObjType.BODY:
     xposref = xipos_in[worldid, refid]
-    xmatref = ximat_in[worldid, refid]
+    xquatref = xiquat_in[worldid, refid]
   elif reftype == ObjType.XBODY:
     xposref = xpos_in[worldid, refid]
-    xmatref = xmat_in[worldid, refid]
+    xquatref = xquat_in[worldid, refid]
   elif reftype == ObjType.GEOM:
     xposref = geom_xpos_in[worldid, refid]
-    xmatref = geom_xmat_in[worldid, refid]
+    xquatref = geom_xquat_in[worldid, refid]
   elif reftype == ObjType.SITE:
     xposref = site_xpos_in[worldid, refid]
-    xmatref = site_xmat_in[worldid, refid]
+    xquatref = site_xquat_in[worldid, refid]
   elif reftype == ObjType.CAMERA:
     xposref = cam_xpos_in[worldid, refid]
     xmatref = cam_xmat_in[worldid, refid]
   else:  # UNKNOWN
     xposref = wp.vec3(0.0)
-    xmatref = wp.identity(3, dtype=float)
+    xquatref = wp.quat(1.0, 0.0, 0.0, 0.0)
 
   cvel, offset = _cvel_offset(
     body_rootid,
@@ -1145,7 +1156,10 @@ def _frame_linvel(
     xlinvelref = clinvelref - wp.cross(offsetref, cangvelref)
     rvec = xpos - xposref
     rel_vel = xlinvel - xlinvelref + wp.cross(rvec, cangvelref)
-    return wp.transpose(xmatref) @ rel_vel
+    if reftype == ObjType.CAMERA:
+      return wp.transpose(xmatref) @ rel_vel
+    else:
+      return math.rot_vec_quat(rel_vel, math.quat_inv(xquatref))
   else:
     return xlinvel
 
@@ -1159,17 +1173,17 @@ def _frame_angvel(
   cam_bodyid: wp.array(dtype=int),
   # Data in:
   xpos_in: wp.array2d(dtype=wp.vec3),
-  xmat_in: wp.array2d(dtype=wp.mat33),
+  xquat_in: wp.array2d(dtype=wp.quat),
   xipos_in: wp.array2d(dtype=wp.vec3),
-  ximat_in: wp.array2d(dtype=wp.mat33),
   geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
   site_xpos_in: wp.array2d(dtype=wp.vec3),
-  site_xmat_in: wp.array2d(dtype=wp.mat33),
   cam_xpos_in: wp.array2d(dtype=wp.vec3),
   cam_xmat_in: wp.array2d(dtype=wp.mat33),
   subtree_com_in: wp.array2d(dtype=wp.vec3),
   cvel_in: wp.array2d(dtype=wp.spatial_vector),
+  xiquat_in: wp.array2d(dtype=wp.quat),
+  geom_xquat_in: wp.array2d(dtype=wp.quat),
+  site_xquat_in: wp.array2d(dtype=wp.quat),
   # In:
   worldid: int,
   objid: int,
@@ -1197,17 +1211,17 @@ def _frame_angvel(
 
   if refid > -1:
     if reftype == ObjType.BODY:
-      xmatref = ximat_in[worldid, refid]
+      xquatref = xiquat_in[worldid, refid]
     elif reftype == ObjType.XBODY:
-      xmatref = xmat_in[worldid, refid]
+      xquatref = xquat_in[worldid, refid]
     elif reftype == ObjType.GEOM:
-      xmatref = geom_xmat_in[worldid, refid]
+      xquatref = geom_xquat_in[worldid, refid]
     elif reftype == ObjType.SITE:
-      xmatref = site_xmat_in[worldid, refid]
+      xquatref = site_xquat_in[worldid, refid]
     elif reftype == ObjType.CAMERA:
       xmatref = cam_xmat_in[worldid, refid]
     else:  # UNKNOWN
-      xmatref = wp.identity(3, dtype=float)
+      xquatref = wp.quat(1.0, 0.0, 0.0, 0.0)
 
     cvelref, _ = _cvel_offset(
       body_rootid,
@@ -1227,7 +1241,10 @@ def _frame_angvel(
     )
     cangvelref = wp.spatial_top(cvelref)
 
-    return wp.transpose(xmatref) @ (cangvel - cangvelref)
+    if reftype == ObjType.CAMERA:
+      return wp.transpose(xmatref) @ (cangvel - cangvelref)
+    else:
+      return math.rot_vec_quat(cangvel - cangvelref, math.quat_inv(xquatref))
   else:
     return cangvel
 
@@ -1262,13 +1279,10 @@ def _sensor_vel(
   # Data in:
   qvel_in: wp.array2d(dtype=float),
   xpos_in: wp.array2d(dtype=wp.vec3),
-  xmat_in: wp.array2d(dtype=wp.mat33),
+  xquat_in: wp.array2d(dtype=wp.quat),
   xipos_in: wp.array2d(dtype=wp.vec3),
-  ximat_in: wp.array2d(dtype=wp.mat33),
   geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
   site_xpos_in: wp.array2d(dtype=wp.vec3),
-  site_xmat_in: wp.array2d(dtype=wp.mat33),
   cam_xpos_in: wp.array2d(dtype=wp.vec3),
   cam_xmat_in: wp.array2d(dtype=wp.mat33),
   subtree_com_in: wp.array2d(dtype=wp.vec3),
@@ -1277,6 +1291,9 @@ def _sensor_vel(
   cvel_in: wp.array2d(dtype=wp.spatial_vector),
   subtree_linvel_in: wp.array2d(dtype=wp.vec3),
   subtree_angmom_in: wp.array2d(dtype=wp.vec3),
+  xiquat_in: wp.array2d(dtype=wp.quat),
+  geom_xquat_in: wp.array2d(dtype=wp.quat),
+  site_xquat_in: wp.array2d(dtype=wp.quat),
   # Data out:
   sensordata_out: wp.array2d(dtype=float),
 ):
@@ -1287,10 +1304,10 @@ def _sensor_vel(
   out = sensordata_out[worldid]
 
   if sensortype == SensorType.VELOCIMETER:
-    vec3 = _velocimeter(body_rootid, site_bodyid, site_xpos_in, site_xmat_in, subtree_com_in, cvel_in, worldid, objid)
+    vec3 = _velocimeter(body_rootid, site_bodyid, site_xpos_in, subtree_com_in, cvel_in, site_xquat_in, worldid, objid)
     _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
   elif sensortype == SensorType.GYRO:
-    vec3 = _gyro(site_bodyid, site_xmat_in, cvel_in, worldid, objid)
+    vec3 = _gyro(site_bodyid, cvel_in, site_xquat_in, worldid, objid)
     _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
   elif sensortype == SensorType.JOINTVEL:
     val = _joint_vel(jnt_dofadr, qvel_in, worldid, objid)
@@ -1314,17 +1331,17 @@ def _sensor_vel(
       site_bodyid,
       cam_bodyid,
       xpos_in,
-      xmat_in,
+      xquat_in,
       xipos_in,
-      ximat_in,
       geom_xpos_in,
-      geom_xmat_in,
       site_xpos_in,
-      site_xmat_in,
       cam_xpos_in,
       cam_xmat_in,
       subtree_com_in,
       cvel_in,
+      xiquat_in,
+      geom_xquat_in,
+      site_xquat_in,
       worldid,
       objid,
       objtype,
@@ -1342,17 +1359,17 @@ def _sensor_vel(
       site_bodyid,
       cam_bodyid,
       xpos_in,
-      xmat_in,
+      xquat_in,
       xipos_in,
-      ximat_in,
       geom_xpos_in,
-      geom_xmat_in,
       site_xpos_in,
-      site_xmat_in,
       cam_xpos_in,
       cam_xmat_in,
       subtree_com_in,
       cvel_in,
+      xiquat_in,
+      geom_xquat_in,
+      site_xquat_in,
       worldid,
       objid,
       objtype,
@@ -1397,13 +1414,10 @@ def sensor_vel(m: Model, d: Data):
       m.sensor_vel_adr,
       d.qvel,
       d.xpos,
-      d.xmat,
+      d.xquat,
       d.xipos,
-      d.ximat,
       d.geom_xpos,
-      d.geom_xmat,
       d.site_xpos,
-      d.site_xmat,
       d.cam_xpos,
       d.cam_xmat,
       d.subtree_com,
@@ -1412,6 +1426,9 @@ def sensor_vel(m: Model, d: Data):
       d.cvel,
       d.subtree_linvel,
       d.subtree_angmom,
+      d.xiquat,
+      d.geom_xquat,
+      d.site_xquat,
     ],
     outputs=[d.sensordata],
   )
@@ -1446,17 +1463,15 @@ def _accelerometer(
   site_bodyid: wp.array(dtype=int),
   # Data in:
   site_xpos_in: wp.array2d(dtype=wp.vec3),
-  site_xmat_in: wp.array2d(dtype=wp.mat33),
   subtree_com_in: wp.array2d(dtype=wp.vec3),
   cvel_in: wp.array2d(dtype=wp.spatial_vector),
   cacc_in: wp.array2d(dtype=wp.spatial_vector),
+  site_xquat_in: wp.array2d(dtype=wp.quat),
   # In:
   worldid: int,
   objid: int,
 ) -> wp.vec3:
   bodyid = site_bodyid[objid]
-  rot = site_xmat_in[worldid, objid]
-  rotT = wp.transpose(rot)
   cvel = cvel_in[worldid, bodyid]
   cvel_top = wp.spatial_top(cvel)
   cvel_bottom = wp.spatial_bottom(cvel)
@@ -1464,9 +1479,10 @@ def _accelerometer(
   cacc_top = wp.spatial_top(cacc)
   cacc_bottom = wp.spatial_bottom(cacc)
   dif = site_xpos_in[worldid, objid] - subtree_com_in[worldid, body_rootid[bodyid]]
-  ang = rotT @ cvel_top
-  lin = rotT @ (cvel_bottom - wp.cross(dif, cvel_top))
-  acc = rotT @ (cacc_bottom - wp.cross(dif, cacc_top))
+  rotT = math.quat_inv(site_xquat_in[worldid, objid])
+  ang = math.rot_vec_quat(cvel_top, rotT)
+  lin = math.rot_vec_quat(cvel_bottom - wp.cross(dif, cvel_top), rotT)
+  acc = math.rot_vec_quat(cacc_bottom - wp.cross(dif, cacc_top), rotT)
   correction = wp.cross(ang, lin)
   return acc + correction
 
@@ -1476,16 +1492,15 @@ def _force(
   # Model:
   site_bodyid: wp.array(dtype=int),
   # Data in:
-  site_xmat_in: wp.array2d(dtype=wp.mat33),
   cfrc_int_in: wp.array2d(dtype=wp.spatial_vector),
+  site_xquat_in: wp.array2d(dtype=wp.quat),
   # In:
   worldid: int,
   objid: int,
 ) -> wp.vec3:
   bodyid = site_bodyid[objid]
   cfrc_int = cfrc_int_in[worldid, bodyid]
-  site_xmat = site_xmat_in[worldid, objid]
-  return wp.transpose(site_xmat) @ wp.spatial_bottom(cfrc_int)
+  return math.rot_vec_quat(wp.spatial_bottom(cfrc_int), math.quat_inv(site_xquat_in[worldid, objid]))
 
 
 @wp.func
@@ -1495,18 +1510,19 @@ def _torque(
   site_bodyid: wp.array(dtype=int),
   # Data in:
   site_xpos_in: wp.array2d(dtype=wp.vec3),
-  site_xmat_in: wp.array2d(dtype=wp.mat33),
   subtree_com_in: wp.array2d(dtype=wp.vec3),
   cfrc_int_in: wp.array2d(dtype=wp.spatial_vector),
+  site_xquat_in: wp.array2d(dtype=wp.quat),
   # In:
   worldid: int,
   objid: int,
 ) -> wp.vec3:
   bodyid = site_bodyid[objid]
   cfrc_int = cfrc_int_in[worldid, bodyid]
-  site_xmat = site_xmat_in[worldid, objid]
   dif = site_xpos_in[worldid, objid] - subtree_com_in[worldid, body_rootid[bodyid]]
-  return wp.transpose(site_xmat) @ (wp.spatial_top(cfrc_int) - wp.cross(dif, wp.spatial_bottom(cfrc_int)))
+  return math.rot_vec_quat(
+    wp.spatial_top(cfrc_int) - wp.cross(dif, wp.spatial_bottom(cfrc_int)), math.quat_inv(site_xquat_in[worldid, objid])
+  )
 
 
 @wp.func
@@ -1709,7 +1725,6 @@ def _sensor_acc(
   xipos_in: wp.array2d(dtype=wp.vec3),
   geom_xpos_in: wp.array2d(dtype=wp.vec3),
   site_xpos_in: wp.array2d(dtype=wp.vec3),
-  site_xmat_in: wp.array2d(dtype=wp.mat33),
   cam_xpos_in: wp.array2d(dtype=wp.vec3),
   subtree_com_in: wp.array2d(dtype=wp.vec3),
   cvel_in: wp.array2d(dtype=wp.spatial_vector),
@@ -1726,6 +1741,7 @@ def _sensor_acc(
   efc_force_in: wp.array2d(dtype=float),
   njmax_in: int,
   nacon_in: wp.array(dtype=int),
+  site_xquat_in: wp.array2d(dtype=wp.quat),
   # In:
   sensor_contact_nmatch_in: wp.array2d(dtype=int),
   sensor_contact_matchid_in: wp.array3d(dtype=int),
@@ -1952,14 +1968,14 @@ def _sensor_acc(
 
   elif sensortype == SensorType.ACCELEROMETER:
     vec3 = _accelerometer(
-      body_rootid, site_bodyid, site_xpos_in, site_xmat_in, subtree_com_in, cvel_in, cacc_in, worldid, objid
+      body_rootid, site_bodyid, site_xpos_in, subtree_com_in, cvel_in, cacc_in, site_xquat_in, worldid, objid
     )
     _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
   elif sensortype == SensorType.FORCE:
-    vec3 = _force(site_bodyid, site_xmat_in, cfrc_int_in, worldid, objid)
+    vec3 = _force(site_bodyid, cfrc_int_in, site_xquat_in, worldid, objid)
     _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
   elif sensortype == SensorType.TORQUE:
-    vec3 = _torque(body_rootid, site_bodyid, site_xpos_in, site_xmat_in, subtree_com_in, cfrc_int_in, worldid, objid)
+    vec3 = _torque(body_rootid, site_bodyid, site_xpos_in, subtree_com_in, cfrc_int_in, site_xquat_in, worldid, objid)
     _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
   elif sensortype == SensorType.ACTUATORFRC:
     val = _actuator_force(actuator_force_in, worldid, objid)
@@ -2014,7 +2030,6 @@ def _sensor_touch(
   sensor_touch_adr: wp.array(dtype=int),
   # Data in:
   site_xpos_in: wp.array2d(dtype=wp.vec3),
-  site_xmat_in: wp.array2d(dtype=wp.mat33),
   contact_pos_in: wp.array(dtype=wp.vec3),
   contact_frame_in: wp.array(dtype=wp.mat33),
   contact_dim_in: wp.array(dtype=int),
@@ -2023,6 +2038,7 @@ def _sensor_touch(
   contact_worldid_in: wp.array(dtype=int),
   efc_force_in: wp.array2d(dtype=float),
   nacon_in: wp.array(dtype=int),
+  site_xquat_in: wp.array2d(dtype=wp.quat),
   # Data out:
   sensordata_out: wp.array2d(dtype=float),
 ):
@@ -2069,7 +2085,7 @@ def _sensor_touch(
     # add if ray-zone intersection (always true when contact.pos inside zone)
     dist, normal = ray.ray_geom(
       site_xpos_in[worldid, objid],
-      site_xmat_in[worldid, objid],
+      site_xquat_in[worldid, objid],
       site_size[objid],
       contact_pos_in[conid],
       conray,
@@ -2112,12 +2128,12 @@ def _sensor_tactile(
   taxel_sensorid: wp.array(dtype=int),
   # Data in:
   geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
   subtree_com_in: wp.array2d(dtype=wp.vec3),
   cvel_in: wp.array2d(dtype=wp.spatial_vector),
   contact_geom_in: wp.array(dtype=wp.vec2i),
   contact_worldid_in: wp.array(dtype=int),
   nacon_in: wp.array(dtype=int),
+  geom_xquat_in: wp.array2d(dtype=wp.quat),
   # Data out:
   sensordata_out: wp.array2d(dtype=float),
 ):
@@ -2151,14 +2167,15 @@ def _sensor_tactile(
   # vertex local position
   vertid = taxel_vertadr[taxelid] - mesh_vertadr[mesh_id]
   pos = mesh_vert[vertid + mesh_vertadr[mesh_id]]
+  quat = geom_xquat_in[worldid, geom_id]
 
   # position in global frame
-  xpos = geom_xmat_in[worldid, geom_id] @ pos
+  xpos = math.rot_vec_quat(pos, quat)
   xpos += geom_xpos_in[worldid, geom_id]
 
   # position in other geom frame
   tmp = xpos - geom_xpos_in[worldid, geom]
-  lpos = wp.transpose(geom_xmat_in[worldid, geom]) @ tmp
+  lpos = math.rot_vec_quat(tmp, math.quat_inv(quat))
 
   plugin_id = geom_plugin_index[geom]
 
@@ -2237,7 +2254,6 @@ def _contact_match(
   sensor_contact_adr: wp.array(dtype=int),
   # Data in:
   site_xpos_in: wp.array2d(dtype=wp.vec3),
-  site_xmat_in: wp.array2d(dtype=wp.mat33),
   contact_dist_in: wp.array(dtype=float),
   contact_pos_in: wp.array(dtype=wp.vec3),
   contact_frame_in: wp.array(dtype=wp.mat33),
@@ -2250,6 +2266,7 @@ def _contact_match(
   efc_force_in: wp.array2d(dtype=float),
   njmax_in: int,
   nacon_in: wp.array(dtype=int),
+  site_xquat_in: wp.array2d(dtype=wp.quat),
   # Out:
   sensor_contact_nmatch_out: wp.array2d(dtype=int),
   sensor_contact_matchid_out: wp.array3d(dtype=int),
@@ -2277,7 +2294,7 @@ def _contact_match(
   # site filter
   if objtype == ObjType.SITE:
     if not inside_geom(
-      site_xpos_in[worldid, objid], site_xmat_in[worldid, objid], site_size[objid], site_type[objid], contact_pos_in[contactid]
+      site_xpos_in[worldid, objid], site_xquat_in[worldid, objid], site_size[objid], site_type[objid], contact_pos_in[contactid]
     ):
       return
 
@@ -2411,7 +2428,6 @@ def sensor_acc(m: Model, d: Data):
       m.sensor_adr,
       m.sensor_touch_adr,
       d.site_xpos,
-      d.site_xmat,
       d.contact.pos,
       d.contact.frame,
       d.contact.dim,
@@ -2420,6 +2436,7 @@ def sensor_acc(m: Model, d: Data):
       d.contact.worldid,
       d.efc.force,
       d.nacon,
+      d.site_xquat,
     ],
     outputs=[
       d.sensordata,
@@ -2453,12 +2470,12 @@ def sensor_acc(m: Model, d: Data):
       m.taxel_vertadr,
       m.taxel_sensorid,
       d.geom_xpos,
-      d.geom_xmat,
       d.subtree_com,
       d.cvel,
       d.contact.geom,
       d.contact.worldid,
       d.nacon,
+      d.geom_xquat,
     ],
     outputs=[
       d.sensordata,
@@ -2492,7 +2509,6 @@ def sensor_acc(m: Model, d: Data):
         m.sensor_intprm,
         m.sensor_contact_adr,
         d.site_xpos,
-        d.site_xmat,
         d.contact.dist,
         d.contact.pos,
         d.contact.frame,
@@ -2505,6 +2521,7 @@ def sensor_acc(m: Model, d: Data):
         d.efc.force,
         d.njmax,
         d.nacon,
+        d.site_xquat,
       ],
       outputs=[sensor_contact_nmatch, sensor_contact_matchid, sensor_contact_criteria, sensor_contact_direction],
     )
@@ -2545,7 +2562,6 @@ def sensor_acc(m: Model, d: Data):
       d.xipos,
       d.geom_xpos,
       d.site_xpos,
-      d.site_xmat,
       d.cam_xpos,
       d.subtree_com,
       d.cvel,
@@ -2562,6 +2578,7 @@ def sensor_acc(m: Model, d: Data):
       d.efc.force,
       d.njmax,
       d.nacon,
+      d.site_xquat,
       sensor_contact_nmatch,
       sensor_contact_matchid,
       sensor_contact_direction,
diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py
index 3e8c03b34..82255a891 100644
--- a/mujoco_warp/_src/smooth.py
+++ b/mujoco_warp/_src/smooth.py
@@ -62,17 +62,15 @@ def _kinematics_level(
   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),
   # In:
   body_tree_: wp.array(dtype=int),
   # 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),
   xanchor_out: wp.array2d(dtype=wp.vec3),
   xaxis_out: wp.array2d(dtype=wp.vec3),
+  xiquat_out: wp.array2d(dtype=wp.quat),
 ):
   worldid, nodeid = wp.tid()
   bodyid = body_tree_[nodeid]
@@ -113,7 +111,7 @@ def _kinematics_level(
       xquat = body_quat[body_quat_id, bodyid]
 
     if pid >= 0:
-      xpos = xmat_in[worldid, pid] @ xpos + xpos_in[worldid, pid]
+      xpos = math.rot_vec_quat(xpos, xquat_in[worldid, pid]) + xpos_in[worldid, pid]
       xquat = math.mul_quat(xquat_in[worldid, pid], xquat)
 
     for _ in range(jntnum):
@@ -145,9 +143,8 @@ def _kinematics_level(
   xpos_out[worldid, bodyid] = xpos
   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]))
+  xiquat_out[worldid, bodyid] = math.mul_quat(xquat, body_iquat[worldid % body_iquat.shape[0], bodyid])
 
 
 @wp.kernel
@@ -164,7 +161,7 @@ def _geom_local_to_global(
   xquat_in: wp.array2d(dtype=wp.quat),
   # Data out:
   geom_xpos_out: wp.array2d(dtype=wp.vec3),
-  geom_xmat_out: wp.array2d(dtype=wp.mat33),
+  geom_xquat_out: wp.array2d(dtype=wp.quat),
 ):
   worldid, geomid = wp.tid()
   bodyid = geom_bodyid[geomid]
@@ -177,7 +174,7 @@ def _geom_local_to_global(
   xpos = xpos_in[worldid, bodyid]
   xquat = xquat_in[worldid, bodyid]
   geom_xpos_out[worldid, geomid] = xpos + math.rot_vec_quat(geom_pos[worldid % geom_pos.shape[0], geomid], xquat)
-  geom_xmat_out[worldid, geomid] = math.quat_to_mat(math.mul_quat(xquat, geom_quat[worldid % geom_quat.shape[0], geomid]))
+  geom_xquat_out[worldid, geomid] = math.mul_quat(xquat, geom_quat[worldid % geom_quat.shape[0], geomid])
 
 
 @wp.kernel
@@ -191,14 +188,14 @@ def _site_local_to_global(
   xquat_in: wp.array2d(dtype=wp.quat),
   # Data out:
   site_xpos_out: wp.array2d(dtype=wp.vec3),
-  site_xmat_out: wp.array2d(dtype=wp.mat33),
+  site_xquat_out: wp.array2d(dtype=wp.quat),
 ):
   worldid, siteid = wp.tid()
   bodyid = site_bodyid[siteid]
   xpos = xpos_in[worldid, bodyid]
   xquat = xquat_in[worldid, bodyid]
   site_xpos_out[worldid, siteid] = xpos + math.rot_vec_quat(site_pos[worldid % site_pos.shape[0], siteid], xquat)
-  site_xmat_out[worldid, siteid] = math.quat_to_mat(math.mul_quat(xquat, site_quat[worldid % site_quat.shape[0], siteid]))
+  site_xquat_out[worldid, siteid] = math.mul_quat(xquat, site_quat[worldid % site_quat.shape[0], siteid])
 
 
 @wp.kernel
@@ -300,24 +297,23 @@ def kinematics(m: Model, d: Data):
         d.mocap_quat,
         d.xpos,
         d.xquat,
-        d.xmat,
         body_tree,
       ],
-      outputs=[d.xpos, d.xquat, d.xmat, d.xipos, d.ximat, d.xanchor, d.xaxis],
+      outputs=[d.xpos, d.xquat, d.xipos, d.xanchor, d.xaxis, d.xiquat],
     )
 
   wp.launch(
     _geom_local_to_global,
     dim=(d.nworld, m.ngeom),
     inputs=[m.body_rootid, m.body_weldid, m.body_mocapid, m.geom_bodyid, m.geom_pos, m.geom_quat, d.xpos, d.xquat],
-    outputs=[d.geom_xpos, d.geom_xmat],
+    outputs=[d.geom_xpos, d.geom_xquat],
   )
 
   wp.launch(
     _site_local_to_global,
     dim=(d.nworld, m.nsite),
     inputs=[m.site_bodyid, m.site_pos, m.site_quat, d.xpos, d.xquat],
-    outputs=[d.site_xpos, d.site_xmat],
+    outputs=[d.site_xpos, d.site_xquat],
   )
 
 
@@ -403,13 +399,13 @@ def _cinert(
   body_inertia: wp.array2d(dtype=wp.vec3),
   # Data in:
   xipos_in: wp.array2d(dtype=wp.vec3),
-  ximat_in: wp.array2d(dtype=wp.mat33),
   subtree_com_in: wp.array2d(dtype=wp.vec3),
+  xiquat_in: wp.array2d(dtype=wp.quat),
   # Data out:
   cinert_out: wp.array2d(dtype=vec10),
 ):
   worldid, bodyid = wp.tid()
-  mat = ximat_in[worldid, bodyid]
+  mat = math.quat_to_mat(xiquat_in[worldid, bodyid])
   inert = body_inertia[worldid % body_inertia.shape[0], bodyid]
   mass = body_mass[worldid % body_mass.shape[0], bodyid]
   dif = xipos_in[worldid, bodyid] - subtree_com_in[worldid, body_rootid[bodyid]]
@@ -449,7 +445,7 @@ def _cdof(
   jnt_dofadr: wp.array(dtype=int),
   jnt_bodyid: wp.array(dtype=int),
   # Data in:
-  xmat_in: wp.array2d(dtype=wp.mat33),
+  xquat_in: wp.array2d(dtype=wp.quat),
   xanchor_in: wp.array2d(dtype=wp.vec3),
   xaxis_in: wp.array2d(dtype=wp.vec3),
   subtree_com_in: wp.array2d(dtype=wp.vec3),
@@ -461,7 +457,7 @@ def _cdof(
   dofid = jnt_dofadr[jntid]
   jnt_type_ = jnt_type[jntid]
   xaxis = xaxis_in[worldid, jntid]
-  xmat = wp.transpose(xmat_in[worldid, bodyid])
+  xmat = wp.transpose(math.quat_to_mat(xquat_in[worldid, bodyid]))
 
   # compute com-anchor vector
   offset = subtree_com_in[worldid, body_rootid[bodyid]] - xanchor_in[worldid, jntid]
@@ -509,13 +505,13 @@ def com_pos(m: Model, d: Data):
   wp.launch(
     _cinert,
     dim=(d.nworld, m.nbody),
-    inputs=[m.body_rootid, m.body_mass, m.body_inertia, d.xipos, d.ximat, d.subtree_com],
+    inputs=[m.body_rootid, m.body_mass, m.body_inertia, d.xipos, d.subtree_com, d.xiquat],
     outputs=[d.cinert],
   )
   wp.launch(
     _cdof,
     dim=(d.nworld, m.njnt),
-    inputs=[m.body_rootid, m.jnt_type, m.jnt_dofadr, m.jnt_bodyid, d.xmat, d.xanchor, d.xaxis, d.subtree_com],
+    inputs=[m.body_rootid, m.jnt_type, m.jnt_dofadr, m.jnt_bodyid, d.xquat, d.xanchor, d.xaxis, d.subtree_com],
     outputs=[d.cdof],
   )
 
@@ -1150,7 +1146,7 @@ def _cfrc_ext_equality(
   eq_data: wp.array2d(dtype=vec11),
   # Data in:
   xpos_in: wp.array2d(dtype=wp.vec3),
-  xmat_in: wp.array2d(dtype=wp.mat33),
+  xquat_in: wp.array2d(dtype=wp.quat),
   subtree_com_in: wp.array2d(dtype=wp.vec3),
   efc_id_in: wp.array2d(dtype=int),
   efc_force_in: wp.array2d(dtype=float),
@@ -1207,7 +1203,7 @@ def _cfrc_ext_equality(
       offset = site_pos[worldid, obj1]
 
     # transform point on body1: local -> global
-    pos = xmat_in[worldid, bodyid1] @ offset + xpos_in[worldid, bodyid1]
+    pos = math.rot_vec_quat(offset, xquat_in[worldid, bodyid1]) + xpos_in[worldid, bodyid1]
 
     # subtree CoM-based torque_force vector
     newpos = subtree_com_in[worldid, body_rootid[bodyid1]]
@@ -1229,7 +1225,7 @@ def _cfrc_ext_equality(
       offset = site_pos[worldid, obj2]
 
     # transform point on body2: local -> global
-    pos = xmat_in[worldid, bodyid2] @ offset + xpos_in[worldid, bodyid2]
+    pos = math.rot_vec_quat(offset, xquat_in[worldid, bodyid2]) + xpos_in[worldid, bodyid2]
 
     # subtree CoM-based torque_force vector
     newpos = subtree_com_in[worldid, body_rootid[bodyid2]]
@@ -1336,7 +1332,7 @@ def rne_postconstraint(m: Model, d: Data):
       m.eq_objtype,
       m.eq_data,
       d.xpos,
-      d.xmat,
+      d.xquat,
       d.subtree_com,
       d.efc.id,
       d.efc.force,
@@ -1777,11 +1773,11 @@ def _transmission(
   qpos_in: wp.array2d(dtype=float),
   xquat_in: wp.array2d(dtype=wp.quat),
   site_xpos_in: wp.array2d(dtype=wp.vec3),
-  site_xmat_in: wp.array2d(dtype=wp.mat33),
   subtree_com_in: wp.array2d(dtype=wp.vec3),
   cdof_in: wp.array2d(dtype=wp.spatial_vector),
   ten_J_in: wp.array3d(dtype=float),
   ten_length_in: wp.array2d(dtype=float),
+  site_xquat_in: wp.array2d(dtype=wp.quat),
   # Data out:
   actuator_length_out: wp.array2d(dtype=float),
   actuator_moment_out: wp.array3d(dtype=float),
@@ -1834,8 +1830,7 @@ def _transmission(
     idslider = trnid[1]
     gear0 = gear[0]
     rod = actuator_cranklength[actid]
-    site_xmat = site_xmat_in[worldid, idslider]
-    axis = wp.vec3(site_xmat[0, 2], site_xmat[1, 2], site_xmat[2, 2])
+    axis = math.rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), site_xquat_in[worldid, idslider])
     site_xpos_id = site_xpos_in[worldid, id]
     site_xpos_idslider = site_xpos_in[worldid, idslider]
     vec = site_xpos_id - site_xpos_idslider
@@ -1922,9 +1917,9 @@ def _transmission(
     # reference site undefined
     if refid == -1:
       # wrench: gear expressed in global frame
-      site_xmat = site_xmat_in[worldid, siteid]
-      wrench_translation = site_xmat @ gear_translation
-      wrench_rotation = site_xmat @ gear_rotational
+      site_xquat = site_xquat_in[worldid, siteid]
+      wrench_translation = math.rot_vec_quat(gear_translation, site_xquat)
+      wrench_rotation = math.rot_vec_quat(gear_rotational, site_xquat)
 
       # moment: global Jacobian projected on wrench
       # TODO(team): parallelize
@@ -1975,16 +1970,16 @@ def _transmission(
 
       site_xpos = site_xpos_in[worldid, siteid]
       ref_xpos = site_xpos_in[worldid, refid]
-      ref_xmat = site_xmat_in[worldid, refid]
+      ref_xquat = site_xquat_in[worldid, refid]
 
       length = float(0.0)
 
       if translational_transmission:
         # vec: site position in reference site frame
-        vec = wp.transpose(ref_xmat) @ (site_xpos - ref_xpos)
+        vec = math.rot_vec_quat(site_xpos - ref_xpos, math.quat_inv(ref_xquat))
         length += wp.dot(vec, gear_translation)
 
-        wrench_translation = ref_xmat @ gear_translation
+        wrench_translation = math.rot_vec_quat(gear_translation, ref_xquat)
 
       if rotational_transmission:
         # get site and refsite quats from parent bodies (avoid converting matrix to quat)
@@ -1995,7 +1990,7 @@ def _transmission(
         vec = math.quat_sub(quat, refquat)
         length += wp.dot(vec, gear_rotational)
 
-        wrench_rotation = ref_xmat @ gear_rotational
+        wrench_rotation = math.rot_vec_quat(gear_rotational, ref_xquat)
 
       actuator_length_out[worldid, actid] = length
 
@@ -2180,11 +2175,11 @@ def transmission(m: Model, d: Data):
       d.qpos,
       d.xquat,
       d.site_xpos,
-      d.site_xmat,
       d.subtree_com,
       d.cdof,
       d.ten_J,
       d.ten_length,
+      d.site_xquat,
     ],
     outputs=[d.actuator_length, d.actuator_moment],
   )
@@ -2443,9 +2438,9 @@ def _subtree_vel_forward(
   body_inertia: wp.array2d(dtype=wp.vec3),
   # Data in:
   xipos_in: wp.array2d(dtype=wp.vec3),
-  ximat_in: wp.array2d(dtype=wp.mat33),
   subtree_com_in: wp.array2d(dtype=wp.vec3),
   cvel_in: wp.array2d(dtype=wp.spatial_vector),
+  xiquat_in: wp.array2d(dtype=wp.quat),
   # Data out:
   subtree_linvel_out: wp.array2d(dtype=wp.vec3),
   subtree_angmom_out: wp.array2d(dtype=wp.vec3),
@@ -2459,18 +2454,18 @@ def _subtree_vel_forward(
   ang = wp.spatial_top(cvel)
   lin = wp.spatial_bottom(cvel)
   xipos = xipos_in[worldid, bodyid]
-  ximat = ximat_in[worldid, bodyid]
+  ximat = math.quat_to_mat(xiquat_in[worldid, bodyid])
   subtree_com_root = subtree_com_in[worldid, body_rootid[bodyid]]
 
   # update linear velocity
   lin -= wp.cross(xipos - subtree_com_root, ang)
 
   subtree_linvel_out[worldid, bodyid] = body_mass[body_mass_id, bodyid] * lin
-  dv = wp.transpose(ximat) @ ang
+  dv = math.rot_vec_quat(ang, math.quat_inv(xiquat_in[worldid, bodyid]))
   dv[0] *= body_inertia[body_inertia_id, bodyid][0]
   dv[1] *= body_inertia[body_inertia_id, bodyid][1]
   dv[2] *= body_inertia[body_inertia_id, bodyid][2]
-  subtree_angmom_out[worldid, bodyid] = ximat @ dv
+  subtree_angmom_out[worldid, bodyid] = math.rot_vec_quat(dv, xiquat_in[worldid, bodyid])
   subtree_bodyvel_out[worldid, bodyid] = wp.spatial_vector(ang, lin)
 
 
@@ -2557,7 +2552,7 @@ def subtree_vel(m: Model, d: Data):
   wp.launch(
     _subtree_vel_forward,
     dim=(d.nworld, m.nbody),
-    inputs=[m.body_rootid, m.body_mass, m.body_inertia, d.xipos, d.ximat, d.subtree_com, d.cvel],
+    inputs=[m.body_rootid, m.body_mass, m.body_inertia, d.xipos, d.subtree_com, d.cvel, d.xiquat],
     outputs=[d.subtree_linvel, d.subtree_angmom, d.subtree_bodyvel],
   )
 
@@ -2692,10 +2687,10 @@ def _spatial_geom_tendon(
   wrap_pulley_scale: wp.array(dtype=float),
   # Data in:
   geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
   site_xpos_in: wp.array2d(dtype=wp.vec3),
   subtree_com_in: wp.array2d(dtype=wp.vec3),
   cdof_in: wp.array2d(dtype=wp.spatial_vector),
+  geom_xquat_in: wp.array2d(dtype=wp.quat),
   # Data out:
   ten_J_out: wp.array3d(dtype=float),
   ten_length_out: wp.array2d(dtype=float),
@@ -2720,7 +2715,7 @@ def _spatial_geom_tendon(
 
   # get geom information
   geom_xpos = geom_xpos_in[worldid, wrap_objid_geom]
-  geom_xmat = geom_xmat_in[worldid, wrap_objid_geom]
+  geom_xquat = geom_xquat_in[worldid, wrap_objid_geom]
   geomsize = geom_size[worldid, wrap_objid_geom][0]
   geom_type = wrap_type[wrap_adr]
 
@@ -2737,7 +2732,7 @@ def _spatial_geom_tendon(
     side = wp.vec3(wp.inf)
 
   # compute geom wrap length and connect points (if wrap occurs)
-  length_geomgeom, geom_pnt0, geom_pnt1 = util_misc.wrap(site_pnt0, site_pnt1, geom_xpos, geom_xmat, geomsize, geom_type, side)
+  length_geomgeom, geom_pnt0, geom_pnt1 = util_misc.wrap(site_pnt0, site_pnt1, geom_xpos, geom_xquat, geomsize, geom_type, side)
 
   # store geom points
   wrap_geom_xpos_out[worldid, elementid] = wp.spatial_vector(geom_pnt0, geom_pnt1)
@@ -3050,10 +3045,10 @@ def tendon(m: Model, d: Data):
       m.wrap_geom_adr,
       m.wrap_pulley_scale,
       d.geom_xpos,
-      d.geom_xmat,
       d.site_xpos,
       d.subtree_com,
       d.cdof,
+      d.geom_xquat,
     ],
     outputs=[d.ten_J, d.ten_length, wrap_geom_xpos],
   )
diff --git a/mujoco_warp/_src/smooth_test.py b/mujoco_warp/_src/smooth_test.py
index cc94e3196..a9b1a0aeb 100644
--- a/mujoco_warp/_src/smooth_test.py
+++ b/mujoco_warp/_src/smooth_test.py
@@ -121,7 +121,7 @@ def test_kinematics(self, make_data):
       mjd = mujoco.MjData(mjm)
       d = mjw.make_data(mjm, nworld=nworld)
 
-    for arr in (d.xpos, d.xipos, d.xquat, d.xmat, d.ximat, d.xanchor, d.xaxis, d.site_xpos, d.site_xmat):
+    for arr in (d.xpos, d.xipos, d.xquat, d.xiquat, d.xanchor, d.xaxis, d.site_xpos, d.site_xquat):
       arr_view = arr[:, 1:]  # skip world body
       arr_view.fill_(wp.inf)
 
@@ -140,18 +140,32 @@ def test_kinematics(self, make_data):
     mujoco.mj_kinematics(mjm, mjd)
     mjw.kinematics(m, d)
 
+    ximat = np.zeros((nworld, m.nbody, 9), dtype=np.float64)
+    for i in range(nworld):
+      for j in range(m.nbody):
+        mujoco.mju_quat2Mat(ximat[i, j], d.xiquat.numpy()[i, j])
+
+    geom_xmat = np.zeros((nworld, m.ngeom, 9), dtype=np.float64)
+    for i in range(nworld):
+      for j in range(m.ngeom):
+        mujoco.mju_quat2Mat(geom_xmat[i, j], d.geom_xquat.numpy()[i, j])
+
+    site_xmat = np.zeros((nworld, m.nsite, 9), dtype=np.float64)
+    for i in range(nworld):
+      for j in range(m.nsite):
+        mujoco.mju_quat2Mat(site_xmat[i, j], d.site_xquat.numpy()[i, j])
+
     for i in range(nworld):
       _assert_eq(d.xanchor.numpy()[i], mjd.xanchor, "xanchor")
       _assert_eq(d.xaxis.numpy()[i], mjd.xaxis, "xaxis")
       _assert_eq(d.xpos.numpy()[i], mjd.xpos, "xpos")
       _assert_eq(d.xquat.numpy()[i], mjd.xquat, "xquat")
-      _assert_eq(d.xmat.numpy()[i], mjd.xmat.reshape((-1, 3, 3)), "xmat")
       _assert_eq(d.xipos.numpy()[i], mjd.xipos, "xipos")
-      _assert_eq(d.ximat.numpy()[i], mjd.ximat.reshape((-1, 3, 3)), "ximat")
+      _assert_eq(ximat[i], mjd.ximat, "ximat")
       _assert_eq(d.geom_xpos.numpy()[i], mjd.geom_xpos, "geom_xpos")
-      _assert_eq(d.geom_xmat.numpy()[i], mjd.geom_xmat.reshape((-1, 3, 3)), "geom_xmat")
+      _assert_eq(geom_xmat[i], mjd.geom_xmat, "geom_xmat")
       _assert_eq(d.site_xpos.numpy()[i], mjd.site_xpos, "site_xpos")
-      _assert_eq(d.site_xmat.numpy()[i], mjd.site_xmat.reshape((-1, 3, 3)), "site_xmat")
+      _assert_eq(site_xmat[i], mjd.site_xmat, "site_xmat")
       _assert_eq(d.mocap_pos.numpy()[i], mjd.mocap_pos, "mocap_pos")
       _assert_eq(d.mocap_quat.numpy()[i], mjd.mocap_quat, "mocap_quat")
 
diff --git a/mujoco_warp/_src/types.py b/mujoco_warp/_src/types.py
index 6ff16d703..4566cdf01 100644
--- a/mujoco_warp/_src/types.py
+++ b/mujoco_warp/_src/types.py
@@ -1571,15 +1571,11 @@ class Data:
     sensordata: sensor data array                               (nworld, nsensordata,)
     xpos: Cartesian position of body frame                      (nworld, nbody, 3)
     xquat: Cartesian orientation of body frame                  (nworld, nbody, 4)
-    xmat: Cartesian orientation of body frame                   (nworld, nbody, 3, 3)
     xipos: Cartesian position of body com                       (nworld, nbody, 3)
-    ximat: Cartesian orientation of body inertia                (nworld, nbody, 3, 3)
     xanchor: Cartesian position of joint anchor                 (nworld, njnt, 3)
     xaxis: Cartesian joint axis                                 (nworld, njnt, 3)
     geom_xpos: Cartesian geom position                          (nworld, ngeom, 3)
-    geom_xmat: Cartesian geom orientation                       (nworld, ngeom, 3, 3)
     site_xpos: Cartesian site position                          (nworld, nsite, 3)
-    site_xmat: Cartesian site orientation                       (nworld, nsite, 3, 3)
     cam_xpos: Cartesian camera position                         (nworld, ncam, 3)
     cam_xmat: Cartesian camera orientation                      (nworld, ncam, 3, 3)
     light_xpos: Cartesian light position                        (nworld, nlight, 3)
@@ -1647,6 +1643,9 @@ class Data:
     collision_pairid: ids from broadphase                       (naconmax, 2)
     collision_worldid: collision world ids from broadphase      (naconmax,)
     ncollision: collision count from broadphase                 (1,)
+    xiquat: Cartesian orientation of body inertia               (nworld, nbody, 4)
+    geom_xquat: Cartesian orientation of geom                   (nworld, ngeom, 4)
+    site_xquat: Cartesian orientation of site                   (nworld, nsite, 4)
   """
 
   solver_niter: array("nworld", int)
@@ -1671,15 +1670,11 @@ class Data:
   sensordata: array("nworld", "nsensordata", float)
   xpos: array("nworld", "nbody", wp.vec3)
   xquat: array("nworld", "nbody", wp.quat)
-  xmat: array("nworld", "nbody", wp.mat33)
   xipos: array("nworld", "nbody", wp.vec3)
-  ximat: array("nworld", "nbody", wp.mat33)
   xanchor: array("nworld", "njnt", wp.vec3)
   xaxis: array("nworld", "njnt", wp.vec3)
   geom_xpos: array("nworld", "ngeom", wp.vec3)
-  geom_xmat: array("nworld", "ngeom", wp.mat33)
   site_xpos: array("nworld", "nsite", wp.vec3)
-  site_xmat: array("nworld", "nsite", wp.mat33)
   cam_xpos: array("nworld", "ncam", wp.vec3)
   cam_xmat: array("nworld", "ncam", wp.mat33)
   light_xpos: array("nworld", "nlight", wp.vec3)
@@ -1745,3 +1740,7 @@ class Data:
   collision_pairid: array("naconmax", wp.vec2i)
   collision_worldid: array("naconmax", int)
   ncollision: array(1, int)
+
+  xiquat: array("nworld", "nbody", wp.quat)
+  geom_xquat: array("nworld", "ngeom", wp.quat)
+  site_xquat: array("nworld", "nsite", wp.quat)
diff --git a/mujoco_warp/_src/util_misc.py b/mujoco_warp/_src/util_misc.py
index 8c48fbe0d..c5507c66e 100644
--- a/mujoco_warp/_src/util_misc.py
+++ b/mujoco_warp/_src/util_misc.py
@@ -324,7 +324,7 @@ def wrap_inside(
 
 @wp.func
 def wrap(
-  x0: wp.vec3, x1: wp.vec3, pos: wp.vec3, mat: wp.mat33, radius: float, geomtype: int, side: wp.vec3
+  x0: wp.vec3, x1: wp.vec3, pos: wp.vec3, quat: wp.quat, radius: float, geomtype: int, side: wp.vec3
 ) -> Tuple[float, wp.vec3, wp.vec3]:
   """Wrap tendons around spheres and cylinders.
 
@@ -332,7 +332,7 @@ def wrap(
     x0: 3D endpoint.
     x1: 3D endpoint.
     pos: Position of geom.
-    mat: Orientation of geom.
+    quat: Orientation of geom.
     radius: Geom radius.
     geomtype: Wrap type (mjtWrap).
     side: 3D position for sidesite, no side point: wp.vec3(wp.inf).
@@ -345,9 +345,9 @@ def wrap(
     return wp.inf, wp.vec3(wp.inf), wp.vec3(wp.inf)
 
   # map sites to wrap object's local frame
-  matT = wp.transpose(mat)
-  p0 = matT @ (x0 - pos)
-  p1 = matT @ (x1 - pos)
+  matT = math.quat_inv(quat)
+  p0 = math.rot_vec_quat(x0 - pos, matT)
+  p1 = math.rot_vec_quat(x1 - pos, matT)
 
   # too close to origin: return
   if (wp.norm_l2(p0) < MJ_MINVAL) or (wp.norm_l2(p1) < MJ_MINVAL):
@@ -403,7 +403,7 @@ def wrap(
 
   if valid_side:
     # side point: apply same projection as x0, x1
-    sidepnt = matT @ (side - pos)
+    sidepnt = math.rot_vec_quat(side - pos, matT)
 
     # side point: project and rescale
     sidepnt_proj = wp.vec2(
@@ -443,8 +443,8 @@ def wrap(
     wlen = wp.sqrt(wlen * wlen + height * height)
 
   # map back to global frame: wpnt
-  wpnt0 = mat @ res0 + pos
-  wpnt1 = mat @ res1 + pos
+  wpnt0 = math.rot_vec_quat(res0, quat) + pos
+  wpnt1 = math.rot_vec_quat(res1, quat) + pos
 
   return wlen, wpnt0, wpnt1
 
@@ -599,7 +599,7 @@ def muscle_dynamics(control: float, activation: float, prm: vec10) -> float:
 
 
 @wp.func
-def inside_geom(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, geomtype: int, point: wp.vec3) -> bool:
+def inside_geom(pos: wp.vec3, quat: wp.quat, size: wp.vec3, geomtype: int, point: wp.vec3) -> bool:
   """Return True if point is inside primitive geom, False otherwise."""
   # vector from geom to point
   vec = point - pos
@@ -609,6 +609,7 @@ def inside_geom(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, geomtype: int, point
     return wp.dot(vec, vec) < size[0] * size[0]
 
   # rotate into local frame
+  mat = math.quat_to_mat(quat)
   plocal = wp.transpose(mat) @ vec
 
   # handle other geom types
diff --git a/mujoco_warp/_src/util_misc_test.py b/mujoco_warp/_src/util_misc_test.py
index 935bef3d0..d3fae1e9c 100644
--- a/mujoco_warp/_src/util_misc_test.py
+++ b/mujoco_warp/_src/util_misc_test.py
@@ -166,7 +166,7 @@ def _wrap(
   x0: np.array,
   x1: np.array,
   xpos: np.array,
-  xmat: np.array,
+  xquat: np.array,
   radius: float,
   geomtype: int,
   side: np.array,
@@ -181,7 +181,7 @@ def wrap(
     x0: wp.vec3,
     x1: wp.vec3,
     pos: wp.vec3,
-    mat: wp.mat33,
+    quat: wp.quat,
     radius: float,
     geomtype: int,
     side: wp.vec3,
@@ -190,7 +190,7 @@ def wrap(
     wpnt0_out: wp.array(dtype=wp.vec3),
     wpnt1_out: wp.array(dtype=wp.vec3),
   ):
-    length_, wpnt0_, wpnt1_ = util_misc.wrap(x0, x1, pos, mat, radius, geomtype, side)
+    length_, wpnt0_, wpnt1_ = util_misc.wrap(x0, x1, pos, quat, radius, geomtype, side)
     length_out[0] = length_
     wpnt0_out[0] = wpnt0_
     wpnt1_out[0] = wpnt1_
@@ -202,17 +202,7 @@ def wrap(
       wp.vec3(x0[0], x0[1], x0[2]),
       wp.vec3(x1[0], x1[1], x1[2]),
       wp.vec3(xpos[0], xpos[1], xpos[2]),
-      wp.mat33(
-        xmat[0, 0],
-        xmat[0, 1],
-        xmat[0, 2],
-        xmat[1, 0],
-        xmat[1, 1],
-        xmat[1, 2],
-        xmat[2, 0],
-        xmat[2, 1],
-        xmat[2, 2],
-      ),
+      wp.quat(xquat[0], xquat[1], xquat[2], xquat[3]),
       radius,
       geomtype,
       wp.vec3(side[0], side[1], side[2]),
@@ -441,11 +431,11 @@ def test_wrap(self, wraptype):
     x0 = np.array([1, 1, 1])
     x1 = np.array([2, 2, 2])
     xpos = np.array([0, 0, 0])
-    xmat = np.eye(3)
+    xquat = np.array([1, 0, 0, 0])
     radius = 0.1
     side = np.array([np.inf, np.inf, np.inf])
 
-    wlen, wpnt0, wpnt1 = _wrap(x0, x1, xpos, xmat, radius, wraptype, side)
+    wlen, wpnt0, wpnt1 = _wrap(x0, x1, xpos, xquat, radius, wraptype, side)
 
     _assert_eq(wlen, -1.0, "wlen")
     _assert_eq(wpnt0, np.array([np.inf, np.inf, np.inf]), "wpnt0")
@@ -455,11 +445,11 @@ def test_wrap(self, wraptype):
     x0 = np.array([0.1, -1.0, 0.0])
     x1 = np.array([0.1, 1.0, 0.0])
     xpos = np.array([0, 0, 0])
-    xmat = np.eye(3)
+    xquat = np.array([1, 0, 0, 0])
     radius = 0.1
     side = np.array([np.inf, np.inf, np.inf])
 
-    wlen, wpnt0, wpnt1 = _wrap(x0, x1, xpos, xmat, radius, wraptype, side)
+    wlen, wpnt0, wpnt1 = _wrap(x0, x1, xpos, xquat, radius, wraptype, side)
 
     _assert_eq(wlen, 0.0, "wlen")
     _assert_eq(wpnt0, np.array([0.1, 0.0, 0.0]), "wpnt0")
@@ -469,11 +459,11 @@ def test_wrap(self, wraptype):
     x0 = np.array([MJ_MINVAL, -100.0, 0.0])
     x1 = np.array([MJ_MINVAL, 100.0, 0.0])
     xpos = np.array([0, 0, 0])
-    xmat = np.eye(3)
+    xquat = np.array([1, 0, 0, 0])
     radius = 0.1
     side = np.array([radius + 10 * MJ_MINVAL, 0, 0])
 
-    wlen, wpnt0, wpnt1 = _wrap(x0, x1, xpos, xmat, radius, wraptype, side)
+    wlen, wpnt0, wpnt1 = _wrap(x0, x1, xpos, xquat, radius, wraptype, side)
 
     # wlen, wpnt0[1], wpnt1[1] are ~0
     _assert_eq(wlen, 0, "wlen")
@@ -484,12 +474,12 @@ def test_wrap(self, wraptype):
     x0 = np.array([0.0, -1.0, 0.0])
     x1 = np.array([0.0, 1.0, 0.0])
     xpos = np.array([0, 0, 0])
-    xmat = np.eye(3)
+    xquat = np.array([1, 0, 0, 0])
     radius = 0.1
     wraptype = WrapType.CYLINDER
     side = np.array([0, 0, 0])
 
-    wlen, wpnt0, wpnt1 = _wrap(x0, x1, xpos, xmat, radius, wraptype, side)
+    wlen, wpnt0, wpnt1 = _wrap(x0, x1, xpos, xquat, radius, wraptype, side)
 
     _assert_eq(wlen, -1.0, "wlen")
     _assert_eq(wpnt0, np.array([np.inf, np.inf, np.inf]), "wpnt0")
@@ -499,11 +489,11 @@ def test_wrap(self, wraptype):
     x0 = np.array([1.0, -1.0, 0.0])
     x1 = np.array([1.0, 1.0, 0.0])
     xpos = np.array([0, 0, 0])
-    xmat = np.eye(3)
+    xquat = np.array([1, 0, 0, 0])
     radius = 0.1
     side = np.array([0.0, 0.0, 0.0])
 
-    wlen, wpnt0, wpnt1 = _wrap(x0, x1, xpos, xmat, radius, wraptype, side)
+    wlen, wpnt0, wpnt1 = _wrap(x0, x1, xpos, xquat, radius, wraptype, side)
 
     _assert_eq(wlen, 0.0, "wlen")
     _assert_eq(wpnt0, np.array([0.1, 0.0, 0.0]), "wpnt0")