Replace SdfLib with linear octree interpolation of TriangleMeshDistance.

Temporary changes to the octree:
- Changed frame from mesh to geom

Before change (tolerance 1e-3):
```
 Simulation time      : 2.34 s
 Steps per second     : 4275
 Realtime factor      : 8.55 x
 Time per step        : 233.9 µs

 Newton iters / step  : 2.26
 Contacts / step      : 3.69
 Constraints / step   : 14.76
 Degrees of freedom   : 12
 Dynamic memory usage : 0.7% of 14M
```

After change (6 octree levels):
```
 Simulation time      : 2.24 s
 Steps per second     : 4472
 Realtime factor      : 8.94 x
 Time per step        : 223.6 µs

 Newton iters / step  : 2.47
 Contacts / step      : 3.37
 Constraints / step   : 13.49
 Degrees of freedom   : 12
 Dynamic memory usage : 0.2% of 14M
```

PiperOrigin-RevId: 781019754
Change-Id: Ib15581244dfe9e571e1c6c2cac4a101cf3d8ba3d

cmake/MujocoDependencies.cmake

@@ -58,9 +58,9 @@ set(MUJOCO_DEP_VERSION_benchmark
     CACHE STRING "Version of `benchmark` to be fetched."
 )
 
-set(MUJOCO_DEP_VERSION_sdflib
+set(MUJOCO_DEP_VERSION_TriangleMeshDistance
-    1927bee6bb8225258a39c8cbf14e18a4d50409ae
+    2cb643de1436e1ba8e2be49b07ec5491ac604457
-    CACHE STRING "Version of `SdfLib` to be fetched."
+    CACHE STRING "Version of `TriangleMeshDistance` to be fetched."
 )
 
 mark_as_advanced(MUJOCO_DEP_VERSION_lodepng)
@@ -73,7 +73,7 @@ mark_as_advanced(MUJOCO_DEP_VERSION_Eigen3)
 mark_as_advanced(MUJOCO_DEP_VERSION_abseil)
 mark_as_advanced(MUJOCO_DEP_VERSION_gtest)
 mark_as_advanced(MUJOCO_DEP_VERSION_benchmark)
-mark_as_advanced(MUJOCO_DEP_VERSION_sdflib)
+mark_as_advanced(MUJOCO_DEP_VERSION_TriangleMeshDistance)
 
 include(FetchContent)
 include(FindOrFetch)
@@ -184,26 +184,19 @@ findorfetch(
   EXCLUDE_FROM_ALL
 )
 
-option(SDFLIB_USE_ASSIMP OFF)
+if(NOT TARGET trianglemeshdistance)
-option(SDFLIB_USE_OPENMP OFF)
+  FetchContent_Declare(
-option(SDFLIB_USE_ENOKI OFF)
+    trianglemeshdistance
-findorfetch(
+    GIT_REPOSITORY https://github.com/InteractiveComputerGraphics/TriangleMeshDistance.git
-  USE_SYSTEM_PACKAGE
+    GIT_TAG ${MUJOCO_DEP_VERSION_TriangleMeshDistance}
-  OFF
+  )
-  PACKAGE_NAME
+
-  sdflib
+  FetchContent_GetProperties(trianglemeshdistance)
-  LIBRARY_NAME
+  if(NOT trianglemeshdistance_POPULATED)
-  sdflib
+    FetchContent_Populate(trianglemeshdistance)
-  GIT_REPO
+    include_directories(${trianglemeshdistance_SOURCE_DIR})
-  https://github.com/UPC-ViRVIG/SdfLib.git
+  endif()
-  GIT_TAG
+endif()
-  ${MUJOCO_DEP_VERSION_sdflib}
-  TARGETS
-  SdfLib
-  EXCLUDE_FROM_ALL
-)
-target_compile_options(SdfLib PRIVATE ${MUJOCO_MACOS_COMPILE_OPTIONS})
-target_link_options(SdfLib PRIVATE ${MUJOCO_MACOS_LINK_OPTIONS})
 
 set(ENABLE_DOUBLE_PRECISION ON)
 set(CCD_HIDE_ALL_SYMBOLS ON)

model/plugin/sdf/octree.xml

@@ -1,31 +0,0 @@
-<mujoco>
-  <compiler texturedir="asset"/>
-
-  <asset>
-    <texture name="texspot" type="2d" file="spot.png"/>
-    <material name="matspot" texture="texspot"/>
-    <mesh name="spot" file="asset/spot.obj"/>
-  </asset>
-
-  <option sdf_iterations="20" sdf_initpoints="40"/>
-
-  <visual>
-    <map force="1000"/>
-  </visual>
-
-  <default>
-    <geom solref="0.01 1" solimp=".95 .99 .0001" friction="0.5"/>
-  </default>
-
-  <statistic meansize="0.2"/>
-
-  <include file="scene.xml"/>
-
-  <worldbody>
-    <body euler="90 0 0" pos="0 0 .7">
-      <geom type="sdf" name="cow1" mesh="spot" material="matspot"/>
-    </body>
-    <light name="left" pos="0 0 1"/>
-    <light name="right" pos="1 0 1"/>
-  </worldbody>
-</mujoco>

plugin/sdf/CMakeLists.txt

@@ -37,7 +37,7 @@ set(MUJOCO_SDF_SRCS
 add_library(sdf_plugin SHARED)
 target_sources(sdf_plugin PRIVATE ${MUJOCO_SDF_SRCS})
 target_include_directories(sdf_plugin PRIVATE ${MUJOCO_SDF_INCLUDE})
-target_link_libraries(sdf_plugin PRIVATE mujoco SdfLib)
+target_link_libraries(sdf_plugin PRIVATE mujoco)
 target_compile_options(
   sdf_plugin
   PRIVATE ${AVX_COMPILE_OPTIONS}

plugin/sdf/README.md

@@ -66,11 +66,11 @@ Parameters:
 
 Implemented in [sdflib.cc](sdflib.cc). Example usage in [cow.xml](../../model/plugin/sdf/cow.xml).
 
-This plugin uses the library [SdfLib](https://github.com/UPC-ViRVIG/SdfLib) to compute a voxel-based approximation of a
+This plugin uses the library [TriangleMeshDistance](https://github.com/InteractiveComputerGraphics/TriangleMeshDistance)
-user-specified mesh. The mesh can be arbitrary and not necessarily convex. This offers an alternative to
+to compute a voxel-based approximation of a user-specified mesh. The mesh can be arbitrary and not necessarily convex.
-convex-decomposed meshes. The performance is likely to be slower than that of analytic SDFs, since a cubic
+This offers an alternative to convex-decomposed meshes. The performance is likely to be slower than that of analytic
-approximation has to be evaluated on the convex grid. However, the SDF generation is done automatically, simplifying the
+SDFs, since a cubic approximation has to be evaluated on the convex grid. However, the SDF generation is done
-task of creating an SDF, which can be difficult for complex shapes.
+automatically, simplifying the task of creating an SDF, which can be difficult for complex shapes.
 
 ### How to make your own SDF
 

plugin/sdf/sdflib.cc

@@ -13,12 +13,12 @@
 // limitations under the License.
 
 #include <cstdint>
+#include <cstring>
 #include <optional>
 #include <utility>
 #include <vector>
 
-#include <SdfLib/utils/Mesh.h>
+#include <TriangleMeshDistance/include/tmd/TriangleMeshDistance.h>
-#include <SdfLib/OctreeSdf.h>
 #include <mujoco/mjplugin.h>
 #include <mujoco/mujoco.h>
 #include "sdf.h"
@@ -27,38 +27,107 @@
 namespace mujoco::plugin::sdf {
 namespace {
 
-inline unsigned int* MakeNonConstUnsigned(const int* ptr) {
+mjtNum boxProjection(mjtNum point[3], const mjtNum box[6]) {
-  return reinterpret_cast<unsigned int*>(const_cast<int*>(ptr));
+  mjtNum r[3] = {point[0] - box[0], point[1] - box[1], point[2] - box[2]};
-}
+  mjtNum q[3] = {mju_abs(r[0]) - box[3], mju_abs(r[1]) - box[4],
-
+                 mju_abs(r[2]) - box[5]};
-mjtNum boxProjection(glm::vec3& point, const sdflib::BoundingBox& box) {
-  glm::vec3 r = point - box.getCenter();
-  glm::vec3 q = glm::abs(r) - 0.5f * box.getSize();
   mjtNum dist_sqr = 0;
   mjtNum eps = 1e-6;
 
   // skip the projection if inside
-  if (q.x <= 0 && q.y <= 0 && q.z <= 0) {
+  if (q[0] <= 0 && q[1] <= 0 && q[2] <= 0) {
-    return glm::max(q.x, glm::max(q.y, q.z));
+    return mju_max(q[0], mju_max(q[1], q[2]));
   }
 
   // in-place projection inside the box if outside
-  if ( q.x >= 0 ) {
+  if ( q[0] >= 0 ) {
-    dist_sqr += q.x * q.x;
+    dist_sqr += q[0] * q[0];
-    point.x -= r.x > 0 ? (q.x+eps) : -(q.x+eps);
+    point[0] -= r[0] > 0 ? (q[0]+eps) : -(q[0]+eps);
   }
-  if ( q.y >= 0 ) {
+  if ( q[1] >= 0 ) {
-    dist_sqr += q.y * q.y;
+    dist_sqr += q[1] * q[1];
-    point.y -= r.y > 0 ? (q.y+eps) : -(q.y+eps);
+    point[1] -= r[1] > 0 ? (q[1]+eps) : -(q[1]+eps);
   }
-  if ( q.z >= 0 ) {
+  if ( q[2] >= 0 ) {
-    dist_sqr += q.z * q.z;
+    dist_sqr += q[2] * q[2];
-    point.z -= r.z > 0 ? (q.z+eps) : -(q.z+eps);
+    point[2] -= r[2] > 0 ? (q[2]+eps) : -(q[2]+eps);
   }
 
   return mju_sqrt(dist_sqr);
 }
 
+// find the octree leaf containing the point p, return the index of the leaf and
+// populate the weights of the interpolated function (if w is not null) and of
+// its gradient (if dw is not null) using the vertices as degrees of freedom for
+// trilinear interpolation.
+static int findOct(mjtNum w[8], mjtNum dw[8][3], const mjtNum* oct_aabb,
+                   const int* oct_child, const mjtNum p[3]) {
+  std::vector<int> stack = {0};
+  mjtNum eps = 1e-8;
+
+  while (!stack.empty()) {
+    int node = stack.back();
+    stack.pop_back();
+    mjtNum vmin[3], vmax[3];
+
+    if (node == -1) {  // SHOULD NOT OCCUR
+      mju_error("Invalid node number");
+      return -1;
+    }
+
+    for (int j = 0; j < 3; j++) {
+      vmin[j] = oct_aabb[6*node+j] - oct_aabb[6*node+3+j];
+      vmax[j] = oct_aabb[6*node+j] + oct_aabb[6*node+3+j];
+    }
+
+    // check if the point is inside the aabb of the octree node
+    if (p[0] + eps < vmin[0] || p[0] - eps > vmax[0] ||
+        p[1] + eps < vmin[1] || p[1] - eps > vmax[1] ||
+        p[2] + eps < vmin[2] || p[2] - eps > vmax[2]) {
+      continue;
+    }
+
+    mjtNum coord[3] = {(p[0] - vmin[0]) / (vmax[0] - vmin[0]),
+                       (p[1] - vmin[1]) / (vmax[1] - vmin[1]),
+                       (p[2] - vmin[2]) / (vmax[2] - vmin[2])};
+
+    // check if the node is a leaf
+    if (oct_child[8*node+0] == -1 && oct_child[8*node+1] == -1 &&
+        oct_child[8*node+2] == -1 && oct_child[8*node+3] == -1 &&
+        oct_child[8*node+4] == -1 && oct_child[8*node+5] == -1 &&
+        oct_child[8*node+6] == -1 && oct_child[8*node+7] == -1) {
+      for (int j = 0; j < 8; j++) {
+        if (w) {
+          w[j] = (j & 1 ? coord[0] : 1 - coord[0]) *
+                 (j & 2 ? coord[1] : 1 - coord[1]) *
+                 (j & 4 ? coord[2] : 1 - coord[2]);
+        }
+        if (dw) {
+          dw[j][0] = (j & 1 ? 1 : -1) *
+                     (j & 2 ? coord[1] : 1 - coord[1]) *
+                     (j & 4 ? coord[2] : 1 - coord[2]);
+          dw[j][1] = (j & 1 ? coord[0] : 1 - coord[0]) *
+                     (j & 2 ? 1 : -1) *
+                     (j & 4 ? coord[2] : 1 - coord[2]);
+          dw[j][2] = (j & 1 ? coord[0] : 1 - coord[0]) *
+                     (j & 2 ? coord[1] : 1 - coord[1]) *
+                     (j & 4 ? 1 : -1);
+        }
+      }
+      return node;
+    }
+
+    // compute which of 8 children to visit next
+    int x = coord[0] < .5 ? 1 : 0;
+    int y = coord[1] < .5 ? 1 : 0;
+    int z = coord[2] < .5 ? 1 : 0;
+    stack.push_back(oct_child[8*node + 4*z + 2*y + x]);
+  }
+
+  mju_error("Node not found");  // SHOULD NOT OCCUR
+  return -1;
+}
+
 }  // namespace
 
 // factory function
@@ -76,30 +145,40 @@ std::optional<SdfLib> SdfLib::Create(const mjModel* m, mjData* d,
   int nface = m->mesh_facenum[meshid];
   int* indices = m->mesh_face + 3*m->mesh_faceadr[meshid];
   float* verts = m->mesh_vert + 3*m->mesh_vertadr[meshid];
-  std::vector<glm::vec3> vertices(nvert);
+  std::vector<double> vertices(3*nvert);
   for (int i = 0; i < nvert; i++) {
     mjtNum vert[3] = {verts[3*i+0], verts[3*i+1], verts[3*i+2]};
     mju_rotVecQuat(vert, vert, m->mesh_quat + 4*meshid);
     mju_addTo3(vert, m->mesh_pos + 3*meshid);
-    vertices[i].x = vert[0];
+    vertices[3*i+0] = vert[0];
-    vertices[i].y = vert[1];
+    vertices[3*i+1] = vert[1];
-    vertices[i].z = vert[2];
+    vertices[3*i+2] = vert[2];
   }
-  sdflib::Mesh mesh(vertices.data(), nvert,
+  tmd::TriangleMeshDistance mesh(vertices.data(), nvert, indices, nface);
-                    MakeNonConstUnsigned(indices), 3*nface);
+  return SdfLib(mesh, m, meshid);
-  mesh.computeBoundingBox();
-  return SdfLib(std::move(mesh));
 }
 
 // plugin constructor
-SdfLib::SdfLib(sdflib::Mesh&& mesh) {
+SdfLib::SdfLib(const tmd::TriangleMeshDistance& sdf, const mjModel* m,
-  sdflib::BoundingBox box = mesh.getBoundingBox();
+               int meshid) {
-  const glm::vec3 modelBBsize = box.getSize();
+  // TODO: do not evaluate the SDF multiple times at the same vertex
-  box.addMargin(
+  // TODO: the value at hanging vertices should be computed from the parent
-      0.1f * glm::max(glm::max(modelBBsize.x, modelBBsize.y), modelBBsize.z));
+  int octadr = m->mesh_octadr[meshid];
-  sdf_func_ =
+  int octnum = m->mesh_octnum[meshid];
-      sdflib::OctreeSdf(mesh, box, 8, 3, 1e-3,
+  oct_aabb_.assign(m->oct_aabb + 6*octadr,
-                        sdflib::OctreeSdf::InitAlgorithm::CONTINUITY, 1);
+                   m->oct_aabb + 6*octadr + 6*octnum);
+  oct_child_.assign(m->oct_child + 8 * octadr,
+                    m->oct_child + 8 * octadr + 8 * octnum);
+  for (int i = 0; i < octnum; ++i) {
+    for (int j = 0; j < 8; j++) {
+        mjtNum v[3];
+        v[0] = oct_aabb_[6*i+0] + (j&1 ? 1 : -1) * oct_aabb_[6*i+3];
+        v[1] = oct_aabb_[6*i+1] + (j&2 ? 1 : -1) * oct_aabb_[6*i+4];
+        v[2] = oct_aabb_[6*i+2] + (j&4 ? 1 : -1) * oct_aabb_[6*i+5];
+        sdf_coeff_.push_back(sdf.signed_distance(v).distance);
+    }
+  }
+  mju_copy(box_, m->oct_aabb + 6*octadr, 6);
 }
 
 // plugin computation
@@ -120,22 +199,35 @@ void SdfLib::Visualize(const mjModel* m, mjData* d, const mjvOption* opt,
 
 // sdf
 mjtNum SdfLib::Distance(const mjtNum p[3]) const {
-  glm::vec3 point(p[0], p[1], p[2]);
+  mjtNum w[8];
-  mjtNum boxDist = boxProjection(point, sdf_func_.getGridBoundingBox());
+  mjtNum sdf = 0;
-  return sdf_func_.getDistance(point) + (boxDist <= 0 ? 0 : boxDist);
+  mjtNum point[3] = {p[0], p[1], p[2]};
+  mjtNum boxDist = boxProjection(point, box_);
+  if (boxDist > 0) {
+    return boxDist;
+  }
+  int node = findOct(w, nullptr, oct_aabb_.data(), oct_child_.data(), point);
+  for (int i = 0; i < 8; ++i) {
+    sdf += w[i] * sdf_coeff_[8*node + i];
+  }
+  return sdf;
 }
 
 // gradient of sdf
 void SdfLib::Gradient(mjtNum grad[3], const mjtNum point[3]) const {
-  glm::vec3 gradient;
+
-  glm::vec3 p(point[0], point[1], point[2]);
+  mjtNum p[3] = {point[0], point[1], point[2]};
 
   // analytic in the interior
-  if (boxProjection(p, sdf_func_.getGridBoundingBox()) <= 0) {
+  if (boxProjection(p, box_) <= 0) {
-    sdf_func_.getDistance(p, gradient);
+    mjtNum dw[8][3];
-    grad[0] = gradient[0];
+    mju_zero3(grad);
-    grad[1] = gradient[1];
+    int node = findOct(nullptr, dw, oct_aabb_.data(), oct_child_.data(), p);
-    grad[2] = gradient[2];
+    for (int i = 0; i < 8; ++i) {
+      grad[0] += dw[i][0] * sdf_coeff_[8*node + i];
+      grad[1] += dw[i][1] * sdf_coeff_[8*node + i];
+      grad[2] += dw[i][2] * sdf_coeff_[8*node + i];
+    }
     return;
   }
 

plugin/sdf/sdflib.h

@@ -16,14 +16,14 @@
 #define MUJOCO_PLUGIN_SDF_SDFLIB_H_
 
 #include <optional>
+#include <vector>
 
-#include <SdfLib/utils/Mesh.h>
-#include <SdfLib/OctreeSdf.h>
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjtnum.h>
 #include <mujoco/mjvisualize.h>
 #include "sdf.h"
+#include <TriangleMeshDistance/include/tmd/TriangleMeshDistance.h>
 
 namespace mujoco::plugin::sdf {
 class SdfLib {
@@ -44,9 +44,13 @@ class SdfLib {
   static void RegisterPlugin();
 
  private:
-  SdfLib(sdflib::Mesh&& mesh);
+  SdfLib(const tmd::TriangleMeshDistance& sdf, const mjModel* m, int meshid);
   SdfVisualizer visualizer_;
-  sdflib::OctreeSdf sdf_func_;
+  std::vector<double> sdf_coeff_;
+
+  mjtNum box_[6];
+  std::vector<mjtNum> oct_aabb_;
+  std::vector<int> oct_child_;
 };
 
 }  // namespace mujoco::plugin::sdf

src/user/user_mesh.cc

@@ -628,7 +628,8 @@ void mjCMesh::TryCompile(const mjVFS* vfs) {
     resource_ = LoadResource(modelfiledir_.Str(), filename.Str(), vfs);
 
     // try loading from cache
-    if (cache != nullptr && LoadCachedMesh(cache, resource_)) {
+    // TODO: move octree to mesh frame so it can be cached
+    if (cache != nullptr && !needoct_ && LoadCachedMesh(cache, resource_)) {
       mju_closeResource(resource_);
       resource_ = nullptr;
       fromCache = true;
@@ -684,14 +685,6 @@ void mjCMesh::TryCompile(const mjVFS* vfs) {
     }
   }
 
-  // make octree
-  if (!needoct_) {
-    octree_.Clear();
-  } else if (octree_.Nodes().empty()) {
-    octree_.SetFace(vert_, face_);
-    octree_.CreateOctree(aamm_);
-  }
-
   // close resource
   if (resource_ != nullptr) {
     mju_closeResource(resource_);
@@ -1539,6 +1532,24 @@ void mjCMesh::Process() {
   boxsz_[1] = 0.5 * std::sqrt(6*(eigval[0] + eigval[2] - eigval[1])/volume);
   boxsz_[2] = 0.5 * std::sqrt(6*(eigval[0] + eigval[1] - eigval[2])/volume);
 
+  // make octree in the geom frame
+  // TODO: make octree in the mesh frame, update engine_collision_sdf
+  if (!needoct_) {
+    octree_.Clear();
+  } else if (octree_.Nodes().empty()) {
+    double aamm[6] = {mjMAXVAL, mjMAXVAL, mjMAXVAL, -mjMAXVAL, -mjMAXVAL, -mjMAXVAL};
+    for (int i = 0; i < nvert(); i++) {
+      aamm[0] = std::min(aamm[0], vert_[3*i + 0]);
+      aamm[3] = std::max(aamm[3], vert_[3*i + 0]);
+      aamm[1] = std::min(aamm[1], vert_[3*i + 1]);
+      aamm[4] = std::max(aamm[4], vert_[3*i + 1]);
+      aamm[2] = std::min(aamm[2], vert_[3*i + 2]);
+      aamm[5] = std::max(aamm[5], vert_[3*i + 2]);
+    }
+    octree_.SetFace(vert_, face_);
+    octree_.CreateOctree(aamm);
+  }
+
   // transform CoM to origin
   for (int i=0; i < nvert(); i++) {
     vert_[3*i + 0] -= CoM[0];

src/user/user_objects.cc

@@ -583,12 +583,16 @@ void mjCOctree::Make(std::vector<Triangle>& elements) {
 
 
 void mjCOctree::CreateOctree(const double aamm[6]) {
+  double aabb[6] = {(aamm[0] + aamm[3]) / 2, (aamm[1] + aamm[4]) / 2, (aamm[2] + aamm[5]) / 2,
+                    (aamm[3] - aamm[0]) / 2, (aamm[4] - aamm[1]) / 2, (aamm[5] - aamm[2]) / 2};
+  double box[6] = {aabb[0] - 1.1 * aabb[3], aabb[1] - 1.1 * aabb[4], aabb[2] - 1.1 * aabb[5],
+                   aabb[0] + 1.1 * aabb[3], aabb[1] + 1.1 * aabb[4], aabb[2] + 1.1 * aabb[5]};
   std::vector<Triangle> elements;
   Make(elements);
   std::vector<Triangle*> elements_ptrs(elements.size());
   std::transform(elements.begin(), elements.end(), elements_ptrs.begin(),
                  [](Triangle& triangle) { return &triangle; });
-  MakeOctree(elements_ptrs, aamm);
+  MakeOctree(elements_ptrs, box);
 }