Use mesh BVH for speeding up ray mesh intersection.

PiperOrigin-RevId: 534141245
Change-Id: I9f5a71ea059a1a232fd2910d1efc0506ab6dbef3

doc/APIreference/APIglobals.rst

@@ -409,6 +409,10 @@ table below. Their names are in the format ``mjKEY_XXX``. They correspond to GLF
        respect this limit, and user-defined functions should also respect it. Such functions are called with a return
        buffer of size ``mjMAXCONPAIR``; attempting to write more contacts in the buffer can cause unpredictable
        behavior.
+   * - ``mjMAXTREEDEPTH``
+     - 50
+     - The maximum depth of each body and mesh bounding volume hierarchy. If this large limit is exceeded, a warning
+       is raised and ray casting may not be possible. For a balanced hierarchy, this implies 1E15 bounding volumes.
    * - ``mjMAXVFS``
      - 200
      - The maximal number of characters in the name of each file in the virtual file system.

doc/changelog.rst

@@ -32,6 +32,7 @@ General
   trajectory optimization. See :ref:`mju_cholFactorBand` documentation for details.
 - Added :ref:`mj_multiRay` function for intersecting multiple rays emanating from a single point.
   This is significantly faster than calling :ref:`mj_ray` multiple times.
+- Ray-mesh collisions are now up to 10x faster, using a bounding volume hierarchy of mesh faces.
 - Increased ``mjMAXUIITEM`` (maximum number of UI elements per section in Simulate) to 100.
 - Added :ref:`documentation<exProvider>` for resource providers.
 - Changed the formula for :ref:`mju_sigmoid`, a finite-support sigmoid :math:`s \colon \mathbf R \rightarrow [0, 1]`.

include/mujoco/mjmodel.h

@@ -24,6 +24,7 @@
 #define mjMINIMP        0.0001    // minimum constraint impedance
 #define mjMAXIMP        0.9999    // maximum constraint impedance
 #define mjMAXCONPAIR    50        // maximum number of contacts per geom pair
+#define mjMAXTREEDEPTH  50        // maximum bounding volume hierarchy depth
 #define mjMAXVFS        2000      // maximum number of files in virtual file system
 #define mjMAXVFSNAME    1000      // maximum filename size in virtual file system
 

src/engine/engine_ray.c

@@ -27,6 +27,8 @@
 #include "engine/engine_util_misc.h"
 #include "engine/engine_util_spatial.h"
 
+
+
 //---------------------------- utility functions ---------------------------------------------------
 
 // map ray to local geom frame
