Repository: google-deepmind/mujoco_warp
Pull request #146: Adds raycast implementation and tests
Author: login: nvtw
Created: 2025-04-11T14:47:30Z | Merged: 2025-05-12T13:40:57Z

See https://github.com/orgs/google-deepmind/projects/5?pane=issue&itemId=104437279&issue=google-deepmind%7Cmujoco_warp%7C94


Diff:
diff --git a/mujoco_warp/__init__.py b/mujoco_warp/__init__.py
index 1ada61bcf..b462855ae 100644
--- a/mujoco_warp/__init__.py
+++ b/mujoco_warp/__init__.py
@@ -34,6 +34,7 @@
 from ._src.io import put_data as put_data
 from ._src.io import put_model as put_model
 from ._src.passive import passive as passive
+from ._src.ray import ray as ray
 from ._src.sensor import sensor_acc as sensor_acc
 from ._src.sensor import sensor_pos as sensor_pos
 from ._src.sensor import sensor_vel as sensor_vel
diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py
index 0a51f1e72..f0a1f8032 100644
--- a/mujoco_warp/_src/io.py
+++ b/mujoco_warp/_src/io.py
@@ -127,17 +127,18 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     tile_beg = tile_corners[i]
     tile_end = mjm.nv if i == len(tile_corners) - 1 else tile_corners[i + 1]
     tiles.setdefault(tile_end - tile_beg, []).append(tile_beg)
+
   qM_tiles = tuple(types.TileSet(adr=wp.array(tiles[sz], dtype=int), size=sz) for sz in sorted(tiles.keys()))
 
-  # subtree_mass is a precalculated arrya used in smooth
+  # subtree_mass is a precalculated array used in smooth
   subtree_mass = np.copy(mjm.body_mass)
   # TODO(team): should this be [mjm.nbody - 1, 0) ?
   for i in range(mjm.nbody - 1, -1, -1):
     subtree_mass[mjm.body_parentid[i]] += subtree_mass[i]
 
   # actuator_moment tiles are grouped by dof size and number of actuators
-  tree_id = mjm.dof_treeid[tile_corners]
-  num_trees = int(np.max(tree_id))
+  tree_id = np.arange(len(tile_corners), dtype=np.int32)
+  num_trees = int(np.max(tree_id)) if len(tree_id) > 0 else 0
   bodyid = []
   for i in range(mjm.nu):
     trntype = mjm.actuator_trntype[i]
@@ -250,7 +251,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     nxn_geom_pair.append((geom1, geom2))
     nxn_pairid.append(pairid)
 
-  return types.Model(
+  m = types.Model(
     nq=mjm.nq,
     nv=mjm.nv,
     nu=mjm.nu,
@@ -270,6 +271,8 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     nwrap=mjm.nwrap,
     nsensor=mjm.nsensor,
     nsensordata=mjm.nsensordata,
+    nmeshvert=mjm.nmeshvert,
+    nmeshface=mjm.nmeshface,
     nlsp=nlsp,
     npair=mjm.npair,
     opt=types.Option(
@@ -375,6 +378,8 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     geom_condim=wp.array(mjm.geom_condim, dtype=int),
     geom_bodyid=wp.array(mjm.geom_bodyid, dtype=int),
     geom_dataid=wp.array(mjm.geom_dataid, dtype=int),
+    geom_group=wp.array(mjm.geom_group, dtype=int),
+    geom_matid=wp.array(mjm.geom_matid, dtype=int),
     geom_priority=wp.array(mjm.geom_priority, dtype=int),
     geom_solmix=wp.array(mjm.geom_solmix, dtype=float),
     geom_solref=wp.array(mjm.geom_solref, dtype=wp.vec2),
@@ -387,6 +392,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     geom_friction=wp.array(mjm.geom_friction, dtype=wp.vec3),
     geom_margin=wp.array(mjm.geom_margin, dtype=float),
     geom_gap=wp.array(mjm.geom_gap, dtype=float),
+    geom_rgba=wp.array(mjm.geom_rgba, dtype=wp.vec4),
     site_bodyid=wp.array(mjm.site_bodyid, dtype=int),
     site_pos=wp.array(mjm.site_pos, dtype=wp.vec3),
     site_quat=wp.array(mjm.site_quat, dtype=wp.quat),
@@ -411,6 +417,8 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     mesh_vertadr=wp.array(mjm.mesh_vertadr, dtype=int),
     mesh_vertnum=wp.array(mjm.mesh_vertnum, dtype=int),
     mesh_vert=wp.array(mjm.mesh_vert, dtype=wp.vec3),
+    mesh_faceadr=wp.array(mjm.mesh_faceadr, dtype=int),
+    mesh_face=wp.array(mjm.mesh_face, dtype=wp.vec3i),
     eq_type=wp.array(mjm.eq_type, dtype=int),
     eq_obj1id=wp.array(mjm.eq_obj1id, dtype=int),
     eq_obj2id=wp.array(mjm.eq_obj2id, dtype=int),
@@ -517,8 +525,11 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
         mujoco.mjtSensor.mjSENS_FRAMEANGACC,
       ],
     ).any(),
+    mat_rgba=wp.array(mjm.mat_rgba, dtype=wp.vec4),
   )
 
+  return m
+
 
 def make_data(mjm: mujoco.MjModel, nworld: int = 1, nconmax: int = -1, njmax: int = -1) -> types.Data:
   # TODO(team): move to Model?
