Add cutoff argument to mj_multiRay.

PiperOrigin-RevId: 532158635
Change-Id: I70f98be9f781c8281e21aa5c17bd85f3654e9383

doc/includes/references.h

@@ -2216,7 +2216,7 @@ 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);
+                 int* geomid, mjtNum* dist, int nray, mjtNum cutoff);
 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

@@ -488,7 +488,7 @@ MJAPI const char* mj_versionString();
 // 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);
+                       int* geomid, mjtNum* dist, int nray, mjtNum cutoff);
 
 // 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.

introspect/functions.py

@@ -2826,6 +2826,10 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
                  name='nray',
                  type=ValueType(name='int'),
              ),
+             FunctionParameterDecl(
+                 name='cutoff',
+                 type=ValueType(name='mjtNum'),
+             ),
          ),
          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
      )),

src/engine/engine_ray.c

@@ -861,16 +861,14 @@ mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum* pnt, const mjtNum
 // 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) {
+                         int bodyexclude, mjtNum cutoff, 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
-    // compute eliminate flag for all geoms
+  for (int geomid=0; geomid<m->ngeom; geomid++)
-    for (int geomid=0; geomid<m->ngeom; geomid++)
+    geom_eliminate[geomid] = ray_eliminate(m, d, geomid, geomgroup, flg_static, bodyexclude);
-      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
@@ -886,6 +884,17 @@ void mju_multiRayPrepare(const mjModel* m, const mjData* d, const mjtNum pnt[3],
       mjtNum* xpos = d->geom_xpos + 3*g;
       mjtNum* xmat = d->geom_xmat + 9*g;
 
+      // skip if eliminated by flags
+      if (geom_eliminate[g]) {
+        continue;
+      }
+
+      // add to geom_eliminate if distance of bounding sphere is above cutoff
+      if (mju_dist3(d->geom_xpos+3*g, pnt)>cutoff+m->geom_rbound[g]) {
+        geom_eliminate[g] = 1;
+        continue;
+      }
+
       if (point_in_box(aabb, xpos, xmat, pnt)) {
         (geom_ba+4*g)[0] = -mjPI;
         (geom_ba+4*g)[1] = 0;
@@ -1009,7 +1018,7 @@ static mjtNum mju_singleRay(const mjModel* m, mjData* d, const mjtNum pnt[3], co
 // 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) {
+                 int* geomid, mjtNum* dist, int nray, mjtNum cutoff) {
   mjMARKSTACK;
 
   // allocate source
@@ -1017,7 +1026,8 @@ void mj_multiRay(const mjModel* m, mjData* d, const mjtNum pnt[3], const mjtNum*
   int* geom_eliminate = mj_stackAllocInt(d, m->ngeom);
 
   // initialize source
-  mju_multiRayPrepare(m, d, pnt, NULL, geomgroup, flg_static, bodyexclude, geom_ba, geom_eliminate);
+  mju_multiRayPrepare(m, d, pnt, NULL, geomgroup, flg_static, bodyexclude,
+                      cutoff, geom_ba, geom_eliminate);
 
   // loop over rays
   for (int i=0; i<nray; i++) {

src/engine/engine_ray.h

@@ -23,15 +23,17 @@
 extern "C" {
 #endif
 
-MJAPI void mju_multiRayPrepare(const mjModel* m, const mjData* d, const mjtNum pnt[3],
+MJAPI void mju_multiRayPrepare(const mjModel* m, const mjData* d,
-                               const mjtNum* ray_xmat, const mjtByte* geomgroup, mjtByte flg_static,
+                               const mjtNum pnt[3], const mjtNum* ray_xmat,
-                               int bodyexclude, mjtNum* geom_ba, int* geom_eliminate);
+                               const mjtByte* geomgroup, mjtByte flg_static,
+                               int bodyexclude, mjtNum cutoff, 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);
+                       int* geomid, mjtNum* dist, int nray, mjtNum cutoff);
 
 // 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

test/engine/engine_ray_test.cc

@@ -29,8 +29,8 @@ namespace {
 static constexpr char kSingleGeomModel[] = R"(
 <mujoco>
   <worldbody>
-    <body pos="-1 0 0">
+    <body pos="-2 0 0">
-      <geom type="sphere" size=".1"/>
+      <geom type="sphere" pos="1 0 0" size=".1"/>
     </body>
   </worldbody>
 </mujoco>
@@ -169,7 +169,8 @@ TEST_F(RayTest, MultiRayEqualsSingleRay) {
   // 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);
+  mj_multiRay(m, d, pnt, vec, NULL, 1, -1, rgeomid_multiray, dist_multiray,
+              N * M, mjMAXVAL);
 
   // compare results with single ray function
   mjtNum dist;
@@ -200,47 +201,57 @@ TEST_F(RayTest, EdgeCases) {
   mjtNum geom_ba[4];
   mjtNum dist;
   int rgeomid;
+  int flags[1] = {0};
 
   // pnt contained in bounding box
   mjtNum pnt1[] = {-1, 0, 0};
-  mju_multiRayPrepare(m, d, pnt1, NULL, NULL, 1, -1, geom_ba, NULL);
+  mju_multiRayPrepare(m, d, pnt1, NULL, NULL, 1, -1, mjMAXVAL, geom_ba, flags);
   EXPECT_FLOAT_EQ(geom_ba[0], -mjPI);
   EXPECT_FLOAT_EQ(geom_ba[1],  0);
   EXPECT_FLOAT_EQ(geom_ba[2],  mjPI);
   EXPECT_FLOAT_EQ(geom_ba[3],  mjPI);
   mjtNum vec1[] = {1, 0, 0};
-  mj_multiRay(m, d, pnt1, vec1, NULL, 1, -1, &rgeomid, &dist, 1);
+  mj_multiRay(m, d, pnt1, vec1, NULL, 1, -1, &rgeomid, &dist, 1, mjMAXVAL);
   EXPECT_FLOAT_EQ(dist, 0.1);
 
   // pnt at phi = Pi, -Pi
   mjtNum pnt2[] = {-.5, 0, 0};
-  mju_multiRayPrepare(m, d, pnt2, NULL, NULL, 1, -1, geom_ba, NULL);
+  mju_multiRayPrepare(m, d, pnt2, NULL, NULL, 1, -1, mjMAXVAL, geom_ba, flags);
   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);
+  mj_multiRay(m, d, pnt2, vec2, NULL, 1, -1, &rgeomid, &dist, 1, mjMAXVAL);
   EXPECT_FLOAT_EQ(dist, 0.4);
 
+  // with cutoff
+  mjtNum cutoff1 = 0.41, cutoff2 = 0.39;
+  mju_multiRayPrepare(m, d, pnt2, NULL, NULL, 1, -1, cutoff1, geom_ba, flags);
+  EXPECT_EQ(flags[0], 0);
+  mju_multiRayPrepare(m, d, pnt2, NULL, NULL, 1, -1, cutoff2, geom_ba, flags);
+  EXPECT_EQ(flags[0], 1);
+  mj_multiRay(m, d, pnt2, vec2, NULL, 1, -1, &rgeomid, &dist, 1, cutoff2);
+  EXPECT_FLOAT_EQ(dist, -1);
+
   // pnt on the boundary of the box
   mjtNum pnt3[] = {.1, .1, .05};
-  mju_multiRayPrepare(m, d, pnt3, NULL, NULL, 1, -1, geom_ba, NULL);
+  mju_multiRayPrepare(m, d, pnt3, NULL, NULL, 1, -1, mjMAXVAL, geom_ba, flags);
   EXPECT_FLOAT_EQ(geom_ba[1], 0);
   EXPECT_FLOAT_EQ(geom_ba[3], mjPI);
   mjtNum vec3[] = {1, 1, 0};
-  mj_multiRay(m, d, pnt3, vec3, NULL, 1, -1, &rgeomid, &dist, 1);
+  mj_multiRay(m, d, pnt3, vec3, NULL, 1, -1, &rgeomid, &dist, 1, mjMAXVAL);
   EXPECT_FLOAT_EQ(dist, -1);
 
   // size 0 geom
   mjtNum pnt4[] = {-2, 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);
+  mju_multiRayPrepare(m, d, pnt4, NULL, NULL, 1, -1, mjMAXVAL, geom_ba, flags);
   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);
+  mj_multiRay(m, d, pnt4, vec4, NULL, 1, -1, &rgeomid, &dist, 1, mjMAXVAL);
   EXPECT_FLOAT_EQ(dist, 0.9);
 
   mj_deleteData(d);

unity/Runtime/Bindings/MjBindings.cs

@@ -3284,7 +3284,7 @@ public static unsafe extern void mj_loadPluginLibrary([MarshalAs(UnmanagedType.L
 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);
+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, double cutoff);
 
 [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);