Add mj_multiRay function to perform multiple ray intersection with the same source.

PiperOrigin-RevId: 528684535
Change-Id: I6e1747a4a0910e6061d6154e19ad8723f56c22e3

doc/APIreference/functions.rst

@@ -843,14 +843,25 @@ where jar = Jac*qacc-aref.
 Ray collisions
 ^^^^^^^^^^^^^^
 
-Ray collision functionality was added in MuJoCo 1.50. This is a new collision detection module that uses analytical
-formulas to intersect a ray ``(p + x*v, x >= 0)`` with a geom, where p is the origin of the ray and v is the vector
-specifying the direction. All functions in this family return the distance to the nearest geom surface, or -1 if there
-is no intersection. Note that if p is inside a geom, the ray will intersect the surface from the inside which still
-counts as an intersection.
+Ray collisions, also known as ray casting, find the distance ``x`` of a ray's intersection with a geom, where a ray is
+a line emanating from the 3D point ``p`` in the direction ``v`` i.e., ``(p + x*v, x >= 0)``. All functions in this
+family return the distance to the nearest geom surface, or -1 if there is no intersection. Note that if ``p`` is inside
+a geom, the ray will intersect the surface from the inside which still counts as an intersection.
 
 All ray collision functions rely on quantities computed by :ref:`mj_kinematics` (see :ref:`mjData`), so must be called
-after  :ref:`mj_kinematics`, or functions that call it (e.g. :ref:`mj_fwdPosition`).
+after  :ref:`mj_kinematics`, or functions that call it (e.g. :ref:`mj_fwdPosition`). The top level functions, which
+intersect with all geoms types, are :ref:`mj_ray` which casts a single ray, and :ref:`mj_multiRay` which casts multiple
+rays from a single point.
+
+.. _mj_multiRay:
+
+mj_multiRay
+~~~~~~~~~~~
+
+.. mujoco-include:: mj_multiRay
+
+Intersect multiple rays emanating from a single point.
+Similar semantics to mj_ray, but vec is an array of (nray x 3) directions.
 
 .. _mj_ray:
 

doc/APIreference/functions_override.rst

@@ -148,14 +148,15 @@ format of qpos.
 
 .. _Raycollisions:
 
-Ray collision functionality was added in MuJoCo 1.50. This is a new collision detection module that uses analytical
-formulas to intersect a ray ``(p + x*v, x >= 0)`` with a geom, where p is the origin of the ray and v is the vector
-specifying the direction. All functions in this family return the distance to the nearest geom surface, or -1 if there
-is no intersection. Note that if p is inside a geom, the ray will intersect the surface from the inside which still
-counts as an intersection.
+Ray collisions, also known as ray casting, find the distance ``x`` of a ray's intersection with a geom, where a ray is
+a line emanating from the 3D point ``p`` in the direction ``v`` i.e., ``(p + x*v, x >= 0)``. All functions in this
+family return the distance to the nearest geom surface, or -1 if there is no intersection. Note that if ``p`` is inside
+a geom, the ray will intersect the surface from the inside which still counts as an intersection.
 
 All ray collision functions rely on quantities computed by :ref:`mj_kinematics` (see :ref:`mjData`), so must be called
-after  :ref:`mj_kinematics`, or functions that call it (e.g. :ref:`mj_fwdPosition`).
+after  :ref:`mj_kinematics`, or functions that call it (e.g. :ref:`mj_fwdPosition`). The top level functions, which
+intersect with all geoms types, are :ref:`mj_ray` which casts a single ray, and :ref:`mj_multiRay` which casts multiple
+rays from a single point.
 
 .. _mj_ray:
 

doc/changelog.rst

@@ -15,7 +15,8 @@ Plugins
 - Added touch-grid sensor plugin. See `documentation <https://github.com/deepmind/mujoco/blob/main/plugin/sensor/README.md>`_
   for details, and associated `touch_grid.xml <https://github.com/deepmind/mujoco/blob/main/model/plugin/touch_grid.xml>`_
   example model. The plugin includes `in-scene visualisation <https://youtu.be/0LOJ3WMnqeA>`_.