diff --git a/mujoco_warp/_src/ray.py b/mujoco_warp/_src/ray.py
new file mode 100644
index 000000000..eb7135268
--- /dev/null
+++ b/mujoco_warp/_src/ray.py
@@ -0,0 +1,740 @@
+# Copyright 2025 The Newton Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import warp as wp
+
+from .types import MJ_MINVAL
+from .types import Data
+from .types import GeomType
+from .types import Model
+from .types import vec6
+
+
+@wp.struct
+class DistanceWithId:
+  dist: wp.float32
+  geom_id: wp.int32
+
+
+@wp.func
+def _ray_quad(a: float, b: float, c: float) -> wp.vec2:
+  """Returns two solutions for quadratic: a*x^2 + 2*b*x + c = 0."""
+  det = b * b - a * c
+  det_2 = wp.sqrt(det)
+
+  x0 = (-b - det_2) / a
+  x1 = (-b + det_2) / a
+  x0 = wp.where((det < MJ_MINVAL) or (x0 < 0.0), wp.inf, x0)
+  x1 = wp.where((det < MJ_MINVAL) or (x1 < 0.0), wp.inf, x1)
+
+  return wp.vec2(x0, x1)
+
+
+@wp.func
+def _ray_plane(
+  # In:
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with a plane."""
+  x = -pnt[2] / vec[2]
+
+  valid = vec[2] <= -MJ_MINVAL  # z-vec pointing towards front face
+  valid = valid and x >= 0.0
+  # only within rendered rectangle
+  p_x = pnt[0] + x * vec[0]
+  p_y = pnt[1] + x * vec[1]
+  valid = valid and ((size[0] <= 0.0) or (wp.abs(p_x) <= size[0]))
+  valid = valid and ((size[1] <= 0.0) or (wp.abs(p_y) <= size[1]))
+
+  return_id = wp.where(valid, geom_id, -1)
+  return DistanceWithId(wp.where(valid, x, wp.inf), return_id)
+
+
+@wp.func
+def _ray_sphere(
+  # In:
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with a sphere."""
+  a = wp.dot(vec, vec)
+  b = wp.dot(vec, pnt)
+  c = wp.dot(pnt, pnt) - size[0] * size[0]
+
+  solutions = _ray_quad(a, b, c)
+  x0 = solutions[0]
+  x1 = solutions[1]
+  x = wp.where(wp.isinf(x0), x1, x0)
+
+  return_id = wp.where(wp.isinf(x), -1, geom_id)
+  return DistanceWithId(x, return_id)
+
+
+@wp.func
+def _ray_capsule(
+  # In:
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with a capsule."""
+
+  # cylinder round side: (x*lvec+lpnt)'*(x*lvec+lpnt) = size[0]*size[0]
+  # For a capsule, we only care about the x,y components when checking cylinder intersection
+  # since the z component is handled separately with the caps
+  vec_2d = wp.vec2(vec[0], vec[1])
+  pnt_2d = wp.vec2(pnt[0], pnt[1])
+  a = wp.dot(vec_2d, vec_2d)
+  b = wp.dot(pnt_2d, vec_2d)
+  c = wp.dot(pnt_2d, pnt_2d) - size[0] * size[0]
+
+  # solve a*x^2 + 2*b*x + c = 0
+  solutions = _ray_quad(a, b, c)
+  x0 = solutions[0]
+  x1 = solutions[1]
+  x = wp.where(wp.isinf(x0), x1, x0)
+
+  # make sure round solution is between flat sides
+  x = wp.where(wp.abs(pnt[2] + x * vec[2]) <= size[1], x, wp.inf)
+
+  # top cap
+  dif = pnt - wp.vec3(0.0, 0.0, size[1])
+  solutions = _ray_quad(wp.dot(vec, vec), wp.dot(vec, dif), wp.dot(dif, dif) - size[0] * size[0])
+  x0 = solutions[0]
+  x1 = solutions[1]
+  # accept only top half of sphere
+  x = wp.where((pnt[2] + x0 * vec[2] >= size[1]) and (x0 < x), x0, x)
+  x = wp.where((pnt[2] + x1 * vec[2] >= size[1]) and (x1 < x), x1, x)
+
+  # bottom cap
+  dif = pnt + wp.vec3(0.0, 0.0, size[1])
+  solutions = _ray_quad(wp.dot(vec, vec), wp.dot(vec, dif), wp.dot(dif, dif) - size[0] * size[0])
+  x0 = solutions[0]
+  x1 = solutions[1]
+  # accept only bottom half of sphere
+  x = wp.where((pnt[2] + x0 * vec[2] <= -size[1]) and (x0 < x), x0, x)
+  x = wp.where((pnt[2] + x1 * vec[2] <= -size[1]) and (x1 < x), x1, x)
+
+  return_id = wp.where(wp.isinf(x), -1, geom_id)
+  return DistanceWithId(x, return_id)
+
+
+@wp.func
+def _ray_ellipsoid(
+  # In:
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with an ellipsoid."""
+
+  # invert size^2
+  s = wp.vec3(1.0 / (size[0] * size[0]), 1.0 / (size[1] * size[1]), 1.0 / (size[2] * size[2]))
+
+  # (x*lvec+lpnt)' * diag(1/size^2) * (x*lvec+lpnt) = 1
+  svec = wp.vec3(s[0] * vec[0], s[1] * vec[1], s[2] * vec[2])
+  a = wp.dot(svec, vec)
+  b = wp.dot(svec, pnt)
+  c = wp.dot(wp.vec3(s[0] * pnt[0], s[1] * pnt[1], s[2] * pnt[2]), pnt) - 1.0
+
+  # solve a*x^2 + 2*b*x + c = 0
+  solutions = _ray_quad(a, b, c)
+  x0 = solutions[0]
+  x1 = solutions[1]
+  x = wp.where(wp.isinf(x0), x1, x0)
+
+  return_id = wp.where(wp.isinf(x), -1, geom_id)
+  return DistanceWithId(x, return_id)
+
+
+@wp.func
+def _ray_box(
+  # In:
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with a box."""
+
+  # Initialize with infinity
+  min_x = wp.inf
+
+  # Check all 6 faces of the box (2 faces per axis)
+  for i in range(wp.static(3)):
+    if vec[i] != 0.0:
+      # Get indices for the other two dimensions
+      j = (i + 1) % 3
+      k = (i + 2) % 3
+
+      for t in range(wp.static(2)):
+        s = wp.where(t == 0, size[i], -size[i])
+        sol = (s - pnt[i]) / vec[i]
+
+        pj = pnt[j] + sol * vec[j]
+        pk = pnt[k] + sol * vec[k]
+        if sol >= 0.0 and wp.abs(pj) <= size[j] and wp.abs(pk) <= size[k]:
+          min_x = wp.min(min_x, sol)
+
+  return_id = wp.where(wp.isinf(min_x), -1, geom_id)
+  return DistanceWithId(min_x, return_id)
+
+
+@wp.struct
+class Triangle:
+  """A struct representing a triangle with 3 vertices."""
+
+  v0: wp.vec3
+  v1: wp.vec3
+  v2: wp.vec3
+
+
+@wp.struct
+class Basis:
+  """A struct representing a basis with 2 vectors."""
+
+  b0: wp.vec3
+  b1: wp.vec3
+
+
+@wp.func
+def _ray_triangle(
+  # In:
+  triangle: Triangle,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  basis: Basis,
+) -> wp.float32:
+  """Returns the distance at which a ray intersects with a triangle."""
+  # dif = v[i] - lpnt
+  dif0 = triangle.v0 - pnt
+  dif1 = triangle.v1 - pnt
+  dif2 = triangle.v2 - pnt
+
+  # project difference vectors in normal plane
+  planar_0_0 = wp.dot(dif0, basis.b0)
+  planar_0_1 = wp.dot(dif0, basis.b1)
+  planar_1_0 = wp.dot(dif1, basis.b0)
+  planar_1_1 = wp.dot(dif1, basis.b1)
+  planar_2_0 = wp.dot(dif2, basis.b0)
+  planar_2_1 = wp.dot(dif2, basis.b1)
+
+  # reject if on the same side of any coordinate axis
+  if (
+    (planar_0_0 > 0.0 and planar_1_0 > 0.0 and planar_2_0 > 0.0)
+    or (planar_0_0 < 0.0 and planar_1_0 < 0.0 and planar_2_0 < 0.0)
+    or (planar_0_1 > 0.0 and planar_1_1 > 0.0 and planar_2_1 > 0.0)
+    or (planar_0_1 < 0.0 and planar_1_1 < 0.0 and planar_2_1 < 0.0)
+  ):
+    return wp.float32(wp.inf)
+
+  # determine if origin is inside planar projection of triangle
+  # A = (p0-p2, p1-p2), b = -p2, solve A*t = b
+  A00 = planar_0_0 - planar_2_0
+  A10 = planar_1_0 - planar_2_0
+  A01 = planar_0_1 - planar_2_1
+  A11 = planar_1_1 - planar_2_1
+
+  b0 = -planar_2_0
+  b1 = -planar_2_1
+
+  det = A00 * A11 - A10 * A01
+  if wp.abs(det) < MJ_MINVAL:
+    return wp.float32(wp.inf)
+
+  t0 = (A11 * b0 - A10 * b1) / det
+  t1 = (-A01 * b0 + A00 * b1) / det
+
+  # check if outside
+  if t0 < 0.0 or t1 < 0.0 or t0 + t1 > 1.0:
+    return wp.float32(wp.inf)
+
+  # intersect ray with plane of triangle
+  dif0 = triangle.v0 - triangle.v2  # v0-v2
+  dif1 = triangle.v1 - triangle.v2  # v1-v2
+  dif2 = pnt - triangle.v2  # lpnt-v2
+  nrm = wp.cross(dif0, dif1)  # normal to triangle plane
+  denom = wp.dot(vec, nrm)
+  if wp.abs(denom) < MJ_MINVAL:
+    return wp.float32(wp.inf)
+
+  dist = -wp.dot(dif2, nrm) / denom
+  return wp.where(dist >= 0.0, dist, wp.float32(wp.inf))
+
+
+@wp.func
+def _ray_mesh(
+  # Model:
+  nmeshface: int,
+  mesh_vertadr: wp.array(dtype=int),
+  mesh_vertnum: wp.array(dtype=int),
+  mesh_vert: wp.array(dtype=wp.vec3),
+  mesh_faceadr: wp.array(dtype=int),
+  mesh_face: wp.array(dtype=wp.vec3i),
+  # In:
+  geom_id: int,
+  data_id: int,
+  unused_size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+) -> DistanceWithId:
+  """Returns the distance and geom_id for ray mesh intersections."""
+
+  # Create basis vectors for the ray
+  basis = Basis()
+
+  # Compute orthogonal basis vectors
+  if wp.abs(vec[0]) < wp.abs(vec[1]):
+    if wp.abs(vec[0]) < wp.abs(vec[2]):
+      basis.b0 = wp.vec3(0.0, vec[2], -vec[1])
+    else:
+      basis.b0 = wp.vec3(vec[1], -vec[0], 0.0)
+  else:
+    if wp.abs(vec[1]) < wp.abs(vec[2]):
+      basis.b0 = wp.vec3(-vec[2], 0.0, vec[0])
+    else:
+      basis.b0 = wp.vec3(vec[1], -vec[0], 0.0)
+
+  # Normalize first basis vector
+  basis.b0 = wp.normalize(basis.b0)
+
+  # Compute second basis vector as cross product
+  basis.b1 = wp.cross(vec, basis.b0)
+  basis.b1 = wp.normalize(basis.b1)
+
+  min_dist = wp.float32(wp.inf)
+  hit_found = int(0)
+
+  # Get mesh vertex data range
+  vert_start = mesh_vertadr[data_id]
+
+  # Get mesh face and vertex data
+  face_start = mesh_faceadr[data_id]
+  face_end = wp.where(data_id + 1 < mesh_faceadr.shape[0], mesh_faceadr[data_id + 1], nmeshface)
+
+  # Iterate through all faces
+  for i in range(face_start, face_end):
+    # Get vertices for this face
+    v_idx = mesh_face[i]
+
+    # Create triangle struct
+    triangle = Triangle()
+    triangle.v0 = mesh_vert[vert_start + v_idx.x]
+    triangle.v1 = mesh_vert[vert_start + v_idx.y]
+    triangle.v2 = mesh_vert[vert_start + v_idx.z]
+
+    # Calculate intersection
+    dist = _ray_triangle(triangle, pnt, vec, basis)
+    if dist < min_dist:
+      min_dist = dist
+      hit_found = 1
+
+  # Return the geom_id if we found a hit, otherwise -1
+  return_id = wp.where(hit_found == 1, geom_id, -1)
+
+  return DistanceWithId(min_dist, return_id)
+
+
+@wp.func
+def _ray_map(
+  # In:
+  pos: wp.vec3,
+  mat: wp.mat33,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+) -> any:
+  """Maps ray to local geom frame coordinates.
+
+  Args:
+      pos: Position of geom frame
+      mat: Orientation of geom frame
+      pnt: Starting point of ray in world coordinates
+      vec: Direction of ray in world coordinates
+
+  Returns:
+      Tuple of (local_pnt, local_vec) mapped to geom frame
+  """
+  # Calculate difference vector
+  dif = pnt - pos
+
+  # Transform point and vector to local coordinates using transpose(mat)
+  local_pnt = wp.transpose(mat) @ dif
+  local_vec = wp.transpose(mat) @ vec
+
+  return local_pnt, local_vec
+
+
+@wp.func
+def _ray_geom(
+  # In:
+  pos: wp.vec3,  # Position of geom frame
+  mat: wp.mat33,  # Orientation of geom frame
+  size: wp.vec3,  # Size parameters of geom
+  pnt: wp.vec3,  # Starting point of ray in world coordinates
+  vec: wp.vec3,  # Direction of ray in world coordinates
+  geomtype: int,  # Type of geometry
+) -> float:
+  """Returns distance along ray to intersection with geom, or infinity if no intersection.
+
+  Matches MuJoCo's mju_rayGeom API for use with touch sensors.
+  Maps inputs to local coordinates before intersection testing.
+  """
+  # Map ray to local coordinates
+  local_pnt, local_vec = _ray_map(pos, mat, pnt, vec)
+
+  # Create DistanceWithId struct to reuse existing ray intersection functions
+  result = DistanceWithId(wp.inf, -1)
+
+  # Call appropriate intersection function based on geom type
+  if geomtype == int(GeomType.PLANE.value):
+    result = _ray_plane(size, local_pnt, local_vec, 0)
+  elif geomtype == int(GeomType.SPHERE.value):
+    result = _ray_sphere(size, local_pnt, local_vec, 0)
+  elif geomtype == int(GeomType.CAPSULE.value):
+    result = _ray_capsule(size, local_pnt, local_vec, 0)
+  elif geomtype == int(GeomType.ELLIPSOID.value):
+    result = _ray_ellipsoid(size, local_pnt, local_vec, 0)
+  elif geomtype == int(GeomType.BOX.value):
+    result = _ray_box(size, local_pnt, local_vec, 0)
+
+  return result.dist
+
+
+@wp.func
+def _ray_geom_with_mesh(
+  # Model:
+  nmeshface: int,
+  geom_type: wp.array(dtype=int),
+  geom_dataid: wp.array(dtype=int),
+  geom_size: wp.array(dtype=wp.vec3),
+  mesh_vertadr: wp.array(dtype=int),
+  mesh_vertnum: wp.array(dtype=int),
+  mesh_vert: wp.array(dtype=wp.vec3),
+  mesh_faceadr: wp.array(dtype=int),
+  mesh_face: wp.array(dtype=wp.vec3i),
+  # In:
+  geom_id: int,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+) -> DistanceWithId:
+  type = geom_type[geom_id]
+  size = geom_size[geom_id]
+
+  # TODO(team): static loop unrolling to remove unnecessary branching
+  if type == int(GeomType.PLANE.value):
+    return _ray_plane(size, pnt, vec, geom_id)
+  elif type == int(GeomType.SPHERE.value):
+    return _ray_sphere(size, pnt, vec, geom_id)
+  elif type == int(GeomType.CAPSULE.value):
+    return _ray_capsule(size, pnt, vec, geom_id)
+  elif type == int(GeomType.ELLIPSOID.value):
+    return _ray_ellipsoid(size, pnt, vec, geom_id)
+  elif type == int(GeomType.BOX.value):
+    return _ray_box(size, pnt, vec, geom_id)
+  elif type == int(GeomType.MESH.value):
+    data_id = geom_dataid[geom_id]
+    return _ray_mesh(
+      nmeshface,
+      mesh_vertadr,
+      mesh_vertnum,
+      mesh_vert,
+      mesh_faceadr,
+      mesh_face,
+      geom_id,
+      data_id,
+      size,
+      pnt,
+      vec,
+    )
+  return DistanceWithId(wp.inf, -1)
+
+
+@wp.struct
+class RayIntersection:
+  dist: wp.float32
+  geom_id: wp.int32
+
+
+snippet = """
+#if defined(__CUDA_ARCH__)
+    return blockDim.x;
+#else    
+    return 1;
+#endif
+    """
+
+
+@wp.func_native(snippet)
+def get_block_dim_x() -> int: ...
+
+
+@wp.func
+def _ray_all_geom(
+  # Model:
+  ngeom: int,
+  nmeshface: int,
+  body_weldid: wp.array(dtype=int),
+  geom_type: wp.array(dtype=int),
+  geom_bodyid: wp.array(dtype=int),
+  geom_dataid: wp.array(dtype=int),
+  geom_group: wp.array(dtype=int),
+  geom_matid: wp.array(dtype=int),
+  geom_size: wp.array(dtype=wp.vec3),
+  geom_rgba: wp.array(dtype=wp.vec4),
+  mesh_vertadr: wp.array(dtype=int),
+  mesh_vertnum: wp.array(dtype=int),
+  mesh_vert: wp.array(dtype=wp.vec3),
+  mesh_faceadr: wp.array(dtype=int),
+  mesh_face: wp.array(dtype=wp.vec3i),
+  mat_rgba: wp.array(dtype=wp.vec4),
+  # Data in:
+  geom_xpos_in: wp.array2d(dtype=wp.vec3),
+  geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  # In:
+  worldid: int,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geomgroup: vec6,
+  has_geomgroup: bool,
+  flg_static: bool,
+  bodyexclude: int,
+  tid: int,
+) -> RayIntersection:
+  num_threads = get_block_dim_x()
+
+  min_val = wp.float32(wp.inf)
+  min_idx = wp.int32(-1)
+
+  upper = ((ngeom + num_threads - 1) // num_threads) * num_threads
+  for geom_id in range(tid, upper, num_threads):
+    if geom_id < ngeom:
+      # Apply all filters combined into a single boolean
+      body_id = geom_bodyid[geom_id]
+
+      # Start with True and apply each filter condition
+      geom_filter = True
+      # Body exclusion filter
+      geom_filter = geom_filter and (body_id != bodyexclude)
+
+      # Static geom filter
+      geom_filter = geom_filter and (flg_static or body_weldid[body_id] != 0)
+
+      # Geom group filter
+      if has_geomgroup:
+        group = geom_group[geom_id]
+        # Clip group index to [0, 5] (mjNGROUP-1)
+        group = wp.max(0, wp.min(5, group))
+        geom_filter = geom_filter and (geomgroup[group] != 0)
+
+      # RGBA filter
+      matid = geom_matid[geom_id]
+      geom_alpha = geom_rgba[geom_id][3]
+      mat_alpha = wp.float32(0.0)
+      if matid != -1:
+        mat_alpha = mat_rgba[matid][3]
+
+      # Geom is visible if either:
+      # 1. No material and non-zero geom alpha, or
+      # 2. Has material and non-zero material alpha
+      geom_visible = (matid == -1 and geom_alpha != 0.0) or (matid != -1 and mat_alpha != 0.0)
+      geom_filter = geom_filter and geom_visible
+
+      if not geom_filter:
+        cur_dist = wp.float32(wp.inf)
+      else:
+        # Get ray in local coordinates
+        pos = geom_xpos_in[worldid, geom_id]
+        rot = geom_xmat_in[worldid, geom_id]
+        local_pnt = wp.transpose(rot) @ (pnt - pos)
+        local_vec = wp.transpose(rot) @ vec
+
+        # Calculate intersection distance
+        result = _ray_geom_with_mesh(
+          nmeshface,
+          geom_type,
+          geom_dataid,
+          geom_size,
+          mesh_vertadr,
+          mesh_vertnum,
+          mesh_vert,
+          mesh_faceadr,
+          mesh_face,
+          geom_id,
+          local_pnt,
+          local_vec,
+        )
+        cur_dist = result.dist
+    else:
+      cur_dist = wp.float32(wp.inf)
+
+    t = wp.tile(cur_dist)
+    local_min_idx = wp.tile_argmin(t)
+    local_min_val = t[local_min_idx[0]]
+
+    id_tile = wp.tile(geom_id)
+
+    if local_min_val < min_val:
+      min_val = local_min_val
+      min_idx = id_tile[local_min_idx[0]]
+
+  min_val = wp.where(min_val == wp.inf, wp.float32(-1.0), min_val)
+
+  return RayIntersection(min_val, min_idx)
+
+
+# One thread block/tile per ray query
+@wp.kernel
+def _ray_all_geom_kernel(
+  # Model:
+  ngeom: int,
+  nmeshface: int,
+  body_weldid: wp.array(dtype=int),
+  geom_type: wp.array(dtype=int),
+  geom_bodyid: wp.array(dtype=int),
+  geom_dataid: wp.array(dtype=int),
+  geom_group: wp.array(dtype=int),
+  geom_matid: wp.array(dtype=int),
+  geom_size: wp.array(dtype=wp.vec3),
+  geom_rgba: wp.array(dtype=wp.vec4),
+  mesh_vertadr: wp.array(dtype=int),
+  mesh_vertnum: wp.array(dtype=int),
+  mesh_vert: wp.array(dtype=wp.vec3),
+  mesh_faceadr: wp.array(dtype=int),
+  mesh_face: wp.array(dtype=wp.vec3i),
+  mat_rgba: wp.array(dtype=wp.vec4),
+  # Data in:
+  geom_xpos_in: wp.array2d(dtype=wp.vec3),
+  geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  # In:
+  pnt: wp.array(dtype=wp.vec3),
+  vec: wp.array(dtype=wp.vec3),
+  geomgroup: vec6,
+  has_geomgroup: bool,
+  flg_static: bool,
+  bodyexclude: int,
+  # Out:
+  dist_out: wp.array(dtype=float, ndim=2),
+  closest_hit_geom_id_out: wp.array(dtype=int, ndim=2),
+):
+  worldid, rayid, tid = wp.tid()
+  intersection = _ray_all_geom(
+    ngeom,
+    nmeshface,
+    body_weldid,
+    geom_type,
+    geom_bodyid,
+    geom_dataid,
+    geom_group,
+    geom_matid,
+    geom_size,
+    geom_rgba,
+    mesh_vertadr,
+    mesh_vertnum,
+    mesh_vert,
+    mesh_faceadr,
+    mesh_face,
+    mat_rgba,
+    geom_xpos_in,
+    geom_xmat_in,
+    worldid,
+    pnt[rayid],
+    vec[rayid],
+    geomgroup,
+    has_geomgroup,
+    flg_static,
+    bodyexclude,
+    tid,
+  )
+
+  # Write intersection results to output arrays
+  dist_out[worldid, rayid] = intersection.dist
+  closest_hit_geom_id_out[worldid, rayid] = intersection.geom_id
+
+
+def ray(
+  m: Model,
+  d: Data,
+  pnt: wp.array(dtype=wp.vec3),
+  vec: wp.array(dtype=wp.vec3),
+  geomgroup: vec6 = None,
+  flg_static: bool = True,
+  bodyexclude: int = -1,
+) -> tuple[wp.array, wp.array]:
+  """Returns the distance at which rays intersect with primitive geoms.
+
+  Args:
+      m: MuJoCo model
+      d: MuJoCo data
+      pnt: ray origin points
+      vec: ray directions
+      geomgroup: group inclusion/exclusion mask (6,), or all zeros to ignore
+      flg_static: if True, allows rays to intersect with static geoms
+      bodyexclude: ignore geoms on specified body id (-1 to disable)
+
+  Returns:
+      dist: distances from ray origins to geom surfaces
+      geom_id: IDs of intersected geoms (-1 if none)
+  """
+  nrays = pnt.shape[0]
+  dist = wp.zeros((d.nworld, nrays), dtype=float)
+  closest_hit_geom_id = wp.zeros((d.nworld, nrays), dtype=int)
+  num_threads = 64
+
+  # Create default geomgroup if None is provided
+  has_geomgroup = geomgroup is not None
+  if geomgroup is None:
+    geomgroup = vec6(0, 0, 0, 0, 0, 0)
+
+  wp.launch_tiled(
+    _ray_all_geom_kernel,
+    dim=(d.nworld, nrays),
+    inputs=[
+      m.ngeom,
+      m.nmeshface,
+      m.body_weldid,
+      m.geom_type,
+      m.geom_bodyid,
+      m.geom_dataid,
+      m.geom_group,
+      m.geom_matid,
+      m.geom_size,
+      m.geom_rgba,
+      m.mesh_vertadr,
+      m.mesh_vertnum,
+      m.mesh_vert,
+      m.mesh_faceadr,
+      m.mesh_face,
+      m.mat_rgba,
+      d.geom_xpos,
+      d.geom_xmat,
+      pnt,
+      vec,
+      geomgroup,
+      has_geomgroup,
+      flg_static,
+      bodyexclude,
+      dist,
+      closest_hit_geom_id,
+    ],
+    block_dim=num_threads,
+  )
+  return dist, closest_hit_geom_id
diff --git a/mujoco_warp/_src/ray_test.py b/mujoco_warp/_src/ray_test.py
new file mode 100644
index 000000000..4bddd6763
--- /dev/null
+++ b/mujoco_warp/_src/ray_test.py
@@ -0,0 +1,285 @@
+# Copyright 2025 The Newton Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for ray functions."""
+
+import mujoco
+import numpy as np
+import warp as wp
+from absl.testing import absltest
+
+import mujoco_warp as mjwarp
+
+from . import test_util
+from .types import vec6
+
+# tolerance for difference between MuJoCo and MJX ray calculations - mostly
+# due to float precision
+_TOLERANCE = 5e-5
+
+
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10  # avoid test noise
+  err_msg = f"mismatch: {name}"
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
+
+
+class RayTest(absltest.TestCase):
+  def test_ray_nothing(self):
+    """Tests that ray returns -1 when nothing is hit."""
+    mjm, mjd, m, d = test_util.fixture("ray.xml")
+
+    pnt = wp.array([wp.vec3(12.146, 1.865, 3.895)], dtype=wp.vec3)
+    vec = wp.array([wp.vec3(0.0, 0.0, -1.0)], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray(m, d, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]  # Extract from [[-1]]
+    dist_np = dist.numpy()[0][0]  # Extract from [[-1.]]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_plane(self):
+    """Tests ray<>plane matches MuJoCo."""
+    mjm, mjd, m, d = test_util.fixture("ray.xml")
+
+    # looking down at a slight angle
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray(m, d, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 0, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(mjm, mjd, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking on wrong side of plane
+    pnt = wp.array([wp.vec3(0.0, 0.0, -0.5)], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray(m, d, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_sphere(self):
+    """Tests ray<>sphere matches MuJoCo."""
+    mjm, mjd, m, d = test_util.fixture("ray.xml")
+
+    # looking down at sphere at a slight angle
+    pnt = wp.array([wp.vec3(0.0, 0.0, 1.6)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray(m, d, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 1, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(mjm, mjd, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+  def test_ray_capsule(self):
+    """Tests ray<>capsule matches MuJoCo."""
+    mjm, mjd, m, d = test_util.fixture("ray.xml")
+
+    # looking down at capsule at a slight angle
+    pnt = wp.array([wp.vec3(0.5, 1.0, 1.6)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray(m, d, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 2, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(mjm, mjd, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking up at capsule from below
+    pnt = wp.array([wp.vec3(-0.5, 1.0, 0.05)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, 1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray(m, d, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 2, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(mjm, mjd, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking at cylinder of capsule from the side
+    pnt = wp.array([wp.vec3(0.0, 1.0, 0.75)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(1.0, 0.0, 0.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray(m, d, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 2, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(mjm, mjd, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+  def test_ray_box(self):
+    """Tests ray<>box matches MuJoCo."""
+    mjm, mjd, m, d = test_util.fixture("ray.xml")
+
+    # looking down at box at a slight angle
+    pnt = wp.array([wp.vec3(1.0, 0.0, 1.6)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray(m, d, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 3, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(mjm, mjd, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking up at box from below
+    pnt = wp.array([wp.vec3(1.0, 0.0, 0.05)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, 1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray(m, d, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 3, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(mjm, mjd, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+  def test_ray_mesh(self):
+    """Tests ray<>mesh matches MuJoCo."""
+    mjm, mjd, m, d = test_util.fixture("ray.xml")
+
+    # look at the tetrahedron
+    pnt = wp.array([wp.vec3(2.0, 2.0, 2.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(-1.0, -1.0, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray(m, d, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 4, "geom_id")
+
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(mjm, mjd, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist-tetrahedron")
+
+    # look away from the dodecahedron
+    pnt = wp.array([wp.vec3(4.0, 2.0, 2.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(2.0, 1.0, 1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray(m, d, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+
+    # look at the dodecahedron
+    pnt = wp.array([wp.vec3(4.0, 2.0, 2.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(-2.0, -1.0, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray(m, d, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 5, "geom_id")
+
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(mjm, mjd, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist-dodecahedron")
+
+  def test_ray_geomgroup(self):
+    """Tests ray geomgroup filter."""
+    mjm, mjd, m, d = test_util.fixture("ray.xml")
+
+    # hits plane with geom_group[0] = 1
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    geomgroup = vec6(1, 0, 0, 0, 0, 0)
+    dist, geomid = mjwarp.ray(m, d, pnt, vec, geomgroup=geomgroup)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 0, "geom_id")
+
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(mjm, mjd, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # nothing hit with geom_group[0] = 0
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    geomgroup = vec6(0, 0, 0, 0, 0, 0)
+    dist, geomid = mjwarp.ray(m, d, pnt, vec, geomgroup=geomgroup)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_flg_static(self):
+    """Tests ray flg_static filter."""
+    mjm, mjd, m, d = test_util.fixture("ray.xml")
+
+    # nothing hit with flg_static = False
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray(m, d, pnt, vec, flg_static=False)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_bodyexclude(self):
+    """Tests ray bodyexclude filter."""
+    mjm, mjd, m, d = test_util.fixture("ray.xml")
+
+    # nothing hit with bodyexclude = 0 (world body)
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray(m, d, pnt, vec, bodyexclude=0)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_invisible(self):
+    """Tests ray doesn't hit transparent geoms."""
+    mjm, mjd, m, d = test_util.fixture("ray.xml")
+
+    # nothing hit with transparent geoms
+    m.geom_rgba = wp.array([wp.vec4(0.0, 0.0, 0.0, 0.0)], dtype=wp.vec4)
+    mujoco.mj_forward(mjm, mjd)
+
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray(m, d, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/mujoco_warp/_src/types.py b/mujoco_warp/_src/types.py
index 4482f25e2..3147d630b 100644
--- a/mujoco_warp/_src/types.py
+++ b/mujoco_warp/_src/types.py
@@ -355,6 +355,10 @@ class vec5f(wp.types.vector(length=5, dtype=float)):
   pass
 
 
+class vec6f(wp.types.vector(length=6, dtype=float)):
+  pass
+
+
 class vec10f(wp.types.vector(length=10, dtype=float)):
   pass
 
@@ -364,6 +368,7 @@ class vec11f(wp.types.vector(length=11, dtype=float)):
 
 
 vec5 = vec5f
+vec6 = vec6f
 vec10 = vec10f
 vec11 = vec11f
 array2df = wp.array2d(dtype=float)
@@ -583,8 +588,10 @@ class Model:
     nwrap: number of wrap objects in all tendon paths        ()
     nsensor: number of sensors                               ()
     nsensordata: number of elements in sensor data vector    ()
-    npair: number of predefined geom pairs                   ()
+    nmeshvert: number of vertices for all meshes             ()
+    nmeshface: number of faces for all meshes                ()
     nlsp: number of step sizes for parallel linsearch        ()
+    npair: number of predefined geom pairs                   ()
     opt: physics options
     stat: model statistics
     qpos0: qpos values at default pose                       (nq,)
@@ -659,6 +666,8 @@ class Model:
     geom_condim: contact dimensionality (1, 3, 4, 6)         (ngeom,)
     geom_bodyid: id of geom's body                           (ngeom,)
     geom_dataid: id of geom's mesh/hfield; -1: none          (ngeom,)
+    geom_group: geom group inclusion/exclusion mask          (ngeom,)
+    geom_matid: material id for rendering                    (ngeom,)
     geom_priority: geom contact priority                     (ngeom,)
     geom_solmix: mixing coef for solref/imp in geom pair     (ngeom,)
     geom_solref: constraint solver reference: contact        (ngeom, mjNREF)
@@ -671,6 +680,7 @@ class Model:
     geom_friction: friction for (slide, spin, roll)          (ngeom, 3)
     geom_margin: detect contact if dist<margin               (ngeom,)
     geom_gap: include in solver if dist<margin-gap           (ngeom,)
+    geom_rgba: rgba when material is omitted                  (ngeom, 4)
     site_bodyid: id of site's body                           (nsite,)
     site_pos: local position offset rel. to body             (nsite, 3)
     site_quat: local orientation offset rel. to body         (nsite, 4)
@@ -695,6 +705,8 @@ class Model:
     mesh_vertadr: first vertex address                       (nmesh,)
     mesh_vertnum: number of vertices                         (nmesh,)
     mesh_vert: vertex positions for all meshes               (nmeshvert, 3)
+    mesh_faceadr: first face address                         (nmesh,)
+    mesh_face: face indices for all meshes                   (nface, 3)
     eq_type: constraint type (mjtEq)                         (neq,)
     eq_obj1id: id of object 1                                (neq,)
     eq_obj2id: id of object 2                                (neq,)
@@ -776,6 +788,7 @@ class Model:
     sensor_subtree_vel: evaluate subtree_vel
     sensor_rne_postconstraint: evaluate rne_postconstraint
     mocap_bodyid: id of body for mocap                       (nmocap,)
+    mat_rgba: rgba                                           (nmat, 4)
   """
 
   nq: int
@@ -797,8 +810,10 @@ class Model:
   nwrap: int
   nsensor: int
   nsensordata: int
-  npair: int
+  nmeshvert: int
+  nmeshface: int
   nlsp: int  # warp only
+  npair: int
   opt: Option
   stat: Statistic
   qpos0: wp.array(dtype=float)
@@ -872,6 +887,8 @@ class Model:
   geom_condim: wp.array(dtype=int)
   geom_bodyid: wp.array(dtype=int)
   geom_dataid: wp.array(dtype=int)
+  geom_group: wp.array(dtype=int)
+  geom_matid: wp.array(dtype=int)
   geom_priority: wp.array(dtype=int)
   geom_solmix: wp.array(dtype=float)
   geom_solref: wp.array(dtype=wp.vec2)
@@ -884,6 +901,7 @@ class Model:
   geom_friction: wp.array(dtype=wp.vec3)
   geom_margin: wp.array(dtype=float)
   geom_gap: wp.array(dtype=float)
+  geom_rgba: wp.array(dtype=wp.vec4)
   site_bodyid: wp.array(dtype=int)
   site_pos: wp.array(dtype=wp.vec3)
   site_quat: wp.array(dtype=wp.quat)
@@ -908,6 +926,8 @@ class Model:
   mesh_vertadr: wp.array(dtype=int)
   mesh_vertnum: wp.array(dtype=int)
   mesh_vert: wp.array(dtype=wp.vec3)
+  mesh_faceadr: wp.array(dtype=int)
+  mesh_face: wp.array(dtype=wp.vec3i)
   eq_type: wp.array(dtype=int)
   eq_obj1id: wp.array(dtype=int)
   eq_obj2id: wp.array(dtype=int)
@@ -989,6 +1009,7 @@ class Model:
   sensor_subtree_vel: bool  # warp only
   sensor_rne_postconstraint: bool  # warp only
   mocap_bodyid: wp.array(dtype=int)  # warp only
+  mat_rgba: wp.array(dtype=wp.vec4)
 
 
 @dataclasses.dataclass
diff --git a/mujoco_warp/test_data/meshes/dodecahedron.stl b/mujoco_warp/test_data/meshes/dodecahedron.stl
new file mode 100644
index 000000000..1a3f9f69a
Binary files /dev/null and b/mujoco_warp/test_data/meshes/dodecahedron.stl differ
diff --git a/mujoco_warp/test_data/meshes/tetrahedron.stl b/mujoco_warp/test_data/meshes/tetrahedron.stl
new file mode 100644
index 000000000..68fed6b85
Binary files /dev/null and b/mujoco_warp/test_data/meshes/tetrahedron.stl differ
diff --git a/mujoco_warp/test_data/ray.xml b/mujoco_warp/test_data/ray.xml
new file mode 100644
index 000000000..bbf42b1e4
--- /dev/null
+++ b/mujoco_warp/test_data/ray.xml
@@ -0,0 +1,18 @@
+<mujoco model="ray">
+  <asset>
+      <mesh name="tetrahedron" file="meshes/tetrahedron.stl" scale="0.4 0.4 0.4" />
+      <mesh name="dodecahedron" file="meshes/dodecahedron.stl" scale="0.04 0.04 0.04" />
+      <texture builtin="checker" height="100" name="texplane" rgb1="0 0 0" rgb2="0.8 0.8 0.8" type="2d" width="100"/>
+      <material name="MatPlane" reflectance="0.5" shininess="1" specular="1" texrepeat="60 60" texture="texplane"/>
+  </asset>
+  <worldbody>
+    <light cutoff="100" diffuse="1 1 1" dir="-0 0 -1.3" directional="true" exponent="1" pos="0 0 1.3" specular=".1 .1 .1"/>
+
+    <geom name="plane" pos="0 0 0" quat="1 0 0 0" size="4 4 4" type="plane" rgba="0.1 0.1 0.1 1"/>
+    <geom name="sphere" pos="0 0 1" quat="1 0 0 0" size="0.5" type="sphere" rgba="1 0 0 1"/>
+    <geom name="capsule" pos="0 1 1" quat="0 0.3826834 0 0.9238795 " size="0.25 0.5" type="capsule" rgba="0 1 0 1"/>
+    <geom name="box" pos="1 0 1" quat="0 0.3826834 0 0.9238795" size="0.5 0.25 0.3" type="box" rgba="0 0 1 1"/>
+    <geom name="mesh" pos="1 1 1" quat="0 0 0.3826834 0.9238795" type="mesh" mesh="tetrahedron" rgba="1 1 0 1"/>
+    <geom name="mesh2" pos="2 1 1" type="mesh" mesh="dodecahedron" rgba="1 0 1 1"/>
+  </worldbody>
+</mujoco>


Reviews:
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T10:26:32Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T10:30:17Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T10:32:13Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T10:32:24Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T10:33:19Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T10:34:33Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T10:35:02Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T10:35:50Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T10:36:33Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T10:36:51Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T10:41:00Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T10:41:54Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T10:42:30Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T10:51:09Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T10:58:12Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T11:00:51Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T11:07:58Z
user: login: mmacklin
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-22T08:37:48Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T11:14:20Z
user: login: thowell
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T11:18:52Z
user: login: thowell
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-12T13:39:05Z

Review comments:
diff_hunk: @@ -363,9 +370,12 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   m.geom_aabb = wp.array(mjm.geom_aabb, dtype=wp.vec3, ndim=3)
   m.geom_rbound = wp.array(mjm.geom_rbound, dtype=wp.float32, ndim=1)
   m.geom_dataid = wp.array(mjm.geom_dataid, dtype=wp.int32, ndim=1)
+  m.geom_group = wp.array(mjm.geom_group, dtype=wp.int32, ndim=1)
   m.mesh_vertadr = wp.array(mjm.mesh_vertadr, dtype=wp.int32, ndim=1)
   m.mesh_vertnum = wp.array(mjm.mesh_vertnum, dtype=wp.int32, ndim=1)
   m.mesh_vert = wp.array(mjm.mesh_vert, dtype=wp.vec3, ndim=1)
+  m.mesh_faceadr = wp.array(mjm.mesh_faceadr, dtype=wp.int32, ndim=1)
+  m.mesh_face = wp.array(mjm.mesh_face, dtype=wp.int32, ndim=2)
path: mujoco_warp/_src/io.py
user: login: thowell
body: since `mesh_face` is `(nmeshface, 3)` this array could be `wp.array(mjm.mesh_face, dtype=wp.vec3i, ndim=1)`
created_at: 2025-04-14T10:26:32Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,286 @@
+# Copyright 2025 The Newton Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for ray functions."""
+
+import mujoco
+import numpy as np
+import warp as wp
+from absl.testing import absltest
+
+import mujoco_warp as mjwarp
+
+from . import test_util
+from .types import vec6
+
+# tolerance for difference between MuJoCo and MJX ray calculations - mostly
+# due to float precision
+_TOLERANCE = 5e-5
+
+
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10  # avoid test noise
+  err_msg = f"mismatch: {name}"
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
+
+
+class RayTest(absltest.TestCase):
+  def test_ray_nothing(self):
+    """Tests that ray returns -1 when nothing is hit."""
+    m, d, _, _ = test_util.fixture("ray.xml")
path: mujoco_warp/_src/ray_test.py
user: login: thowell
body: following the convention of other tests:

`mjm, mjd, m, d = ...`

then it is not necessary to call `mujoco.mj_forward`, `mjwarp.put_model`, and `mjwarp.put_data` below since these functions are called by `test_util.fixture`.
created_at: 2025-04-14T10:30:17Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,286 @@
+# Copyright 2025 The Newton Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for ray functions."""
+
+import mujoco
+import numpy as np
+import warp as wp
+from absl.testing import absltest
+
+import mujoco_warp as mjwarp
+
+from . import test_util
+from .types import vec6
+
+# tolerance for difference between MuJoCo and MJX ray calculations - mostly
+# due to float precision
+_TOLERANCE = 5e-5
+
+
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10  # avoid test noise
+  err_msg = f"mismatch: {name}"
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
+
+
+class RayTest(absltest.TestCase):
+  def test_ray_nothing(self):
+    """Tests that ray returns -1 when nothing is hit."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    pnt = wp.array([wp.vec3(12.146, 1.865, 3.895)], dtype=wp.vec3)
+    vec = wp.array([wp.vec3(0.0, 0.0, -1.0)], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]  # Extract from [[-1]]
+    dist_np = dist.numpy()[0][0]  # Extract from [[-1.]]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_plane(self):
+    """Tests ray<>plane matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
path: mujoco_warp/_src/ray_test.py
user: login: thowell
body: `mjm, mjd, m, d = ...`
created_at: 2025-04-14T10:32:13Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,286 @@
+# Copyright 2025 The Newton Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for ray functions."""
+
+import mujoco
+import numpy as np
+import warp as wp
+from absl.testing import absltest
+
+import mujoco_warp as mjwarp
+
+from . import test_util
+from .types import vec6
+
+# tolerance for difference between MuJoCo and MJX ray calculations - mostly
+# due to float precision
+_TOLERANCE = 5e-5
+
+
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10  # avoid test noise
+  err_msg = f"mismatch: {name}"
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
+
+
+class RayTest(absltest.TestCase):
+  def test_ray_nothing(self):
+    """Tests that ray returns -1 when nothing is hit."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    pnt = wp.array([wp.vec3(12.146, 1.865, 3.895)], dtype=wp.vec3)
+    vec = wp.array([wp.vec3(0.0, 0.0, -1.0)], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]  # Extract from [[-1]]
+    dist_np = dist.numpy()[0][0]  # Extract from [[-1.]]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_plane(self):
+    """Tests ray<>plane matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at a slight angle
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 0, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking on wrong side of plane
+    pnt = wp.array([wp.vec3(0.0, 0.0, -0.5)], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_sphere(self):
+    """Tests ray<>sphere matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
path: mujoco_warp/_src/ray_test.py
user: login: thowell
body: `mjm, mjd, m, d = ...`
created_at: 2025-04-14T10:32:24Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,286 @@
+# Copyright 2025 The Newton Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for ray functions."""
+
+import mujoco
+import numpy as np
+import warp as wp
+from absl.testing import absltest
+
+import mujoco_warp as mjwarp
+
+from . import test_util
+from .types import vec6
+
+# tolerance for difference between MuJoCo and MJX ray calculations - mostly
+# due to float precision
+_TOLERANCE = 5e-5
+
+
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10  # avoid test noise
+  err_msg = f"mismatch: {name}"
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
+
+
+class RayTest(absltest.TestCase):
+  def test_ray_nothing(self):
+    """Tests that ray returns -1 when nothing is hit."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    pnt = wp.array([wp.vec3(12.146, 1.865, 3.895)], dtype=wp.vec3)
+    vec = wp.array([wp.vec3(0.0, 0.0, -1.0)], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]  # Extract from [[-1]]
+    dist_np = dist.numpy()[0][0]  # Extract from [[-1.]]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_plane(self):
+    """Tests ray<>plane matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at a slight angle
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 0, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking on wrong side of plane
+    pnt = wp.array([wp.vec3(0.0, 0.0, -0.5)], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_sphere(self):
+    """Tests ray<>sphere matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at sphere at a slight angle
+    pnt = wp.array([wp.vec3(0.0, 0.0, 1.6)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 1, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+  def test_ray_capsule(self):
+    """Tests ray<>capsule matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
path: mujoco_warp/_src/ray_test.py
user: login: thowell
body: `mjm, mjd, m, d = ...`
created_at: 2025-04-14T10:33:19Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,286 @@
+# Copyright 2025 The Newton Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for ray functions."""
+
+import mujoco
+import numpy as np
+import warp as wp
+from absl.testing import absltest
+
+import mujoco_warp as mjwarp
+
+from . import test_util
+from .types import vec6
+
+# tolerance for difference between MuJoCo and MJX ray calculations - mostly
+# due to float precision
+_TOLERANCE = 5e-5
+
+
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10  # avoid test noise
+  err_msg = f"mismatch: {name}"
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
+
+
+class RayTest(absltest.TestCase):
+  def test_ray_nothing(self):
+    """Tests that ray returns -1 when nothing is hit."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    pnt = wp.array([wp.vec3(12.146, 1.865, 3.895)], dtype=wp.vec3)
+    vec = wp.array([wp.vec3(0.0, 0.0, -1.0)], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]  # Extract from [[-1]]
+    dist_np = dist.numpy()[0][0]  # Extract from [[-1.]]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_plane(self):
+    """Tests ray<>plane matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at a slight angle
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 0, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking on wrong side of plane
+    pnt = wp.array([wp.vec3(0.0, 0.0, -0.5)], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_sphere(self):
+    """Tests ray<>sphere matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at sphere at a slight angle
+    pnt = wp.array([wp.vec3(0.0, 0.0, 1.6)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 1, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+  def test_ray_capsule(self):
+    """Tests ray<>capsule matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at capsule at a slight angle
+    pnt = wp.array([wp.vec3(0.5, 1.0, 1.6)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 2, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking up at capsule from below
+    pnt = wp.array([wp.vec3(-0.5, 1.0, 0.05)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, 1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 2, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking at cylinder of capsule from the side
+    pnt = wp.array([wp.vec3(0.0, 1.0, 0.75)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(1.0, 0.0, 0.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 2, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+  def test_ray_box(self):
+    """Tests ray<>box matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
path: mujoco_warp/_src/ray_test.py
user: login: thowell
body: `mjm, mjd, m, d = ...`
created_at: 2025-04-14T10:34:33Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,286 @@
+# Copyright 2025 The Newton Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for ray functions."""
+
+import mujoco
+import numpy as np
+import warp as wp
+from absl.testing import absltest
+
+import mujoco_warp as mjwarp
+
+from . import test_util
+from .types import vec6
+
+# tolerance for difference between MuJoCo and MJX ray calculations - mostly
+# due to float precision
+_TOLERANCE = 5e-5
+
+
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10  # avoid test noise
+  err_msg = f"mismatch: {name}"
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
+
+
+class RayTest(absltest.TestCase):
+  def test_ray_nothing(self):
+    """Tests that ray returns -1 when nothing is hit."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    pnt = wp.array([wp.vec3(12.146, 1.865, 3.895)], dtype=wp.vec3)
+    vec = wp.array([wp.vec3(0.0, 0.0, -1.0)], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]  # Extract from [[-1]]
+    dist_np = dist.numpy()[0][0]  # Extract from [[-1.]]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_plane(self):
+    """Tests ray<>plane matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at a slight angle
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 0, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking on wrong side of plane
+    pnt = wp.array([wp.vec3(0.0, 0.0, -0.5)], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_sphere(self):
+    """Tests ray<>sphere matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at sphere at a slight angle
+    pnt = wp.array([wp.vec3(0.0, 0.0, 1.6)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 1, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+  def test_ray_capsule(self):
+    """Tests ray<>capsule matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at capsule at a slight angle
+    pnt = wp.array([wp.vec3(0.5, 1.0, 1.6)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 2, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking up at capsule from below
+    pnt = wp.array([wp.vec3(-0.5, 1.0, 0.05)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, 1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 2, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking at cylinder of capsule from the side
+    pnt = wp.array([wp.vec3(0.0, 1.0, 0.75)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(1.0, 0.0, 0.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 2, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+  def test_ray_box(self):
+    """Tests ray<>box matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at box at a slight angle
+    pnt = wp.array([wp.vec3(1.0, 0.0, 1.6)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 3, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking up at box from below
+    pnt = wp.array([wp.vec3(1.0, 0.0, 0.05)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, 1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 3, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+  def test_ray_mesh(self):
+    """Tests ray<>mesh matches MuJoCo."""
path: mujoco_warp/_src/ray_test.py
user: login: thowell
body: `mjm, mjd, m, d = ...`
created_at: 2025-04-14T10:35:02Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,286 @@
+# Copyright 2025 The Newton Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for ray functions."""
+
+import mujoco
+import numpy as np
+import warp as wp
+from absl.testing import absltest
+
+import mujoco_warp as mjwarp
+
+from . import test_util
+from .types import vec6
+
+# tolerance for difference between MuJoCo and MJX ray calculations - mostly
+# due to float precision
+_TOLERANCE = 5e-5
+
+
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10  # avoid test noise
+  err_msg = f"mismatch: {name}"
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
+
+
+class RayTest(absltest.TestCase):
+  def test_ray_nothing(self):
+    """Tests that ray returns -1 when nothing is hit."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    pnt = wp.array([wp.vec3(12.146, 1.865, 3.895)], dtype=wp.vec3)
+    vec = wp.array([wp.vec3(0.0, 0.0, -1.0)], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]  # Extract from [[-1]]
+    dist_np = dist.numpy()[0][0]  # Extract from [[-1.]]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_plane(self):
+    """Tests ray<>plane matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at a slight angle
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 0, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking on wrong side of plane
+    pnt = wp.array([wp.vec3(0.0, 0.0, -0.5)], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_sphere(self):
+    """Tests ray<>sphere matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at sphere at a slight angle
+    pnt = wp.array([wp.vec3(0.0, 0.0, 1.6)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 1, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+  def test_ray_capsule(self):
+    """Tests ray<>capsule matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at capsule at a slight angle
+    pnt = wp.array([wp.vec3(0.5, 1.0, 1.6)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 2, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking up at capsule from below
+    pnt = wp.array([wp.vec3(-0.5, 1.0, 0.05)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, 1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 2, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking at cylinder of capsule from the side
+    pnt = wp.array([wp.vec3(0.0, 1.0, 0.75)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(1.0, 0.0, 0.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 2, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+  def test_ray_box(self):
+    """Tests ray<>box matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at box at a slight angle
+    pnt = wp.array([wp.vec3(1.0, 0.0, 1.6)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 3, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking up at box from below
+    pnt = wp.array([wp.vec3(1.0, 0.0, 0.05)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, 1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 3, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+  def test_ray_mesh(self):
+    """Tests ray<>mesh matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # look at the tetrahedron
+    pnt = wp.array([wp.vec3(2.0, 2.0, 2.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(-1.0, -1.0, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 4, "geom_id")
+
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist-tetrahedron")
+
+    # look away from the dodecahedron
+    pnt = wp.array([wp.vec3(4.0, 2.0, 2.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(2.0, 1.0, 1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+
+    # look at the dodecahedron
+    pnt = wp.array([wp.vec3(4.0, 2.0, 2.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(-2.0, -1.0, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 5, "geom_id")
+
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist-dodecahedron")
+
+  def test_ray_geomgroup(self):
+    """Tests ray geomgroup filter."""
+    m, d, _, _ = test_util.fixture("ray.xml")
path: mujoco_warp/_src/ray_test.py
user: login: thowell
body: `mjm, mjd, m, d = ...`
created_at: 2025-04-14T10:35:50Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,286 @@
+# Copyright 2025 The Newton Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for ray functions."""
+
+import mujoco
+import numpy as np
+import warp as wp
+from absl.testing import absltest
+
+import mujoco_warp as mjwarp
+
+from . import test_util
+from .types import vec6
+
+# tolerance for difference between MuJoCo and MJX ray calculations - mostly
+# due to float precision
+_TOLERANCE = 5e-5
+
+
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10  # avoid test noise
+  err_msg = f"mismatch: {name}"
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
+
+
+class RayTest(absltest.TestCase):
+  def test_ray_nothing(self):
+    """Tests that ray returns -1 when nothing is hit."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    pnt = wp.array([wp.vec3(12.146, 1.865, 3.895)], dtype=wp.vec3)
+    vec = wp.array([wp.vec3(0.0, 0.0, -1.0)], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]  # Extract from [[-1]]
+    dist_np = dist.numpy()[0][0]  # Extract from [[-1.]]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_plane(self):
+    """Tests ray<>plane matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at a slight angle
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 0, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking on wrong side of plane
+    pnt = wp.array([wp.vec3(0.0, 0.0, -0.5)], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_sphere(self):
+    """Tests ray<>sphere matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at sphere at a slight angle
+    pnt = wp.array([wp.vec3(0.0, 0.0, 1.6)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 1, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+  def test_ray_capsule(self):
+    """Tests ray<>capsule matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at capsule at a slight angle
+    pnt = wp.array([wp.vec3(0.5, 1.0, 1.6)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 2, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking up at capsule from below
+    pnt = wp.array([wp.vec3(-0.5, 1.0, 0.05)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, 1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 2, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking at cylinder of capsule from the side
+    pnt = wp.array([wp.vec3(0.0, 1.0, 0.75)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(1.0, 0.0, 0.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 2, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+  def test_ray_box(self):
+    """Tests ray<>box matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at box at a slight angle
+    pnt = wp.array([wp.vec3(1.0, 0.0, 1.6)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 3, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking up at box from below
+    pnt = wp.array([wp.vec3(1.0, 0.0, 0.05)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, 1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 3, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+  def test_ray_mesh(self):
+    """Tests ray<>mesh matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # look at the tetrahedron
+    pnt = wp.array([wp.vec3(2.0, 2.0, 2.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(-1.0, -1.0, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 4, "geom_id")
+
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist-tetrahedron")
+
+    # look away from the dodecahedron
+    pnt = wp.array([wp.vec3(4.0, 2.0, 2.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(2.0, 1.0, 1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+
+    # look at the dodecahedron
+    pnt = wp.array([wp.vec3(4.0, 2.0, 2.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(-2.0, -1.0, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 5, "geom_id")
+
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist-dodecahedron")
+
+  def test_ray_geomgroup(self):
+    """Tests ray geomgroup filter."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # hits plane with geom_group[0] = 1
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    geomgroup = vec6(1, 0, 0, 0, 0, 0)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec, geomgroup=geomgroup)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 0, "geom_id")
+
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # nothing hit with geom_group[0] = 0
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    geomgroup = vec6(0, 0, 0, 0, 0, 0)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec, geomgroup=geomgroup)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_flg_static(self):
+    """Tests ray flg_static filter."""
+    m, d, _, _ = test_util.fixture("ray.xml")
path: mujoco_warp/_src/ray_test.py
user: login: thowell
body: `mjm, mjd, m, d = ...`
created_at: 2025-04-14T10:36:33Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,286 @@
+# Copyright 2025 The Newton Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for ray functions."""
+
+import mujoco
+import numpy as np
+import warp as wp
+from absl.testing import absltest
+
+import mujoco_warp as mjwarp
+
+from . import test_util
+from .types import vec6
+
+# tolerance for difference between MuJoCo and MJX ray calculations - mostly
+# due to float precision
+_TOLERANCE = 5e-5
+
+
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10  # avoid test noise
+  err_msg = f"mismatch: {name}"
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
+
+
+class RayTest(absltest.TestCase):
+  def test_ray_nothing(self):
+    """Tests that ray returns -1 when nothing is hit."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    pnt = wp.array([wp.vec3(12.146, 1.865, 3.895)], dtype=wp.vec3)
+    vec = wp.array([wp.vec3(0.0, 0.0, -1.0)], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]  # Extract from [[-1]]
+    dist_np = dist.numpy()[0][0]  # Extract from [[-1.]]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_plane(self):
+    """Tests ray<>plane matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at a slight angle
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 0, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking on wrong side of plane
+    pnt = wp.array([wp.vec3(0.0, 0.0, -0.5)], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_sphere(self):
+    """Tests ray<>sphere matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at sphere at a slight angle
+    pnt = wp.array([wp.vec3(0.0, 0.0, 1.6)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 1, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+  def test_ray_capsule(self):
+    """Tests ray<>capsule matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at capsule at a slight angle
+    pnt = wp.array([wp.vec3(0.5, 1.0, 1.6)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 2, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking up at capsule from below
+    pnt = wp.array([wp.vec3(-0.5, 1.0, 0.05)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, 1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 2, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking at cylinder of capsule from the side
+    pnt = wp.array([wp.vec3(0.0, 1.0, 0.75)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(1.0, 0.0, 0.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 2, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+  def test_ray_box(self):
+    """Tests ray<>box matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at box at a slight angle
+    pnt = wp.array([wp.vec3(1.0, 0.0, 1.6)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 3, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking up at box from below
+    pnt = wp.array([wp.vec3(1.0, 0.0, 0.05)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, 1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 3, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+  def test_ray_mesh(self):
+    """Tests ray<>mesh matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # look at the tetrahedron
+    pnt = wp.array([wp.vec3(2.0, 2.0, 2.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(-1.0, -1.0, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 4, "geom_id")
+
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist-tetrahedron")
+
+    # look away from the dodecahedron
+    pnt = wp.array([wp.vec3(4.0, 2.0, 2.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(2.0, 1.0, 1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+
+    # look at the dodecahedron
+    pnt = wp.array([wp.vec3(4.0, 2.0, 2.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(-2.0, -1.0, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 5, "geom_id")
+
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist-dodecahedron")
+
+  def test_ray_geomgroup(self):
+    """Tests ray geomgroup filter."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # hits plane with geom_group[0] = 1
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    geomgroup = vec6(1, 0, 0, 0, 0, 0)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec, geomgroup=geomgroup)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 0, "geom_id")
+
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # nothing hit with geom_group[0] = 0
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    geomgroup = vec6(0, 0, 0, 0, 0, 0)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec, geomgroup=geomgroup)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_flg_static(self):
+    """Tests ray flg_static filter."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # nothing hit with flg_static = False
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec, flg_static=False)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_bodyexclude(self):
+    """Tests ray bodyexclude filter."""
+    m, d, _, _ = test_util.fixture("ray.xml")
path: mujoco_warp/_src/ray_test.py
user: login: thowell
body: `mjm, mjd, m, d = ...`
created_at: 2025-04-14T10:36:51Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,286 @@
+# Copyright 2025 The Newton Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for ray functions."""
+
+import mujoco
+import numpy as np
+import warp as wp
+from absl.testing import absltest
+
+import mujoco_warp as mjwarp
+
+from . import test_util
+from .types import vec6
+
+# tolerance for difference between MuJoCo and MJX ray calculations - mostly
+# due to float precision
+_TOLERANCE = 5e-5
+
+
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10  # avoid test noise
+  err_msg = f"mismatch: {name}"
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
+
+
+class RayTest(absltest.TestCase):
+  def test_ray_nothing(self):
+    """Tests that ray returns -1 when nothing is hit."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    pnt = wp.array([wp.vec3(12.146, 1.865, 3.895)], dtype=wp.vec3)
+    vec = wp.array([wp.vec3(0.0, 0.0, -1.0)], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]  # Extract from [[-1]]
+    dist_np = dist.numpy()[0][0]  # Extract from [[-1.]]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_plane(self):
+    """Tests ray<>plane matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at a slight angle
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 0, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking on wrong side of plane
+    pnt = wp.array([wp.vec3(0.0, 0.0, -0.5)], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_sphere(self):
+    """Tests ray<>sphere matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at sphere at a slight angle
+    pnt = wp.array([wp.vec3(0.0, 0.0, 1.6)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 1, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+  def test_ray_capsule(self):
+    """Tests ray<>capsule matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at capsule at a slight angle
+    pnt = wp.array([wp.vec3(0.5, 1.0, 1.6)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 2, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking up at capsule from below
+    pnt = wp.array([wp.vec3(-0.5, 1.0, 0.05)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, 1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 2, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking at cylinder of capsule from the side
+    pnt = wp.array([wp.vec3(0.0, 1.0, 0.75)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(1.0, 0.0, 0.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 2, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+  def test_ray_box(self):
+    """Tests ray<>box matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # looking down at box at a slight angle
+    pnt = wp.array([wp.vec3(1.0, 0.0, 1.6)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 3, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # looking up at box from below
+    pnt = wp.array([wp.vec3(1.0, 0.0, 0.05)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.0, 0.05, 1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 3, "geom_id")
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+  def test_ray_mesh(self):
+    """Tests ray<>mesh matches MuJoCo."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # look at the tetrahedron
+    pnt = wp.array([wp.vec3(2.0, 2.0, 2.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(-1.0, -1.0, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 4, "geom_id")
+
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist-tetrahedron")
+
+    # look away from the dodecahedron
+    pnt = wp.array([wp.vec3(4.0, 2.0, 2.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(2.0, 1.0, 1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+
+    # look at the dodecahedron
+    pnt = wp.array([wp.vec3(4.0, 2.0, 2.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(-2.0, -1.0, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 5, "geom_id")
+
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist-dodecahedron")
+
+  def test_ray_geomgroup(self):
+    """Tests ray geomgroup filter."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # hits plane with geom_group[0] = 1
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    geomgroup = vec6(1, 0, 0, 0, 0, 0)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec, geomgroup=geomgroup)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, 0, "geom_id")
+
+    pnt_np, vec_np = pnt.numpy()[0], vec.numpy()[0]
+    unused = np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt_np, vec_np, None, 1, -1, unused)
+    _assert_eq(dist_np, mj_dist, "dist")
+
+    # nothing hit with geom_group[0] = 0
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    geomgroup = vec6(0, 0, 0, 0, 0, 0)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec, geomgroup=geomgroup)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_flg_static(self):
+    """Tests ray flg_static filter."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # nothing hit with flg_static = False
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec, flg_static=False)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
+  def test_ray_bodyexclude(self):
+    """Tests ray bodyexclude filter."""
+    m, d, _, _ = test_util.fixture("ray.xml")
+    mujoco.mj_forward(m, d)
+    mx, dx = mjwarp.put_model(m), mjwarp.put_data(m, d)
+
+    # nothing hit with bodyexclude = 0 (world body)
+    pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3)
+    vec = wp.array([wp.normalize(wp.vec3(0.1, 0.2, -1.0))], dtype=wp.vec3)
+    dist, geomid = mjwarp.ray_geom(mx, dx, pnt, vec, bodyexclude=0)
+    wp.synchronize()
+    geomid_np = geomid.numpy()[0][0]
+    dist_np = dist.numpy()[0][0]
+    _assert_eq(geomid_np, -1, "geom_id")
+    _assert_eq(dist_np, -1, "dist")
+
path: mujoco_warp/_src/ray_test.py
user: login: thowell
body: please add a test for transparent geoms. mjx reference: https://github.com/google-deepmind/mujoco/blob/8e7457d7c6c1cbda42ec39d2297d4f90b3396040/mjx/mujoco/mjx/_src/ray_test.py#L238
created_at: 2025-04-14T10:41:00Z
author_association: COLLABORATOR
diff_hunk: @@ -576,6 +583,7 @@ class Model:
     geom_condim: contact dimensionality (1, 3, 4, 6)         (ngeom,)
     geom_bodyid: id of geom's body                           (ngeom,)
     geom_dataid: id of geom's mesh/hfield; -1: none          (ngeom,)
+    geom_group: geom group inclusion/exclusion mask          (ngeom, 1)
path: mujoco_warp/_src/types.py
user: login: thowell
body: `(ngeom,)`
created_at: 2025-04-14T10:41:53Z
author_association: COLLABORATOR
diff_hunk: @@ -788,6 +801,8 @@ class Model:
   mesh_vertadr: wp.array(dtype=wp.int32, ndim=1)
   mesh_vertnum: wp.array(dtype=wp.int32, ndim=1)
   mesh_vert: wp.array(dtype=wp.vec3, ndim=1)
+  mesh_faceadr: wp.array(dtype=wp.int32, ndim=1)
+  mesh_face: wp.array(dtype=wp.int32, ndim=2)
path: mujoco_warp/_src/types.py
user: login: thowell
body: `wp.array(dtype=wp.vec3i, ndim=1)`
created_at: 2025-04-14T10:42:30Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,588 @@
+# Copyright 2025 The Newton Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import warp as wp
+
+from .types import MJ_MINVAL
+from .types import Data
+from .types import GeomType
+from .types import Model
+from .types import vec6
+
+
+@wp.struct
+class DistanceWithId:
+  dist: wp.float32
+  geom_id: wp.int32
+
+
+@wp.func
+def _ray_quad(a: float, b: float, c: float) -> wp.vec2:
+  """Returns two solutions for quadratic: a*x^2 + 2*b*x + c = 0."""
+  det = b * b - a * c
+  det_2 = wp.sqrt(det)
+
+  x0 = (-b - det_2) / a
+  x1 = (-b + det_2) / a
+  x0 = wp.where((det < MJ_MINVAL) or (x0 < 0.0), wp.inf, x0)
+  x1 = wp.where((det < MJ_MINVAL) or (x1 < 0.0), wp.inf, x1)
+
+  return wp.vec2(x0, x1)
+
+
+@wp.func
+def _ray_plane(
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with a plane."""
+  x = -pnt[2] / vec[2]
+
+  valid = vec[2] <= -MJ_MINVAL  # z-vec pointing towards front face
+  valid = valid and x >= 0.0
+  # only within rendered rectangle
+  p_x = pnt[0] + x * vec[0]
+  p_y = pnt[1] + x * vec[1]
+  valid = valid and ((size[0] <= 0.0) or (wp.abs(p_x) <= size[0]))
+  valid = valid and ((size[1] <= 0.0) or (wp.abs(p_y) <= size[1]))
+
+  return_id = wp.where(valid, geom_id, -1)
+  return DistanceWithId(wp.where(valid, x, wp.inf), return_id)
+
+
+@wp.func
+def _ray_sphere(
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with a sphere."""
+  a = wp.dot(vec, vec)
+  b = wp.dot(vec, pnt)
+  c = wp.dot(pnt, pnt) - size[0] * size[0]
+
+  solutions = _ray_quad(a, b, c)
+  x0 = solutions[0]
+  x1 = solutions[1]
+  x = wp.where(wp.isinf(x0), x1, x0)
+
+  return_id = wp.where(wp.isinf(x), -1, geom_id)
+  return DistanceWithId(x, return_id)
+
+
+@wp.func
+def _ray_capsule(
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with a capsule."""
+
+  # cylinder round side: (x*lvec+lpnt)'*(x*lvec+lpnt) = size[0]*size[0]
+  # For a capsule, we only care about the x,y components when checking cylinder intersection
+  # since the z component is handled separately with the caps
+  a = wp.dot(wp.vec2(vec[0], vec[1]), wp.vec2(vec[0], vec[1]))
path: mujoco_warp/_src/ray.py
user: login: thowell
body: would it make sense to create variables for `wp.vec2([vec[0], vec[1])` and `wp.vec2(pnt[0], pnt[1])` and utilize these variables to compute, `a`, `b`, and `c`?

created_at: 2025-04-14T10:51:09Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,588 @@
+# Copyright 2025 The Newton Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import warp as wp
+
+from .types import MJ_MINVAL
+from .types import Data
+from .types import GeomType
+from .types import Model
+from .types import vec6
+
+
+@wp.struct
+class DistanceWithId:
+  dist: wp.float32
+  geom_id: wp.int32
+
+
+@wp.func
+def _ray_quad(a: float, b: float, c: float) -> wp.vec2:
+  """Returns two solutions for quadratic: a*x^2 + 2*b*x + c = 0."""
+  det = b * b - a * c
+  det_2 = wp.sqrt(det)
+
+  x0 = (-b - det_2) / a
+  x1 = (-b + det_2) / a
+  x0 = wp.where((det < MJ_MINVAL) or (x0 < 0.0), wp.inf, x0)
+  x1 = wp.where((det < MJ_MINVAL) or (x1 < 0.0), wp.inf, x1)
+
+  return wp.vec2(x0, x1)
+
+
+@wp.func
+def _ray_plane(
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with a plane."""
+  x = -pnt[2] / vec[2]
+
+  valid = vec[2] <= -MJ_MINVAL  # z-vec pointing towards front face
+  valid = valid and x >= 0.0
+  # only within rendered rectangle
+  p_x = pnt[0] + x * vec[0]
+  p_y = pnt[1] + x * vec[1]
+  valid = valid and ((size[0] <= 0.0) or (wp.abs(p_x) <= size[0]))
+  valid = valid and ((size[1] <= 0.0) or (wp.abs(p_y) <= size[1]))
+
+  return_id = wp.where(valid, geom_id, -1)
+  return DistanceWithId(wp.where(valid, x, wp.inf), return_id)
+
+
+@wp.func
+def _ray_sphere(
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with a sphere."""
+  a = wp.dot(vec, vec)
+  b = wp.dot(vec, pnt)
+  c = wp.dot(pnt, pnt) - size[0] * size[0]
+
+  solutions = _ray_quad(a, b, c)
+  x0 = solutions[0]
+  x1 = solutions[1]
+  x = wp.where(wp.isinf(x0), x1, x0)
+
+  return_id = wp.where(wp.isinf(x), -1, geom_id)
+  return DistanceWithId(x, return_id)
+
+
+@wp.func
+def _ray_capsule(
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with a capsule."""
+
+  # cylinder round side: (x*lvec+lpnt)'*(x*lvec+lpnt) = size[0]*size[0]
+  # For a capsule, we only care about the x,y components when checking cylinder intersection
+  # since the z component is handled separately with the caps
+  a = wp.dot(wp.vec2(vec[0], vec[1]), wp.vec2(vec[0], vec[1]))
+  b = wp.dot(wp.vec2(pnt[0], pnt[1]), wp.vec2(vec[0], vec[1]))
+  c = wp.dot(wp.vec2(pnt[0], pnt[1]), wp.vec2(pnt[0], pnt[1])) - size[0] * size[0]
+
+  # solve a*x^2 + 2*b*x + c = 0
+  solutions = _ray_quad(a, b, c)
+  x0 = solutions[0]
+  x1 = solutions[1]
+  x = wp.where(wp.isinf(x0), x1, x0)
+
+  # make sure round solution is between flat sides
+  x = wp.where(wp.abs(pnt[2] + x * vec[2]) <= size[1], x, wp.inf)
+
+  # top cap
+  dif = pnt - wp.vec3(0.0, 0.0, size[1])
+  solutions = _ray_quad(
+    wp.dot(vec, vec), wp.dot(vec, dif), wp.dot(dif, dif) - size[0] * size[0]
+  )
+  x0 = solutions[0]
+  x1 = solutions[1]
+  # accept only top half of sphere
+  x = wp.where((pnt[2] + x0 * vec[2] >= size[1]) and (x0 < x), x0, x)
+  x = wp.where((pnt[2] + x1 * vec[2] >= size[1]) and (x1 < x), x1, x)
+
+  # bottom cap
+  dif = pnt + wp.vec3(0.0, 0.0, size[1])
+  solutions = _ray_quad(
+    wp.dot(vec, vec), wp.dot(vec, dif), wp.dot(dif, dif) - size[0] * size[0]
+  )
+  x0 = solutions[0]
+  x1 = solutions[1]
+  # accept only bottom half of sphere
+  x = wp.where((pnt[2] + x0 * vec[2] <= -size[1]) and (x0 < x), x0, x)
+  x = wp.where((pnt[2] + x1 * vec[2] <= -size[1]) and (x1 < x), x1, x)
+
+  return_id = wp.where(wp.isinf(x), -1, geom_id)
+  return DistanceWithId(x, return_id)
+
+
+@wp.func
+def _ray_ellipsoid(
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with an ellipsoid."""
+
+  # invert size^2
+  s = wp.vec3(
+    1.0 / (size[0] * size[0]), 1.0 / (size[1] * size[1]), 1.0 / (size[2] * size[2])
+  )
+
+  # (x*lvec+lpnt)' * diag(1/size^2) * (x*lvec+lpnt) = 1
+  svec = wp.vec3(s[0] * vec[0], s[1] * vec[1], s[2] * vec[2])
+  a = wp.dot(svec, vec)
+  b = wp.dot(svec, pnt)
+  c = wp.dot(wp.vec3(s[0] * pnt[0], s[1] * pnt[1], s[2] * pnt[2]), pnt) - 1.0
+
+  # solve a*x^2 + 2*b*x + c = 0
+  solutions = _ray_quad(a, b, c)
+  x0 = solutions[0]
+  x1 = solutions[1]
+  x = wp.where(wp.isinf(x0), x1, x0)
+
+  return_id = wp.where(wp.isinf(x), -1, geom_id)
+  return DistanceWithId(x, return_id)
+
+
+@wp.func
+def _ray_box(
path: mujoco_warp/_src/ray.py
user: login: thowell
body: this function can probably follow the mujoco reference more closely
https://github.com/google-deepmind/mujoco/blob/f317bd17c3494b954982de7bcf14030978d5013b/src/engine/engine_ray.c#L387
created_at: 2025-04-14T10:58:12Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,588 @@
+# Copyright 2025 The Newton Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import warp as wp
+
+from .types import MJ_MINVAL
+from .types import Data
+from .types import GeomType
+from .types import Model
+from .types import vec6
+
+
+@wp.struct
+class DistanceWithId:
+  dist: wp.float32
+  geom_id: wp.int32
+
+
+@wp.func
+def _ray_quad(a: float, b: float, c: float) -> wp.vec2:
+  """Returns two solutions for quadratic: a*x^2 + 2*b*x + c = 0."""
+  det = b * b - a * c
+  det_2 = wp.sqrt(det)
+
+  x0 = (-b - det_2) / a
+  x1 = (-b + det_2) / a
+  x0 = wp.where((det < MJ_MINVAL) or (x0 < 0.0), wp.inf, x0)
+  x1 = wp.where((det < MJ_MINVAL) or (x1 < 0.0), wp.inf, x1)
+
+  return wp.vec2(x0, x1)
+
+
+@wp.func
+def _ray_plane(
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with a plane."""
+  x = -pnt[2] / vec[2]
+
+  valid = vec[2] <= -MJ_MINVAL  # z-vec pointing towards front face
+  valid = valid and x >= 0.0
+  # only within rendered rectangle
+  p_x = pnt[0] + x * vec[0]
+  p_y = pnt[1] + x * vec[1]
+  valid = valid and ((size[0] <= 0.0) or (wp.abs(p_x) <= size[0]))
+  valid = valid and ((size[1] <= 0.0) or (wp.abs(p_y) <= size[1]))
+
+  return_id = wp.where(valid, geom_id, -1)
+  return DistanceWithId(wp.where(valid, x, wp.inf), return_id)
+
+
+@wp.func
+def _ray_sphere(
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with a sphere."""
+  a = wp.dot(vec, vec)
+  b = wp.dot(vec, pnt)
+  c = wp.dot(pnt, pnt) - size[0] * size[0]
+
+  solutions = _ray_quad(a, b, c)
+  x0 = solutions[0]
+  x1 = solutions[1]
+  x = wp.where(wp.isinf(x0), x1, x0)
+
+  return_id = wp.where(wp.isinf(x), -1, geom_id)
+  return DistanceWithId(x, return_id)
+
+
+@wp.func
+def _ray_capsule(
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with a capsule."""
+
+  # cylinder round side: (x*lvec+lpnt)'*(x*lvec+lpnt) = size[0]*size[0]
+  # For a capsule, we only care about the x,y components when checking cylinder intersection
+  # since the z component is handled separately with the caps
+  a = wp.dot(wp.vec2(vec[0], vec[1]), wp.vec2(vec[0], vec[1]))
+  b = wp.dot(wp.vec2(pnt[0], pnt[1]), wp.vec2(vec[0], vec[1]))
+  c = wp.dot(wp.vec2(pnt[0], pnt[1]), wp.vec2(pnt[0], pnt[1])) - size[0] * size[0]
+
+  # solve a*x^2 + 2*b*x + c = 0
+  solutions = _ray_quad(a, b, c)
+  x0 = solutions[0]
+  x1 = solutions[1]
+  x = wp.where(wp.isinf(x0), x1, x0)
+
+  # make sure round solution is between flat sides
+  x = wp.where(wp.abs(pnt[2] + x * vec[2]) <= size[1], x, wp.inf)
+
+  # top cap
+  dif = pnt - wp.vec3(0.0, 0.0, size[1])
+  solutions = _ray_quad(
+    wp.dot(vec, vec), wp.dot(vec, dif), wp.dot(dif, dif) - size[0] * size[0]
+  )
+  x0 = solutions[0]
+  x1 = solutions[1]
+  # accept only top half of sphere
+  x = wp.where((pnt[2] + x0 * vec[2] >= size[1]) and (x0 < x), x0, x)
+  x = wp.where((pnt[2] + x1 * vec[2] >= size[1]) and (x1 < x), x1, x)
+
+  # bottom cap
+  dif = pnt + wp.vec3(0.0, 0.0, size[1])
+  solutions = _ray_quad(
+    wp.dot(vec, vec), wp.dot(vec, dif), wp.dot(dif, dif) - size[0] * size[0]
+  )
+  x0 = solutions[0]
+  x1 = solutions[1]
+  # accept only bottom half of sphere
+  x = wp.where((pnt[2] + x0 * vec[2] <= -size[1]) and (x0 < x), x0, x)
+  x = wp.where((pnt[2] + x1 * vec[2] <= -size[1]) and (x1 < x), x1, x)
+
+  return_id = wp.where(wp.isinf(x), -1, geom_id)
+  return DistanceWithId(x, return_id)
+
+
+@wp.func
+def _ray_ellipsoid(
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with an ellipsoid."""
+
+  # invert size^2
+  s = wp.vec3(
+    1.0 / (size[0] * size[0]), 1.0 / (size[1] * size[1]), 1.0 / (size[2] * size[2])
+  )
+
+  # (x*lvec+lpnt)' * diag(1/size^2) * (x*lvec+lpnt) = 1
+  svec = wp.vec3(s[0] * vec[0], s[1] * vec[1], s[2] * vec[2])
+  a = wp.dot(svec, vec)
+  b = wp.dot(svec, pnt)
+  c = wp.dot(wp.vec3(s[0] * pnt[0], s[1] * pnt[1], s[2] * pnt[2]), pnt) - 1.0
+
+  # solve a*x^2 + 2*b*x + c = 0
+  solutions = _ray_quad(a, b, c)
+  x0 = solutions[0]
+  x1 = solutions[1]
+  x = wp.where(wp.isinf(x0), x1, x0)
+
+  return_id = wp.where(wp.isinf(x), -1, geom_id)
+  return DistanceWithId(x, return_id)
+
+
+@wp.func
+def _ray_box(
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with a box."""
+
+  # Initialize with infinity
+  min_x = wp.inf
+
+  # Check all 6 faces of the box
+  # X faces (i=0)
+  if vec[0] != 0.0:
+    # +X face
+    x_pos = (size[0] - pnt[0]) / vec[0]
+    p1 = pnt[1] + x_pos * vec[1]
+    p2 = pnt[2] + x_pos * vec[2]
+    if x_pos >= 0.0 and wp.abs(p1) <= size[1] and wp.abs(p2) <= size[2]:
+      min_x = wp.min(min_x, x_pos)
+
+    # -X face
+    x_neg = (-size[0] - pnt[0]) / vec[0]
+    p1 = pnt[1] + x_neg * vec[1]
+    p2 = pnt[2] + x_neg * vec[2]
+    if x_neg >= 0.0 and wp.abs(p1) <= size[1] and wp.abs(p2) <= size[2]:
+      min_x = wp.min(min_x, x_neg)
+
+  # Y faces (i=1)
+  if vec[1] != 0.0:
+    # +Y face
+    y_pos = (size[1] - pnt[1]) / vec[1]
+    p0 = pnt[0] + y_pos * vec[0]
+    p2 = pnt[2] + y_pos * vec[2]
+    if y_pos >= 0.0 and wp.abs(p0) <= size[0] and wp.abs(p2) <= size[2]:
+      min_x = wp.min(min_x, y_pos)
+
+    # -Y face
+    y_neg = (-size[1] - pnt[1]) / vec[1]
+    p0 = pnt[0] + y_neg * vec[0]
+    p2 = pnt[2] + y_neg * vec[2]
+    if y_neg >= 0.0 and wp.abs(p0) <= size[0] and wp.abs(p2) <= size[2]:
+      min_x = wp.min(min_x, y_neg)
+
+  # Z faces (i=2)
+  if vec[2] != 0.0:
+    # +Z face
+    z_pos = (size[2] - pnt[2]) / vec[2]
+    p0 = pnt[0] + z_pos * vec[0]
+    p1 = pnt[1] + z_pos * vec[1]
+    if z_pos >= 0.0 and wp.abs(p0) <= size[0] and wp.abs(p1) <= size[1]:
+      min_x = wp.min(min_x, z_pos)
+
+    # -Z face
+    z_neg = (-size[2] - pnt[2]) / vec[2]
+    p0 = pnt[0] + z_neg * vec[0]
+    p1 = pnt[1] + z_neg * vec[1]
+    if z_neg >= 0.0 and wp.abs(p0) <= size[0] and wp.abs(p1) <= size[1]:
+      min_x = wp.min(min_x, z_neg)
+
+  return_id = wp.where(wp.isinf(min_x), -1, geom_id)
+  return DistanceWithId(min_x, return_id)
+
+
+@wp.struct
+class Triangle:
+  """A struct representing a triangle with 3 vertices."""
+
+  v0: wp.vec3
+  v1: wp.vec3
+  v2: wp.vec3
+
+
+@wp.struct
+class Basis:
+  """A struct representing a basis with 2 vectors."""
+
+  b0: wp.vec3
+  b1: wp.vec3
+
+
+@wp.func
+def _ray_triangle(
+  triangle: Triangle,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  basis: Basis,
+) -> wp.float32:
+  """Returns the distance at which a ray intersects with a triangle."""
+  # dif = v[i] - lpnt
+  dif0 = triangle.v0 - pnt
+  dif1 = triangle.v1 - pnt
+  dif2 = triangle.v2 - pnt
+
+  # project difference vectors in normal plane
+  planar_0_0 = wp.dot(dif0, basis.b0)
+  planar_0_1 = wp.dot(dif0, basis.b1)
+  planar_1_0 = wp.dot(dif1, basis.b0)
+  planar_1_1 = wp.dot(dif1, basis.b1)
+  planar_2_0 = wp.dot(dif2, basis.b0)
+  planar_2_1 = wp.dot(dif2, basis.b1)
+
+  # reject if on the same side of any coordinate axis
+  if (
+    (planar_0_0 > 0.0 and planar_1_0 > 0.0 and planar_2_0 > 0.0)
+    or (planar_0_0 < 0.0 and planar_1_0 < 0.0 and planar_2_0 < 0.0)
+    or (planar_0_1 > 0.0 and planar_1_1 > 0.0 and planar_2_1 > 0.0)
+    or (planar_0_1 < 0.0 and planar_1_1 < 0.0 and planar_2_1 < 0.0)
+  ):
+    return wp.float32(wp.inf)
+
+  # determine if origin is inside planar projection of triangle
+  # A = (p0-p2, p1-p2), b = -p2, solve A*t = b
+  A00 = planar_0_0 - planar_2_0
+  A10 = planar_1_0 - planar_2_0
+  A01 = planar_0_1 - planar_2_1
+  A11 = planar_1_1 - planar_2_1
+
+  b0 = -planar_2_0
+  b1 = -planar_2_1
+
+  det = A00 * A11 - A10 * A01
+  if wp.abs(det) < MJ_MINVAL:
+    return wp.float32(wp.inf)
+
+  t0 = (A11 * b0 - A10 * b1) / det
+  t1 = (-A01 * b0 + A00 * b1) / det
+
+  # check if outside
+  if t0 < 0.0 or t1 < 0.0 or t0 + t1 > 1.0:
+    return wp.float32(wp.inf)
+
+  # intersect ray with plane of triangle
+  dif0 = triangle.v0 - triangle.v2  # v0-v2
+  dif1 = triangle.v1 - triangle.v2  # v1-v2
+  dif2 = pnt - triangle.v2  # lpnt-v2
+  nrm = wp.cross(dif0, dif1)  # normal to triangle plane
+  denom = wp.dot(vec, nrm)
+  if wp.abs(denom) < MJ_MINVAL:
+    return wp.float32(wp.inf)
+
+  dist = -wp.dot(dif2, nrm) / denom
+  return wp.where(dist >= 0.0, dist, wp.float32(wp.inf))
+
+
+@wp.func
+def _ray_mesh(
+  m: Model,
+  geom_id: int,
+  unused_size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+) -> DistanceWithId:
+  """Returns the distance and geom_id for ray mesh intersections."""
+  data_id = m.geom_dataid[geom_id]
+
+  # Create basis vectors for the ray
+  basis = Basis()
+
+  # Compute orthogonal basis vectors
+  if wp.abs(vec[0]) < wp.abs(vec[1]):
+    if wp.abs(vec[0]) < wp.abs(vec[2]):
+      basis.b0 = wp.vec3(0.0, vec[2], -vec[1])
+    else:
+      basis.b0 = wp.vec3(vec[1], -vec[0], 0.0)
+  else:
+    if wp.abs(vec[1]) < wp.abs(vec[2]):
+      basis.b0 = wp.vec3(-vec[2], 0.0, vec[0])
+    else:
+      basis.b0 = wp.vec3(vec[1], -vec[0], 0.0)
+
+  # Normalize first basis vector
+  basis.b0 = wp.normalize(basis.b0)
+
+  # Compute second basis vector as cross product
+  basis.b1 = wp.cross(vec, basis.b0)
+  basis.b1 = wp.normalize(basis.b1)
+
+  min_dist = wp.float32(wp.inf)
+  hit_found = int(0)
+
+  # Get mesh vertex data range
+  vert_start = m.mesh_vertadr[data_id]
+  # vert_end = wp.where(
path: mujoco_warp/_src/ray.py
user: login: thowell
body: can we remove the commented out lines?
created_at: 2025-04-14T11:00:51Z
author_association: COLLABORATOR
diff_hunk: @@ -34,6 +34,7 @@
 from ._src.io import put_data as put_data
 from ._src.io import put_model as put_model
 from ._src.passive import passive as passive
+from ._src.ray import ray_geom as ray_geom
path: mujoco_warp/__init__.py
user: login: thowell
body: i think the api function should be `ray`
mujoco reference `mj_ray`: https://github.com/google-deepmind/mujoco/blob/f317bd17c3494b954982de7bcf14030978d5013b/src/engine/engine_ray.c#L1145
mjx reference `ray`: https://github.com/google-deepmind/mujoco/blob/f317bd17c3494b954982de7bcf14030978d5013b/mjx/mujoco/mjx/__init__.py#L37
created_at: 2025-04-14T11:07:58Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,557 @@
+# Copyright 2025 The Newton Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import warp as wp
+
+from .types import MJ_MINVAL
+from .types import Data
+from .types import GeomType
+from .types import Model
+from .types import vec6
+
+
+@wp.struct
+class DistanceWithId:
+  dist: wp.float32
+  geom_id: wp.int32
+
+
+@wp.func
+def _ray_quad(a: float, b: float, c: float) -> wp.vec2:
+  """Returns two solutions for quadratic: a*x^2 + 2*b*x + c = 0."""
+  det = b * b - a * c
+  det_2 = wp.sqrt(det)
+
+  x0 = (-b - det_2) / a
+  x1 = (-b + det_2) / a
+  x0 = wp.where((det < MJ_MINVAL) or (x0 < 0.0), wp.inf, x0)
+  x1 = wp.where((det < MJ_MINVAL) or (x1 < 0.0), wp.inf, x1)
+
+  return wp.vec2(x0, x1)
+
+
+@wp.func
+def _ray_plane(
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with a plane."""
+  x = -pnt[2] / vec[2]
+
+  valid = vec[2] <= -MJ_MINVAL  # z-vec pointing towards front face
+  valid = valid and x >= 0.0
+  # only within rendered rectangle
+  p_x = pnt[0] + x * vec[0]
+  p_y = pnt[1] + x * vec[1]
+  valid = valid and ((size[0] <= 0.0) or (wp.abs(p_x) <= size[0]))
+  valid = valid and ((size[1] <= 0.0) or (wp.abs(p_y) <= size[1]))
+
+  return_id = wp.where(valid, geom_id, -1)
+  return DistanceWithId(wp.where(valid, x, wp.inf), return_id)
+
+
+@wp.func
+def _ray_sphere(
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with a sphere."""
+  a = wp.dot(vec, vec)
+  b = wp.dot(vec, pnt)
+  c = wp.dot(pnt, pnt) - size[0] * size[0]
+
+  solutions = _ray_quad(a, b, c)
+  x0 = solutions[0]
+  x1 = solutions[1]
+  x = wp.where(wp.isinf(x0), x1, x0)
+
+  return_id = wp.where(wp.isinf(x), -1, geom_id)
+  return DistanceWithId(x, return_id)
+
+
+@wp.func
+def _ray_capsule(
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with a capsule."""
+
+  # cylinder round side: (x*lvec+lpnt)'*(x*lvec+lpnt) = size[0]*size[0]
+  # For a capsule, we only care about the x,y components when checking cylinder intersection
+  # since the z component is handled separately with the caps
+  vec_2d = wp.vec2(vec[0], vec[1])
+  pnt_2d = wp.vec2(pnt[0], pnt[1])
+  a = wp.dot(vec_2d, vec_2d)
+  b = wp.dot(pnt_2d, vec_2d)
+  c = wp.dot(pnt_2d, pnt_2d) - size[0] * size[0]
+
+  # solve a*x^2 + 2*b*x + c = 0
+  solutions = _ray_quad(a, b, c)
+  x0 = solutions[0]
+  x1 = solutions[1]
+  x = wp.where(wp.isinf(x0), x1, x0)
+
+  # make sure round solution is between flat sides
+  x = wp.where(wp.abs(pnt[2] + x * vec[2]) <= size[1], x, wp.inf)
+
+  # top cap
+  dif = pnt - wp.vec3(0.0, 0.0, size[1])
+  solutions = _ray_quad(
+    wp.dot(vec, vec), wp.dot(vec, dif), wp.dot(dif, dif) - size[0] * size[0]
+  )
+  x0 = solutions[0]
+  x1 = solutions[1]
+  # accept only top half of sphere
+  x = wp.where((pnt[2] + x0 * vec[2] >= size[1]) and (x0 < x), x0, x)
+  x = wp.where((pnt[2] + x1 * vec[2] >= size[1]) and (x1 < x), x1, x)
+
+  # bottom cap
+  dif = pnt + wp.vec3(0.0, 0.0, size[1])
+  solutions = _ray_quad(
+    wp.dot(vec, vec), wp.dot(vec, dif), wp.dot(dif, dif) - size[0] * size[0]
+  )
+  x0 = solutions[0]
+  x1 = solutions[1]
+  # accept only bottom half of sphere
+  x = wp.where((pnt[2] + x0 * vec[2] <= -size[1]) and (x0 < x), x0, x)
+  x = wp.where((pnt[2] + x1 * vec[2] <= -size[1]) and (x1 < x), x1, x)
+
+  return_id = wp.where(wp.isinf(x), -1, geom_id)
+  return DistanceWithId(x, return_id)
+
+
+@wp.func
+def _ray_ellipsoid(
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with an ellipsoid."""
+
+  # invert size^2
+  s = wp.vec3(
+    1.0 / (size[0] * size[0]), 1.0 / (size[1] * size[1]), 1.0 / (size[2] * size[2])
+  )
+
+  # (x*lvec+lpnt)' * diag(1/size^2) * (x*lvec+lpnt) = 1
+  svec = wp.vec3(s[0] * vec[0], s[1] * vec[1], s[2] * vec[2])
+  a = wp.dot(svec, vec)
+  b = wp.dot(svec, pnt)
+  c = wp.dot(wp.vec3(s[0] * pnt[0], s[1] * pnt[1], s[2] * pnt[2]), pnt) - 1.0
+
+  # solve a*x^2 + 2*b*x + c = 0
+  solutions = _ray_quad(a, b, c)
+  x0 = solutions[0]
+  x1 = solutions[1]
+  x = wp.where(wp.isinf(x0), x1, x0)
+
+  return_id = wp.where(wp.isinf(x), -1, geom_id)
+  return DistanceWithId(x, return_id)
+
+
+@wp.func
+def _ray_box(
+  size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geom_id: int,
+) -> DistanceWithId:
+  """Returns the distance at which a ray intersects with a box."""
+
+  # Initialize with infinity
+  min_x = wp.inf
+
+  # Check all 6 faces of the box (2 faces per axis)
+  for i in range(wp.static(3)):
+    if vec[i] != 0.0:
+      # Get indices for the other two dimensions
+      j = (i + 1) % 3
+      k = (i + 2) % 3
+
+      for t in range(wp.static(2)):
+        s = wp.where(t == 0, size[i], -size[i])
+        sol = (s - pnt[i]) / vec[i]
+
+        pj = pnt[j] + sol * vec[j]
+        pk = pnt[k] + sol * vec[k]
+        if sol >= 0.0 and wp.abs(pj) <= size[j] and wp.abs(pk) <= size[k]:
+          min_x = wp.min(min_x, sol)
+
+  return_id = wp.where(wp.isinf(min_x), -1, geom_id)
+  return DistanceWithId(min_x, return_id)
+
+
+@wp.struct
+class Triangle:
+  """A struct representing a triangle with 3 vertices."""
+
+  v0: wp.vec3
+  v1: wp.vec3
+  v2: wp.vec3
+
+
+@wp.struct
+class Basis:
+  """A struct representing a basis with 2 vectors."""
+
+  b0: wp.vec3
+  b1: wp.vec3
+
+
+@wp.func
+def _ray_triangle(
+  triangle: Triangle,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  basis: Basis,
+) -> wp.float32:
+  """Returns the distance at which a ray intersects with a triangle."""
+  # dif = v[i] - lpnt
+  dif0 = triangle.v0 - pnt
+  dif1 = triangle.v1 - pnt
+  dif2 = triangle.v2 - pnt
+
+  # project difference vectors in normal plane
+  planar_0_0 = wp.dot(dif0, basis.b0)
+  planar_0_1 = wp.dot(dif0, basis.b1)
+  planar_1_0 = wp.dot(dif1, basis.b0)
+  planar_1_1 = wp.dot(dif1, basis.b1)
+  planar_2_0 = wp.dot(dif2, basis.b0)
+  planar_2_1 = wp.dot(dif2, basis.b1)
+
+  # reject if on the same side of any coordinate axis
+  if (
+    (planar_0_0 > 0.0 and planar_1_0 > 0.0 and planar_2_0 > 0.0)
+    or (planar_0_0 < 0.0 and planar_1_0 < 0.0 and planar_2_0 < 0.0)
+    or (planar_0_1 > 0.0 and planar_1_1 > 0.0 and planar_2_1 > 0.0)
+    or (planar_0_1 < 0.0 and planar_1_1 < 0.0 and planar_2_1 < 0.0)
+  ):
+    return wp.float32(wp.inf)
+
+  # determine if origin is inside planar projection of triangle
+  # A = (p0-p2, p1-p2), b = -p2, solve A*t = b
+  A00 = planar_0_0 - planar_2_0
+  A10 = planar_1_0 - planar_2_0
+  A01 = planar_0_1 - planar_2_1
+  A11 = planar_1_1 - planar_2_1
+
+  b0 = -planar_2_0
+  b1 = -planar_2_1
+
+  det = A00 * A11 - A10 * A01
+  if wp.abs(det) < MJ_MINVAL:
+    return wp.float32(wp.inf)
+
+  t0 = (A11 * b0 - A10 * b1) / det
+  t1 = (-A01 * b0 + A00 * b1) / det
+
+  # check if outside
+  if t0 < 0.0 or t1 < 0.0 or t0 + t1 > 1.0:
+    return wp.float32(wp.inf)
+
+  # intersect ray with plane of triangle
+  dif0 = triangle.v0 - triangle.v2  # v0-v2
+  dif1 = triangle.v1 - triangle.v2  # v1-v2
+  dif2 = pnt - triangle.v2  # lpnt-v2
+  nrm = wp.cross(dif0, dif1)  # normal to triangle plane
+  denom = wp.dot(vec, nrm)
+  if wp.abs(denom) < MJ_MINVAL:
+    return wp.float32(wp.inf)
+
+  dist = -wp.dot(dif2, nrm) / denom
+  return wp.where(dist >= 0.0, dist, wp.float32(wp.inf))
+
+
+@wp.func
+def _ray_mesh(
+  m: Model,
+  geom_id: int,
+  unused_size: wp.vec3,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+) -> DistanceWithId:
+  """Returns the distance and geom_id for ray mesh intersections."""
+  data_id = m.geom_dataid[geom_id]
+
+  # Create basis vectors for the ray
+  basis = Basis()
+
+  # Compute orthogonal basis vectors
+  if wp.abs(vec[0]) < wp.abs(vec[1]):
+    if wp.abs(vec[0]) < wp.abs(vec[2]):
+      basis.b0 = wp.vec3(0.0, vec[2], -vec[1])
+    else:
+      basis.b0 = wp.vec3(vec[1], -vec[0], 0.0)
+  else:
+    if wp.abs(vec[1]) < wp.abs(vec[2]):
+      basis.b0 = wp.vec3(-vec[2], 0.0, vec[0])
+    else:
+      basis.b0 = wp.vec3(vec[1], -vec[0], 0.0)
+
+  # Normalize first basis vector
+  basis.b0 = wp.normalize(basis.b0)
+
+  # Compute second basis vector as cross product
+  basis.b1 = wp.cross(vec, basis.b0)
+  basis.b1 = wp.normalize(basis.b1)
+
+  min_dist = wp.float32(wp.inf)
+  hit_found = int(0)
+
+  # Get mesh vertex data range
+  vert_start = m.mesh_vertadr[data_id]
+
+  # Get mesh face and vertex data
+  face_start = m.mesh_faceadr[data_id]
+  face_end = wp.where(
+    data_id + 1 < m.mesh_faceadr.shape[0], m.mesh_faceadr[data_id + 1], m.nmeshface
+  )
+
+  # Iterate through all faces
+  for i in range(face_start, face_end):
+    # Get vertices for this face
+    v_idx = m.mesh_face[i]
+
+    # Create triangle struct
+    triangle = Triangle()
+    triangle.v0 = m.mesh_vert[vert_start + v_idx.x]
+    triangle.v1 = m.mesh_vert[vert_start + v_idx.y]
+    triangle.v2 = m.mesh_vert[vert_start + v_idx.z]
+
+    # Calculate intersection
+    dist = _ray_triangle(triangle, pnt, vec, basis)
path: mujoco_warp/_src/ray.py
user: login: mmacklin
body: Btw Warp has BVH accelerated versions of these routines: https://nvidia.github.io/warp/modules/functions.html#warp.mesh_query_ray, also distance and sign queries: https://nvidia.github.io/warp/modules/functions.html#warp.mesh_query_point.
created_at: 2025-04-22T08:37:48Z
author_association: COLLABORATOR
diff_hunk: @@ -127,17 +127,29 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     tile_beg = tile_corners[i]
     tile_end = mjm.nv if i == len(tile_corners) - 1 else tile_corners[i + 1]
     tiles.setdefault(tile_end - tile_beg, []).append(tile_beg)
+
+  # Handle empty tiles case
path: mujoco_warp/_src/io.py
user: login: thowell
body: is this block of code still utilized?
created_at: 2025-05-12T11:14:19Z
author_association: COLLABORATOR
path: mujoco_warp/test_data/meshes/pyramid.stl
user: login: thowell
body: is this file needed?
created_at: 2025-05-12T11:18:51Z
author_association: COLLABORATOR

PR comments:
user: login: nvtw
created_at: 2025-04-11T15:05:56Z
body: I assume the tests currently fail because wp.tile_argmin was only merged yesterday into warp and might not yet be in the warp build we are using. Will investigate more on Monday.
author_association: COLLABORATOR
user: login: nvtw
created_at: 2025-04-22T08:52:10Z
body: I don't understand why the ray unit tests fail on github. They succeed on my Linux and on my Windows machine.
author_association: COLLABORATOR
user: login: erikfrey
created_at: 2025-04-23T20:41:27Z
body: @nvtw Is it a CPU vs GPU issue?  Do you have a way to force CPU on one of your dev machines and rerun the tests?
author_association: COLLABORATOR
user: login: thowell
created_at: 2025-04-29T11:55:45Z
body: @nvtw it should be possible to reproduce the test failure locally with #200
author_association: COLLABORATOR
user: login: thowell
created_at: 2025-05-01T09:38:19Z
body: @erikfrey is it ok to add these .stl files? it looks like they are from https://github.com/google-deepmind/mujoco/tree/main/mjx/mujoco/mjx/test_data/meshes. 

the alternative would be to define meshes directly in xml?
author_association: COLLABORATOR
user: login: thowell
created_at: 2025-05-05T18:52:37Z
body: request that would be useful for the touch sensor implementation #213:

- update `_ray_geom` to follow the mujoco implementation api, something like:
`ray_geom(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3, geomtype: int) -> float`

here is the mujoco reference for `mju_rayGeom`: https://github.com/google-deepmind/mujoco/blob/253abad210823254f0f238ef3ad4cc8d975372d3/src/engine/engine_ray.c#L818

here is the mujoco reference for `mjSENS_TOUCH`: https://github.com/google-deepmind/mujoco/blob/253abad210823254f0f238ef3ad4cc8d975372d3/src/engine/engine_sensor.c#L777 where `mju_rayGeom` is utilized

thanks!
author_association: COLLABORATOR
user: login: erikfrey
created_at: 2025-05-09T19:41:55Z
body: @thowell yes it's OK to bring those meshes in.
author_association: COLLABORATOR
user: login: nvtw
created_at: 2025-05-12T11:05:59Z
body: Should be all done now
author_association: COLLABORATOR
user: login: thowell
created_at: 2025-05-12T11:24:35Z
body: @nvtw please take a look at the comment about the block of code in io.py and the comment about the mesh file. thanks!
author_association: COLLABORATOR