@@ -130,8 +132,8 @@ static mjtNum ray_quad(mjtNum a, mjtNum b, mjtNum c, mjtNum* x) {
 
 
 // intersect ray with triangle
-static mjtNum ray_triangle(mjtNum v[][3], const mjtNum* lpnt, const mjtNum* lvec,
+mjtNum ray_triangle(mjtNum v[][3], const mjtNum* lpnt, const mjtNum* lvec,
-                           const mjtNum* b0, const mjtNum* b1) {
+                    const mjtNum* b0, const mjtNum* b1) {
   // dif = v[i] - lpnt
   mjtNum dif[3][3];
   for (int i=0; i<3; i++) {
@@ -186,8 +188,6 @@ static mjtNum ray_triangle(mjtNum v[][3], const mjtNum* lpnt, const mjtNum* lvec
   return (-mju_dot3(dif[2], nrm) / denom);
 }
 
-
-
 //---------------------------- geom-specific intersection functions --------------------------------
 
 // plane
@@ -599,19 +599,53 @@ mjtNum mj_rayHfield(const mjModel* m, const mjData* d, int id,
 
 
 
-// intersect ray with mesh
+// ray vs axis-aligned bounding box using slab method
-mjtNum mj_rayMesh(const mjModel* m, const mjData* d, int id,
+// see Ericson, Real-time Collision Detection section 5.3.3.
-                  const mjtNum* pnt, const mjtNum* vec) {
+int mju_raySlab(const mjtNum aabb[6], const mjtNum xpos[3],
-  // check geom type
+                const mjtNum xmat[9], const mjtNum* pnt, const mjtNum* vec) {
-  if (m->geom_type[id]!=mjGEOM_MESH) {
+  mjtNum tmin = 0.0, tmax = INFINITY;
-    mju_error("mj_rayMesh: geom with mesh type expected");
+
+  // compute min and max
+  mjtNum min[3] = {aabb[0]-aabb[3], aabb[1]-aabb[4], aabb[2]-aabb[5]};
+  mjtNum max[3] = {aabb[0]+aabb[3], aabb[1]+aabb[4], aabb[2]+aabb[5]};
+
+  // compute ray in local coordinates
+  mjtNum src[3], dir[3];
+  ray_map(xpos, xmat, pnt, vec, src, dir);
+
+  // check intersections
+  mjtNum invdir[3] = { 1.0 / dir[0], 1.0 / dir[1], 1.0 / dir[2] };
+  for (int d = 0; d < 3; ++d) {
+    mjtNum t1 = (min[d] - src[d]) * invdir[d];
+    mjtNum t2 = (max[d] - src[d]) * invdir[d];
+    mjtNum minval = t1 < t2 ? t1 : t2;
+    mjtNum maxval = t1 < t2 ? t2 : t1;
+    tmin = tmin > minval ? tmin : minval;
+    tmax = tmax < maxval ? tmax : maxval;
   }
 
-  // bounding box test
+  return tmin < tmax;
-  if (ray_box(d->geom_xpos+3*id, d->geom_xmat+9*id, m->geom_size+3*id, pnt, vec, NULL)<0) {
+}
-    return -1;
+
+// ray vs tree intersection
+mjtNum mju_rayTree(const mjModel* m, const mjData* d, int id, const mjtNum* pnt,
+                   const mjtNum* vec) {
+  const int meshid = m->geom_dataid[id];
+  const int bvhadr = m->mesh_bvhadr[meshid];
+  const int* faceid = m->bvh_geomid + bvhadr;
+  const mjtNum* bvh = m->bvh_aabb + 6*bvhadr;
+  const int* child = m->bvh_child + 2*bvhadr;
+
+  if (meshid==-1) {
+    mju_error("mju_rayTree: mesh id of geom %d is -1", meshid);  // SHOULD NOT OCCUR
   }
 
+  // initialize stack
+  int stack[mjMAXTREEDEPTH];
+  int nstack = 0;
+  stack[nstack] = 0;
+  nstack++;
+
   // map to local frame
   mjtNum lpnt[3], lvec[3];
   ray_map(d->geom_xpos+3*id, d->geom_xmat+9*id, pnt, vec, lpnt, lvec);
@@ -633,37 +667,79 @@ mjtNum mj_rayMesh(const mjModel* m, const mjData* d, int id,
   // init solution
   mjtNum x = -1, sol;
 
-  // process all triangles
+  while (nstack) {
-  int face, meshid = m->geom_dataid[id];
+    // pop from stack
-  for (face = m->mesh_faceadr[meshid];
+    nstack--;
-       face < m->mesh_faceadr[meshid] + m->mesh_facenum[meshid];
+    int node = stack[nstack];
-       face++) {
-    // get float vertices
-    float* vf[3];
-    vf[0] = m->mesh_vert + 3*(m->mesh_face[3*face]   + m->mesh_vertadr[meshid]);
-    vf[1] = m->mesh_vert + 3*(m->mesh_face[3*face+1] + m->mesh_vertadr[meshid]);
-    vf[2] = m->mesh_vert + 3*(m->mesh_face[3*face+2] + m->mesh_vertadr[meshid]);
 
-    // convert to mjtNum
+    // intersection test
-    mjtNum v[3][3];
+    int intersect = mju_raySlab(bvh+6*node, d->geom_xpos+3*id, d->geom_xmat+9*id, pnt, vec);
-    for (int i=0; i<3; i++) {
+
-      for (int j=0; j<3; j++) {
+    // if no intersection, skip
-        v[i][j] = (mjtNum)vf[i][j];
+    if (!intersect) {
-      }
+      continue;
     }
 
-    // solve
+    // node1 is a leaf
-    sol = ray_triangle(v, lpnt, lvec, b0, b1);
+    if (faceid[node] != -1) {
+      int face = faceid[node] + m->mesh_faceadr[meshid];
 
-    // update
+      // get float vertices
-    if (sol>=0 && (x<0 || sol<x)) {
+      float* vf[3];
-      x = sol;
+      vf[0] = m->mesh_vert + 3*(m->mesh_face[3*face+0] + m->mesh_vertadr[meshid]);
+      vf[1] = m->mesh_vert + 3*(m->mesh_face[3*face+1] + m->mesh_vertadr[meshid]);
+      vf[2] = m->mesh_vert + 3*(m->mesh_face[3*face+2] + m->mesh_vertadr[meshid]);
+
+      // convert to mjtNum
+      mjtNum v[3][3];
+      for (int i=0; i<3; i++) {
+        for (int j=0; j<3; j++) {
+          v[i][j] = (mjtNum)vf[i][j];
+        }
+      }
+
+      // solve
+      sol = ray_triangle(v, lpnt, lvec, b0, b1);
+
+      // update
+      if (sol>=0 && (x<0 || sol<x)) {
+        x = sol;
+      }
+      continue;
+    }
+
+    // used for rendering
+    d->bvh_active[node + bvhadr] = 1;
+
+    // add children to the stack
+    for (int i=0; i<2; i++) {
+      if (child[2*node+i] != -1) {
+        if (nstack >= mjMAXTREEDEPTH) mju_error("BVH stack depth exceeded in geom %d.", id);
+        stack[nstack] = child[2*node+i];
+        nstack++;
+      }
     }
   }
 
   return x;
 }
 
+// intersect ray with mesh
+mjtNum mj_rayMesh(const mjModel* m, const mjData* d, int id,
+                  const mjtNum* pnt, const mjtNum* vec) {
+  // check geom type
+  if (m->geom_type[id]!=mjGEOM_MESH) {
+    mju_error("mj_rayMesh: geom with mesh type expected");
+  }
+
+  // bounding box test
+  if (ray_box(d->geom_xpos+3*id, d->geom_xmat+9*id, m->geom_size+3*id, pnt, vec, NULL)<0) {
+    return -1;
+  }
+
+  return mju_rayTree(m, d, id, pnt, vec);
+}
+
 
 
 // intersect ray with pure geom, no meshes or hfields

src/engine/engine_ray.h

@@ -46,6 +46,10 @@ MJAPI mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum* pnt, const
 MJAPI mjtNum mj_rayHfield(const mjModel* m, const mjData* d, int geomid,
                           const mjtNum* pnt, const mjtNum* vec);
 
+// intersect ray with triangle
+MJAPI mjtNum ray_triangle(mjtNum v[][3], const mjtNum* lpnt, const mjtNum* lvec,
+                          const mjtNum* b0, const mjtNum* b1);
+
 // intersect ray with mesh
 MJAPI mjtNum mj_rayMesh(const mjModel* m, const mjData* d, int geomid,
                         const mjtNum* pnt, const mjtNum* vec);

src/user/user_objects.cc

@@ -34,6 +34,7 @@
 #include "engine/engine_macro.h"
 #include "engine/engine_passive.h"
 #include "engine/engine_plugin.h"
+#include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"
 #include "engine/engine_util_solve.h"
@@ -338,6 +339,14 @@ int mjCBoundingVolumeHierarchy::MakeBVH(std::vector<mjCBoundingVolume>& elements
     }
   }
 
+  // inflate flat AABBs
+  for (int i=0; i<3; i++) {
+    if (mju_abs(AABB[i]-AABB[i+3])<mjEPS) {
+      AABB[i+0] -= mjEPS;
+      AABB[i+3] += mjEPS;
+    }
+  }
+
   // store current index
   int index = nbvh++;
   child.push_back(-1);
@@ -432,6 +441,10 @@ int mjCBoundingVolumeHierarchy::MakeBVH(std::vector<mjCBoundingVolume>& elements
               name_.c_str(), nelements);
   }
 
+  if (lev>mjMAXTREEDEPTH) {
+    mju_warning("max tree depth exceeded in body=%s", name_.c_str());
+  }
+
   return index;
 }
 

test/engine/engine_ray_test.cc

@@ -51,6 +51,28 @@ static constexpr char kRayCastingModel[] = R"(
 </mujoco>
 )";
 
+static constexpr char kCubeletModel[] = R"(
+<mujoco>
+  <asset>
+    <mesh name="cubelet"
+    vertex="0.0085  -0.01     0.0085   -0.0085  -0.01    -0.0085    0.0085  -0.01    -0.0085
+            0.01     0.0085   0.0085    0.01    -0.0085  -0.0085    0.01     0.0085  -0.0085
+           -0.0085   0.0085   0.01      0.0085  -0.0085   0.01      0.0085   0.0085   0.01
+           -0.01    -0.0085   0.0085   -0.01     0.0085  -0.0085   -0.01    -0.0085  -0.0085
+           -0.0085   0.01     0.0085    0.0085   0.01    -0.0085   -0.0085   0.01    -0.0085
+           -0.0085  -0.0085  -0.01     -0.0085  -0.01     0.0085   -0.0085  -0.0085   0.01
+           -0.0085   0.0085  -0.01     -0.01     0.0085   0.0085    0.0085  -0.0085  -0.01
+            0.01    -0.0085   0.0085    0.0085   0.0085  -0.01      0.0085   0.01     0.0085"/>
+  </asset>
+
+  <worldbody>
+    <body pos="1 0 0">
+      <geom type="mesh" mesh="cubelet"/>
+    </body>
+  </worldbody>
+</mujoco>
+)";
+
 using ::testing::NotNull;
 using RayTest = MujocoTest;
 
@@ -140,6 +162,8 @@ TEST_F(RayTest, ExcludeStatic) {
   mj_deleteModel(model);
 }
 
+// ------------------------------- mj_multiRay --------------------------------
+
 TEST_F(RayTest, MultiRayEqualsSingleRay) {
   mjModel* m = LoadModelFromString(kRayCastingModel);
   ASSERT_THAT(m, NotNull());
@@ -258,5 +282,126 @@ TEST_F(RayTest, EdgeCases) {
   mj_deleteModel(m);
 }
 
+// ------------------------------- mj_rayMesh ---------------------------------
+
+// old ray mesh intersection
+mjtNum _rayMesh(const mjModel* m, const mjData* d, int geomid,
+                const mjtNum* pnt, const mjtNum* vec) {
+  // check geom type
+  if (m->geom_type[geomid] != mjGEOM_MESH) {
+    mju_error("mj_rayMesh: geom with mesh type expected");
+  }
+
+  // map to local frame
+  mjtNum lpnt[3], lvec[3];
+  const mjtNum* pos = d->geom_xpos+3*geomid;
+  const mjtNum dif[3] = {pnt[0]-pos[0], pnt[1]-pos[1], pnt[2]-pos[2]};
+  mju_rotVecMatT(lpnt, dif, d->geom_xmat+9*geomid);
+  mju_rotVecMatT(lvec, vec, d->geom_xmat+9*geomid);
+
+  // construct basis vectors of normal plane
+  mjtNum b0[3] = {1, 1, 1}, b1[3];
+  if (mju_abs(lvec[0]) >= mju_abs(lvec[1]) &&
+      mju_abs(lvec[0]) >= mju_abs(lvec[2])) {
+    b0[0] = 0;
+  } else if (mju_abs(lvec[1]) >= mju_abs(lvec[2])) {
+    b0[1] = 0;
+  } else {
+    b0[2] = 0;
+  }
+  mju_addScl3(b1, b0, lvec, -mju_dot3(lvec, b0)/mju_dot3(lvec, lvec));
+  mju_normalize3(b1);
+  mju_cross(b0, b1, lvec);
+  mju_normalize3(b0);
+
+  // init solution
+  mjtNum x = -1, sol;
+
+  // process all triangles
+  int face, meshid = m->geom_dataid[geomid];
+  for (face = m->mesh_faceadr[meshid];
+       face < m->mesh_faceadr[meshid] + m->mesh_facenum[meshid];
+       face++) {
+    // get float vertices
+    float* vf[3];
+    vf[0] = m->mesh_vert + 3*(m->mesh_face[3*face]   + m->mesh_vertadr[meshid]);
+    vf[1] = m->mesh_vert + 3*(m->mesh_face[3*face+1] + m->mesh_vertadr[meshid]);
+    vf[2] = m->mesh_vert + 3*(m->mesh_face[3*face+2] + m->mesh_vertadr[meshid]);
+
+    // convert to mjtNum
+    mjtNum v[3][3];
+    for (int i=0; i < 3; i++) {
+      for (int j=0; j < 3; j++) {
+        v[i][j] = (mjtNum)vf[i][j];
+      }
+    }
+
+    // solve
+    sol = ray_triangle(v, lpnt, lvec, b0, b1);
+
+    // update
+    if (sol >= 0 && (x < 0 || sol < x)) {
+      x = sol;
+    }
+  }
+
+  return x;
+}
+
+// performs a ray mesh test using a given mjModel
+void _rayMeshTest(const mjModel* m) {
+  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, 0};
+  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);
+      }
+    }
+  }
+  // compare results with single ray function
+  mjtNum dist_new, dist_old;
+
+  for (int i = 0; i < N; ++i) {
+    for (int j = 0; j < M; ++j) {
+      int idx = i * M + j;
+      dist_old = _rayMesh(m, d, /*geomid=*/0, pnt, vec + 3 * idx);
+      dist_new = mj_rayMesh(m, d, /*geomid=*/0, pnt, vec + 3 * idx);
+      EXPECT_FLOAT_EQ(dist_new, dist_old);
+    }
+  }
+
+  mj_deleteData(d);
+}
+
+TEST_F(RayTest, RayMeshPruning) {
+  char error[1024] = {0};
+  const std::string xml_path =
+      GetTestDataFilePath("engine/testdata/ray/stanford_bunny.xml");
+
+  mjModel* m = mj_loadXML(xml_path.c_str(), NULL, error, sizeof(error));
+  ASSERT_THAT(m, NotNull()) << error;
+  _rayMeshTest(m);
+  mj_deleteModel(m);
+
+  m = LoadModelFromString(kCubeletModel);
+  ASSERT_THAT(m, NotNull());
+  _rayMeshTest(m);
+  mj_deleteModel(m);
+}
+
 }  // namespace
 }  // namespace mujoco

test/engine/testdata/ray/stanford_bunny.xml

@@ -0,0 +1,8 @@
+<mujoco>
+  <asset>
+    <mesh file="stanford_bunny.obj"/>
+  </asset>
+  <worldbody>
+    <geom type="mesh" mesh="stanford_bunny" pos="1 0 0" euler="0 90 0"/>
+  </worldbody>
+</mujoco>

unity/Runtime/Bindings/MjBindings.cs

@@ -39,6 +39,7 @@ public const double mjMINMU = 1e-05;
 public const double mjMINIMP = 0.0001;
 public const double mjMAXIMP = 0.9999;
 public const int mjMAXCONPAIR = 50;
+public const int mjMAXTREEDEPTH = 50;
 public const int mjMAXVFS = 2000;
 public const int mjMAXVFSNAME = 1000;
 public const int mjNEQDATA = 11;