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
2025-07-09 06:20:36 -07:00
parent e6c5715903
commit b81f1db8af
8 changed files with 194 additions and 121 deletions
@@ -58,9 +58,9 @@ set(MUJOCO_DEP_VERSION_benchmark
    CACHE STRING "Version of `benchmark` to be fetched."
 )

-set(MUJOCO_DEP_VERSION_sdflib
-    1927bee6bb8225258a39c8cbf14e18a4d50409ae
-    CACHE STRING "Version of `SdfLib` to be fetched."
+set(MUJOCO_DEP_VERSION_TriangleMeshDistance
+    2cb643de1436e1ba8e2be49b07ec5491ac604457
+    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)
-option(SDFLIB_USE_OPENMP OFF)
-option(SDFLIB_USE_ENOKI OFF)
-findorfetch(
-  USE_SYSTEM_PACKAGE
-  OFF
-  PACKAGE_NAME
-  sdflib
-  LIBRARY_NAME
-  sdflib
-  GIT_REPO
-  https://github.com/UPC-ViRVIG/SdfLib.git
-  GIT_TAG
-  ${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})
+if(NOT TARGET trianglemeshdistance)
+  FetchContent_Declare(
+    trianglemeshdistance
+    GIT_REPOSITORY https://github.com/InteractiveComputerGraphics/TriangleMeshDistance.git
+    GIT_TAG ${MUJOCO_DEP_VERSION_TriangleMeshDistance}
+  )
+
+  FetchContent_GetProperties(trianglemeshdistance)
+  if(NOT trianglemeshdistance_POPULATED)
+    FetchContent_Populate(trianglemeshdistance)
+    include_directories(${trianglemeshdistance_SOURCE_DIR})
+  endif()
+endif()

 set(ENABLE_DOUBLE_PRECISION ON)
 set(CCD_HIDE_ALL_SYMBOLS ON)
@@ -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>
@@ -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}
@@ -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
-user-specified mesh. The mesh can be arbitrary and not necessarily convex. This offers an alternative to
-convex-decomposed meshes. The performance is likely to be slower than that of analytic SDFs, since a cubic
-approximation has to be evaluated on the convex grid. However, the SDF generation is done automatically, simplifying the
-task of creating an SDF, which can be difficult for complex shapes.
+This plugin uses the library [TriangleMeshDistance](https://github.com/InteractiveComputerGraphics/TriangleMeshDistance)
+to compute a voxel-based approximation of a user-specified mesh. The mesh can be arbitrary and not necessarily convex.
+This offers an alternative to convex-decomposed meshes. The performance is likely to be slower than that of analytic
+SDFs, since a cubic approximation has to be evaluated on the convex grid. However, the SDF generation is done
+automatically, simplifying the task of creating an SDF, which can be difficult for complex shapes.

 ### How to make your own SDF

@@ -13,12 +13,12 @@
 // limitations under the License.

 #include <cstdint>
+#include <cstring>
 #include <optional>
 #include <utility>
 #include <vector>

-#include <SdfLib/utils/Mesh.h>
-#include <SdfLib/OctreeSdf.h>
+#include <TriangleMeshDistance/include/tmd/TriangleMeshDistance.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) {
-  return reinterpret_cast<unsigned int*>(const_cast<int*>(ptr));
-}
-
-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 boxProjection(mjtNum point[3], const mjtNum box[6]) {
+  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 dist_sqr = 0;
  mjtNum eps = 1e-6;

  // skip the projection if inside
-  if (q.x <= 0 && q.y <= 0 && q.z <= 0) {
-    return glm::max(q.x, glm::max(q.y, q.z));
+  if (q[0] <= 0 && q[1] <= 0 && q[2] <= 0) {
+    return mju_max(q[0], mju_max(q[1], q[2]));
  }

  // in-place projection inside the box if outside
-  if ( q.x >= 0 ) {
-    dist_sqr += q.x * q.x;
-    point.x -= r.x > 0 ? (q.x+eps) : -(q.x+eps);
+  if ( q[0] >= 0 ) {
+    dist_sqr += q[0] * q[0];
+    point[0] -= r[0] > 0 ? (q[0]+eps) : -(q[0]+eps);
  }
-  if ( q.y >= 0 ) {
-    dist_sqr += q.y * q.y;
-    point.y -= r.y > 0 ? (q.y+eps) : -(q.y+eps);
+  if ( q[1] >= 0 ) {
+    dist_sqr += q[1] * q[1];
+    point[1] -= r[1] > 0 ? (q[1]+eps) : -(q[1]+eps);
  }
-  if ( q.z >= 0 ) {
-    dist_sqr += q.z * q.z;
-    point.z -= r.z > 0 ? (q.z+eps) : -(q.z+eps);
+  if ( q[2] >= 0 ) {
+    dist_sqr += q[2] * q[2];
+    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[i].y = vert[1];
-    vertices[i].z = vert[2];
+    vertices[3*i+0] = vert[0];
+    vertices[3*i+1] = vert[1];
+    vertices[3*i+2] = vert[2];
  }
-  sdflib::Mesh mesh(vertices.data(), nvert,
-                    MakeNonConstUnsigned(indices), 3*nface);
-  mesh.computeBoundingBox();
-  return SdfLib(std::move(mesh));
+  tmd::TriangleMeshDistance mesh(vertices.data(), nvert, indices, nface);
+  return SdfLib(mesh, m, meshid);
 }

 // plugin constructor
-SdfLib::SdfLib(sdflib::Mesh&& mesh) {
-  sdflib::BoundingBox box = mesh.getBoundingBox();
-  const glm::vec3 modelBBsize = box.getSize();
-  box.addMargin(
-      0.1f * glm::max(glm::max(modelBBsize.x, modelBBsize.y), modelBBsize.z));
-  sdf_func_ =
-      sdflib::OctreeSdf(mesh, box, 8, 3, 1e-3,
-                        sdflib::OctreeSdf::InitAlgorithm::CONTINUITY, 1);
+SdfLib::SdfLib(const tmd::TriangleMeshDistance& sdf, const mjModel* m,
+               int meshid) {
+  // TODO: do not evaluate the SDF multiple times at the same vertex
+  // TODO: the value at hanging vertices should be computed from the parent
+  int octadr = m->mesh_octadr[meshid];
+  int octnum = m->mesh_octnum[meshid];
+  oct_aabb_.assign(m->oct_aabb + 6*octadr,
+                   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 boxDist = boxProjection(point, sdf_func_.getGridBoundingBox());
-  return sdf_func_.getDistance(point) + (boxDist <= 0 ? 0 : boxDist);
+  mjtNum w[8];
+  mjtNum sdf = 0;
+  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) {
-    sdf_func_.getDistance(p, gradient);
-    grad[0] = gradient[0];
-    grad[1] = gradient[1];
-    grad[2] = gradient[2];
+  if (boxProjection(p, box_) <= 0) {
+    mjtNum dw[8][3];
+    mju_zero3(grad);
+    int node = findOct(nullptr, dw, oct_aabb_.data(), oct_child_.data(), p);
+    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;
  }

@@ -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
@@ -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];
@@ -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);
 }