-
+- Add ``mj_multiRay`` function for intersecting multiple rays emanating from a single point. This is significantly
+  faster than calling ``mj_ray`` multiple times.
 
 Version 2.3.5 (April 25, 2023)
 ------------------------------

doc/includes/references.h

@@ -2207,6 +2207,9 @@ void mj_loadPluginLibrary(const char* path);
 void mj_loadAllPluginLibraries(const char* directory, mjfPluginLibraryLoadCallback callback);
 int mj_version(void);
 const char* mj_versionString();
+void mj_multiRay(const mjModel* m, mjData* d, const mjtNum pnt[3], const mjtNum* vec,
+                 const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
+                 int* geomid, mjtNum* dist, int nray);
 mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum pnt[3], const mjtNum vec[3],
               const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
               int geomid[1]);

include/mujoco/mujoco.h

@@ -484,6 +484,12 @@ MJAPI const char* mj_versionString();
 
 //---------------------------------- Ray collisions ------------------------------------------------
 
+// Intersect multiple rays emanating from a single point.
+// Similar semantics to mj_ray, but vec is an array of (nray x 3) directions.
+MJAPI void mj_multiRay(const mjModel* m, mjData* d, const mjtNum pnt[3], const mjtNum* vec,
+                       const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
+                       int* geomid, mjtNum* dist, int nray);
+
 // Intersect ray (pnt+x*vec, x>=0) with visible geoms, except geoms in bodyexclude.
 // Return distance (x) to nearest surface, or -1 if no intersection and output geomid.
 // geomgroup, flg_static are as in mjvOption; geomgroup==NULL skips group exclusion.

introspect/functions.py

@@ -2766,6 +2766,69 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          parameters=(),
          doc='Return the current version of MuJoCo as a null-terminated string.',  # pylint: disable=line-too-long
      )),
+    ('mj_multiRay',
+     FunctionDecl(
+         name='mj_multiRay',
+         return_type=ValueType(name='void'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='m',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjModel', is_const=True),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='d',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjData'),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='pnt',
+                 type=ArrayType(
+                     inner_type=ValueType(name='mjtNum', is_const=True),
+                     extents=(3,),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='vec',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjtNum', is_const=True),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='geomgroup',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjtByte', is_const=True),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='flg_static',
+                 type=ValueType(name='mjtByte'),
+             ),
+             FunctionParameterDecl(
+                 name='bodyexclude',
+                 type=ValueType(name='int'),
+             ),
+             FunctionParameterDecl(
+                 name='geomid',
+                 type=PointerType(
+                     inner_type=ValueType(name='int'),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='dist',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjtNum'),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='nray',
+                 type=ValueType(name='int'),
+             ),
+         ),
+         doc='Intersect multiple rays emanating from a single point. Similar semantics to mj_ray, but vec is an array of (nray x 3) directions.',  # pylint: disable=line-too-long
+     )),
     ('mj_ray',
      FunctionDecl(
          name='mj_ray',

python/mujoco/functions.cc

@@ -546,6 +546,7 @@ PYBIND11_MODULE(_functions, pymodule) {
   Def<traits::mj_versionString>(pymodule);
 
   // Ray collision
+  Def<traits::mj_multiRay>(pymodule);
   Def<traits::mj_ray>(
       pymodule,
       [](const raw::MjModel* m, const raw::MjData* d, const mjtNum(*pnt)[3],

src/engine/engine_ray.c

@@ -20,6 +20,7 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjvisualize.h>
+#include "engine/engine_io.h"
 #include "engine/engine_macro.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
@@ -46,6 +47,20 @@ static void ray_map(const mjtNum* pos, const mjtNum* mat, const mjtNum* pnt, con
 
 
 
+// map to azimuth angle in spherical coordinates
+static mjtNum longitude(const mjtNum vec[3]) {
+  return mju_atan2(vec[1], vec[0]);
+}
+
+
+
+// map to elevation angle in spherical coordinates
+static mjtNum latitude(const mjtNum vec[3]) {
+  return mju_atan2(mju_sqrt(vec[0]*vec[0] + vec[1]*vec[1]), vec[2]);
+}
+
+
+
 // eliminate geom
 static int ray_eliminate(const mjModel* m, const mjData* d, int geomid,
                          const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude) {
@@ -207,13 +222,13 @@ static mjtNum ray_plane(const mjtNum* pos, const mjtNum* mat, const mjtNum* size
 
 
 // sphere
-static mjtNum ray_sphere(const mjtNum* pos, const mjtNum* mat, const mjtNum* size,
+static mjtNum ray_sphere(const mjtNum* pos, const mjtNum* mat, mjtNum dist_sqr,
                          const mjtNum* pnt, const mjtNum* vec) {
   // (x*vec+pnt-pos)'*(x*vec+pnt-pos) = size[0]*size[0]
   mjtNum dif[3] = {pnt[0]-pos[0], pnt[1]-pos[1], pnt[2]-pos[2]};
   mjtNum a = vec[0]*vec[0] + vec[1]*vec[1] + vec[2]*vec[2];
   mjtNum b = vec[0]*dif[0] + vec[1]*dif[1] + vec[2]*dif[2];
-  mjtNum c = dif[0]*dif[0] + dif[1]*dif[1] + dif[2]*dif[2] - size[0]*size[0];
+  mjtNum c = dif[0]*dif[0] + dif[1]*dif[1] + dif[2]*dif[2] - dist_sqr;
 
   // solve a*x^2 + 2*b*x + c = 0
   mjtNum xx[2];
@@ -227,7 +242,7 @@ static mjtNum ray_capsule(const mjtNum* pos, const mjtNum* mat, const mjtNum* si
                           const mjtNum* pnt, const mjtNum* vec) {
   // bounding sphere test
   mjtNum ssz = size[0] + size[1];
-  if (ray_sphere(pos, NULL, &ssz, pnt, vec)<0) {
+  if (ray_sphere(pos, NULL, ssz*ssz, pnt, vec)<0) {
     return -1;
   }
 
@@ -315,8 +330,8 @@ static mjtNum ray_ellipsoid(const mjtNum* pos, const mjtNum* mat, const mjtNum*
 static mjtNum ray_cylinder(const mjtNum* pos, const mjtNum* mat, const mjtNum* size,
                            const mjtNum* pnt, const mjtNum* vec) {
   // bounding sphere test
-  mjtNum ssz = mju_sqrt(size[0]*size[0] + size[1]*size[1]);
-  if (ray_sphere(pos, NULL, &ssz, pnt, vec)<0) {
+  mjtNum ssz = size[0]*size[0] + size[1]*size[1];
+  if (ray_sphere(pos, NULL, ssz, pnt, vec)<0) {
     return -1;
   }
 
@@ -382,8 +397,8 @@ static mjtNum ray_box(const mjtNum* pos, const mjtNum* mat, const mjtNum* size,
   }
 
   // bounding sphere test
-  mjtNum ssz = mju_sqrt(size[0]*size[0] + size[1]*size[1] + size[2]*size[2]);
-  if (ray_sphere(pos, NULL, &ssz, pnt, vec)<0) {
+  mjtNum ssz = size[0]*size[0] + size[1]*size[1] + size[2]*size[2];
+  if (ray_sphere(pos, NULL, ssz, pnt, vec)<0) {
     return -1;
   }
 
@@ -659,7 +674,7 @@ mjtNum mju_rayGeom(const mjtNum* pos, const mjtNum* mat, const mjtNum* size,
     return ray_plane(pos, mat, size, pnt, vec);
 
   case mjGEOM_SPHERE:
-    return ray_sphere(pos, mat, size, pnt, vec);
+    return ray_sphere(pos, mat, size[0]*size[0], pnt, vec);
 
   case mjGEOM_CAPSULE:
     return ray_capsule(pos, mat, size, pnt, vec);
@@ -774,6 +789,28 @@ mjtNum mju_raySkin(int nface, int nvert, const int* face, const float* vert,
 
 
 
+// return 1 if point is inside object-aligned bounding box, 0 otherwise
+static int point_in_box(const mjtNum aabb[6], const mjtNum xpos[3],
+                        const mjtNum xmat[9], const mjtNum pnt[3]) {
+  mjtNum point[3];
+
+  // compute point in local coordinates of the box
+  mju_sub3(point, pnt, xpos);
+  mju_rotVecMatT(point, point, xmat);
+  mju_subFrom3(point, aabb);
+
+  // check intersections
+  for (int j=0; j<3; j++) {  // directions
+    if (mju_abs(point[j]) > aabb[3+j]) {
+      return 0;
+    }
+  }
+
+  return 1;
+}
+
+
+
 //---------------------------- main entry point ---------------------------------------------------
 
 // intersect ray (pnt+x*vec, x>=0) with visible geoms, except geoms on bodyexclude
@@ -819,3 +856,169 @@ mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum* pnt, const mjtNum
 
   return dist;
 }
+
+
+// Initializes spherical bounding angles (geom_ba) and flag vector for a given source
+void mju_multiRayPrepare(const mjModel* m, const mjData* d, const mjtNum pnt[3],
+                         const mjtNum* ray_xmat, const mjtByte* geomgroup, mjtByte flg_static,
+                         int bodyexclude, mjtNum* geom_ba, int* geom_eliminate) {
+  if (ray_xmat) {
+    mju_error("ray_xmat is currently unused, should be NULL");
+  }
+
+  if (geom_eliminate) {
+    // compute eliminate flag for all geoms
+    for (int geomid=0; geomid<m->ngeom; geomid++)
+      geom_eliminate[geomid] = ray_eliminate(m, d, geomid, geomgroup, flg_static, bodyexclude);
+  }
+
+  for (int b=0; b<m->nbody; b++) {
+    // skip precomputation if no bounding volume is available
+    if (m->body_bvhadr[b] == -1) {
+      continue;
+    }
+
+    // loop over child geoms, compute bounding angles
+    for (int i=0; i<m->body_geomnum[b]; i++) {
+      int g = i + m->body_geomadr[b];
+      mjtNum AABB[4] = {mjMAXVAL, mjMAXVAL, -mjMAXVAL, -mjMAXVAL};
+      mjtNum* aabb = m->geom_aabb + 6*g;
+      mjtNum* xpos = d->geom_xpos + 3*g;
+      mjtNum* xmat = d->geom_xmat + 9*g;
+
+      if (point_in_box(aabb, xpos, xmat, pnt)) {
+        (geom_ba+4*g)[0] = -mjPI;
+        (geom_ba+4*g)[1] = -mjPI/2;
+        (geom_ba+4*g)[2] = mjPI;
+        (geom_ba+4*g)[3] = mjPI/2;
+        continue;
+      }
+
+      // loop over box vertices, compute spherical aperture
+      for (int v=0; v<8; v++) {
+        mjtNum vert[3], box[3];
+        vert[0] = (v&1 ? aabb[0]+aabb[3] : aabb[0]-aabb[3]);
+        vert[1] = (v&2 ? aabb[1]+aabb[4] : aabb[1]-aabb[4]);
+        vert[2] = (v&4 ? aabb[2]+aabb[5] : aabb[2]-aabb[5]);
+
+        // rotate to the world frame
+        mju_rotVecMat(box, vert, xmat);
+        mju_addTo3(box, xpos);
+
+        // spherical coordinates
+        mju_sub3(vert, box, pnt);
+        mjtNum azimuth = longitude(vert);
+        mjtNum elevation = latitude(vert);
+
+        // update bounds
+        AABB[0] = mju_min(AABB[0], azimuth);
+        AABB[1] = mju_min(AABB[1], elevation);
+        AABB[2] = mju_max(AABB[2], azimuth);
+        AABB[3] = mju_max(AABB[3], elevation);
+      }
+
+      if (AABB[2]-AABB[0] > mjPI) {
+        AABB[0] = -mjPI;
+        AABB[2] =  mjPI;
+      }
+
+      if (AABB[3]-AABB[1] > mjPI) {  // SHOULD NOT OCCUR
+        mju_error("mj_ray: discontinuity in azimuth angle");
+      }
+
+      mju_copy(geom_ba+4*g, AABB, 4);
+    }
+  }
+}
+
+
+// Performs single ray intersection
+static mjtNum mju_singleRay(const mjModel* m, mjData* d, const mjtNum pnt[3], const mjtNum vec[3],
+                            int* ray_eliminate, mjtNum* geom_ba, int geomid[1]) {
+  mjtNum dist, newdist;
+
+  // check vector length
+  if (mju_norm3(vec)<mjMINVAL) {
+    mju_error("mj_ray: vector length is too small");
+  }
+
+  // clear result
+  dist = -1;
+  *geomid = -1;
+
+  // get ray spherical coordinates
+  mjtNum azimuth = longitude(vec);
+  mjtNum elevation = latitude(vec);
+
+  // loop over bodies not eliminated by bodyexclude
+  for (int b=0; b<m->nbody; b++) {
+    // exclude body using bounding sphere test
+    if (m->body_bvhadr[b] != -1) {
+      mjtNum* pos = m->bvh_aabb + 6*m->body_bvhadr[b];
+      mjtNum* size = pos + 3;
+      mjtNum ssz = size[0]*size[0] + size[1]*size[1] + size[2]*size[2];
+      if (ray_sphere(pos, NULL, ssz, pnt, vec)<0) {
+        continue;
+      }
+    }
+
+    // loop over geoms if bounding sphere test fails
+    for (int g=0; g<m->body_geomnum[b]; g++) {
+      int i = m->body_geomadr[b] + g;
+      if (ray_eliminate[i]) {
+        continue;
+      }
+
+      // exclude geom using bounding angles
+      if (m->body_bvhadr[b] != -1) {
+        if (azimuth<(geom_ba+4*i)[0] || elevation<(geom_ba+4*i)[1] ||
+            azimuth>(geom_ba+4*i)[2] || elevation>(geom_ba+4*i)[3]) {
+          continue;
+        }
+      }
+
+      // handle mesh and hfield separately
+      if (m->geom_type[i]==mjGEOM_MESH) {
+        newdist = mj_rayMesh(m, d, i, pnt, vec);
+      } else if (m->geom_type[i]==mjGEOM_HFIELD) {
+        newdist = mj_rayHfield(m, d, i, pnt, vec);
+      }
+
+      // otherwise general dispatch
+      else {
+        newdist = mju_rayGeom(d->geom_xpos+3*i, d->geom_xmat+9*i,
+                              m->geom_size+3*i, pnt, vec, m->geom_type[i]);
+      }
+
+      // update if closer intersection found
+      if (newdist>=0 && (newdist<dist || dist<0)) {
+        dist = newdist;
+        *geomid = i;
+      }
+    }
+  }
+
+  return dist;
+}
+
+
+// Performs multiple ray intersections with the precomputes bv and flags
+void mj_multiRay(const mjModel* m, mjData* d, const mjtNum pnt[3], const mjtNum* vec,
+                 const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
+                 int* geomid, mjtNum* dist, int nray) {
+  mjMARKSTACK;
+
+  // allocate source
+  mjtNum* geom_ba = mj_stackAlloc(d, 4*m->ngeom);
+  int* geom_eliminate = mj_stackAllocInt(d, m->ngeom);
+
+  // initialize source
+  mju_multiRayPrepare(m, d, pnt, NULL, geomgroup, flg_static, bodyexclude, geom_ba, geom_eliminate);
+
+  // loop over rays
+  for (int i=0; i<nray; i++) {
+    dist[i] = mju_singleRay(m, d, pnt, vec+3*i, geom_eliminate, geom_ba, geomid+i);
+  }
+
+  mjFREESTACK;
+}

src/engine/engine_ray.h

@@ -23,6 +23,16 @@
 extern "C" {
 #endif
 
+MJAPI void mju_multiRayPrepare(const mjModel* m, const mjData* d, const mjtNum pnt[3],
+                               const mjtNum* ray_xmat, const mjtByte* geomgroup, mjtByte flg_static,
+                               int bodyexclude, mjtNum* geom_ba, int* geom_eliminate);
+
+// Intersect multiple rays emanating from a single source
+// Similar semantics to mj_ray, but vec is an array of (nray x 3) directions.
+MJAPI void mj_multiRay(const mjModel* m, mjData* d, const mjtNum pnt[3], const mjtNum* vec,
+                       const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
+                       int* geomid, mjtNum* dist, int nray);
+
 // intersect ray (pnt+x*vec, x>=0) with visible geoms, except geoms on bodyexclude
 //  return geomid and distance (x) to nearest surface, or -1 if no intersection
 //  geomgroup, flg_static are as in mjvOption; geomgroup==NULL skips group exclusion

test/engine/engine_ray_test.cc

@@ -18,12 +18,22 @@
 #include <gtest/gtest.h>
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
+#include <mujoco/mjtnum.h>
 #include <mujoco/mujoco.h>
+#include "src/engine/engine_ray.h"
 #include "test/fixture.h"
 
 namespace mujoco {
 namespace {
 
+static constexpr char kSingleGeomModel[] = R"(
+<mujoco>
+  <worldbody>
+    <geom type="sphere" size=".1" pos="0 0 0"/>
+  </worldbody>
+</mujoco>
+)";
+
 static constexpr char kRayCastingModel[] = R"(
 <mujoco>
   <worldbody>
@@ -128,5 +138,112 @@ TEST_F(RayTest, ExcludeStatic) {
   mj_deleteModel(model);
 }
 
+TEST_F(RayTest, MultiRayEqualsSingleRay) {
+  mjModel* m = LoadModelFromString(kRayCastingModel);
+  ASSERT_THAT(m, NotNull());
+  mjData* d = mj_makeData(m);
+  ASSERT_THAT(d, NotNull());
+  mj_forward(m, d);
+
+  // create ray array
+  constexpr int N = 80;
+  constexpr int M = 60;
+  mjtNum vec[3*N*M];
+  mjtNum pnt[3] = {1, 2, 3};
+  mjtNum cone[4][3] = {{1, 1, -1}, {1, 1, 1}, {1, -1, -1}, {1, -1, 1}};
+  memset(vec, 0, 3*N*M*sizeof(mjtNum));
+
+  for (int i = 0; i < N; ++i) {
+    for (int j = 0; j < M; ++j) {
+      for (int k = 0; k < 3; ++k) {
+        vec[3 * (i * M + j) + k] =           i * cone[0][k] / (N - 1) +
+                                             j * cone[1][1] / (M - 1) +
+                                   (N - i - 1) * cone[2][k] / (N - 1) +
+                                   (M - j - 1) * cone[3][k] / (M - 1);
+      }
+    }
+  }
+
+  // compute intersections with multiray functions
+  mjtNum dist_multiray[3*N*M];
+  int rgeomid_multiray[N*M];
+  mj_multiRay(m, d, pnt, vec, NULL, 1, -1, rgeomid_multiray, dist_multiray, N*M);
+
+  // compare results with single ray function
+  mjtNum dist;
+  int rgeomid;
+
+  for (int i = 0; i < N; ++i) {
+    for (int j = 0; j < M; ++j) {
+      int idx = i * M + j;
+      dist = mj_ray(m, d, pnt, vec + 3 * idx, NULL, 1, -1, &rgeomid);
+      EXPECT_FLOAT_EQ(dist, dist_multiray[idx]);
+      EXPECT_EQ(rgeomid, rgeomid_multiray[idx]);
+    }
+  }
+
+  mj_deleteData(d);
+  mj_deleteModel(m);
+}
+
+TEST_F(RayTest, EdgeCases) {
+  mjModel* m = LoadModelFromString(kSingleGeomModel);
+  ASSERT_THAT(m, NotNull());
+  ASSERT_THAT(m->nbvh, 1);
+  mjData* d = mj_makeData(m);
+  ASSERT_THAT(d, NotNull());
+  mj_forward(m, d);
+
+  // spherical bounding box and result arrays
+  mjtNum geom_ba[4];
+  mjtNum dist;
+  int rgeomid;
+
+  // pnt contained in bounding box
+  mjtNum pnt1[] = {0, 0, 0};
+  mju_multiRayPrepare(m, d, pnt1, NULL, NULL, 1, -1, geom_ba, NULL);
+  EXPECT_FLOAT_EQ(geom_ba[0], -mjPI);
+  EXPECT_FLOAT_EQ(geom_ba[1], -mjPI/2);
+  EXPECT_FLOAT_EQ(geom_ba[2],  mjPI);
+  EXPECT_FLOAT_EQ(geom_ba[3],  mjPI/2);
+  mjtNum vec1[] = {1, 0, 0};
+  mj_multiRay(m, d, pnt1, vec1, NULL, 1, -1, &rgeomid, &dist, 1);
+  EXPECT_FLOAT_EQ(dist, 0.1);
+
+  // pnt at phi = Pi, -Pi
+  mjtNum pnt2[] = {1, 0, 0};
+  mju_multiRayPrepare(m, d, pnt2, NULL, NULL, 1, -1, geom_ba, NULL);
+  EXPECT_FLOAT_EQ(geom_ba[0], -mjPI);  // atan(y<0, x<0)
+  EXPECT_FLOAT_EQ(geom_ba[2],  mjPI);  // atan(y>0, x<0)
+  mjtNum vec2[] = {-1, 0, 0};
+  mj_multiRay(m, d, pnt2, vec2, NULL, 1, -1, &rgeomid, &dist, 1);
+  EXPECT_FLOAT_EQ(dist, 0.9);
+
+  // pnt on the boundary of the box
+  mjtNum pnt3[] = {.1, .1, .05};
+  mju_multiRayPrepare(m, d, pnt3, NULL, NULL, 1, -1, geom_ba, NULL);
+  EXPECT_FLOAT_EQ(geom_ba[1], -mjPI/2);
+  EXPECT_FLOAT_EQ(geom_ba[3],  mjPI/2);
+  mjtNum vec3[] = {1, 1, 0};
+  mj_multiRay(m, d, pnt3, vec3, NULL, 1, -1, &rgeomid, &dist, 1);
+  EXPECT_FLOAT_EQ(dist, -1);
+
+  // size 0 geom
+  mjtNum pnt4[] = {-1, 0, 0};
+  m->geom_aabb[0] = m->geom_aabb[1] = m->geom_aabb[2] = 0;
+  m->geom_aabb[3] = m->geom_aabb[4] = m->geom_aabb[5] = 0;
+  mju_multiRayPrepare(m, d, pnt4, NULL, NULL, 1, -1, geom_ba, NULL);
+  EXPECT_FLOAT_EQ(geom_ba[0], 0);
+  EXPECT_FLOAT_EQ(geom_ba[1], mjPI/2);
+  EXPECT_FLOAT_EQ(geom_ba[2], 0);
+  EXPECT_FLOAT_EQ(geom_ba[3], mjPI/2);
+  mjtNum vec4[] = {1, 0, 0};
+  mj_multiRay(m, d, pnt4, vec4, NULL, 1, -1, &rgeomid, &dist, 1);
+  EXPECT_FLOAT_EQ(dist, 0.9);
+
+  mj_deleteData(d);
+  mj_deleteModel(m);
+}
+
 }  // namespace
 }  // namespace mujoco

unity/Runtime/Bindings/MjBindings.cs

@@ -3280,6 +3280,9 @@ public static unsafe extern void mj_loadPluginLibrary([MarshalAs(UnmanagedType.L
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern int mj_version();
 
+[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
+public static unsafe extern void mj_multiRay(mjModel_* m, mjData_* d, double* pnt, double* vec, byte* geomgroup, byte flg_static, int bodyexclude, int* geomid, double* dist, int nray);
+
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern double mj_ray(mjModel_* m, mjData_* d, double* pnt, double* vec, byte* geomgroup, byte flg_static, int bodyexclude, int* geomid);