Move SDF precomputation to the mujoco compiler.

Before:
```
 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
```

After:
```
 Simulation time      : 1.71 s
 Steps per second     : 5854
 Realtime factor      : 11.71 x
 Time per step        : 170.8 µs

 Newton iters / step  : 2.19
 Contacts / step      : 3.45
 Constraints / step   : 13.79
 Degrees of freedom   : 12
 Dynamic memory usage : 0.2% of 14M
```

PiperOrigin-RevId: 781075182
Change-Id: Ie509047ff581ab0df4b10bbd394c27d350ab5a13

doc/includes/references.h

@@ -1098,6 +1098,7 @@ struct mjModel_ {
   int*      oct_depth;            // depth in the octree                      (noct x 1)
   int*      oct_child;            // children of octree node                  (noct x 8)
   mjtNum*   oct_aabb;             // octree node bounding box (center, size)  (noct x 6)
+  mjtNum*   oct_coeff;            // octree interpolation coefficients        (noct x 8)
 
   // joints
   int*      jnt_type;             // type of joint (mjtJoint)                 (njnt x 1)

include/mujoco/mjmodel.h

@@ -771,6 +771,7 @@ struct mjModel_ {
   int*      oct_depth;            // depth in the octree                      (noct x 1)
   int*      oct_child;            // children of octree node                  (noct x 8)
   mjtNum*   oct_aabb;             // octree node bounding box (center, size)  (noct x 6)
+  mjtNum*   oct_coeff;            // octree interpolation coefficients        (noct x 8)
 
   // joints
   int*      jnt_type;             // type of joint (mjtJoint)                 (njnt x 1)

include/mujoco/mjxmacro.h

@@ -220,6 +220,7 @@
     X   ( int,     oct_depth,             noct,          1                    ) \
     X   ( int,     oct_child,             noct,          8                    ) \
     X   ( mjtNum,  oct_aabb,              noct,          6                    ) \
+    X   ( mjtNum,  oct_coeff,             noct,          8                    ) \
     X   ( int,     jnt_type,              njnt,          1                    ) \
     X   ( int,     jnt_qposadr,           njnt,          1                    ) \
     X   ( int,     jnt_dofadr,            njnt,          1                    ) \

plugin/sdf/sdflib.cc

@@ -18,7 +18,6 @@
 #include <utility>
 #include <vector>
 
-#include <TriangleMeshDistance/include/tmd/TriangleMeshDistance.h>
 #include <mujoco/mjplugin.h>
 #include <mujoco/mujoco.h>
 #include "sdf.h"
@@ -140,44 +139,19 @@ std::optional<SdfLib> SdfLib::Create(const mjModel* m, mjData* d,
       break;
     }
   }
-  int meshid = m->geom_dataid[geomid];
-  int nvert = m->mesh_vertnum[meshid];
-  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<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[3*i+0] = vert[0];
-    vertices[3*i+1] = vert[1];
-    vertices[3*i+2] = vert[2];
-  }
-  tmd::TriangleMeshDistance mesh(vertices.data(), nvert, indices, nface);
-  return SdfLib(mesh, m, meshid);
+  return SdfLib(m, m->geom_dataid[geomid]);
 }
 
 // plugin constructor
-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
+SdfLib::SdfLib(const mjModel* m, int meshid) {
   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);
-    }
-  }
+  sdf_coeff_.assign(8 * octnum, 0);
+  memcpy(sdf_coeff_.data(), m->oct_coeff + 8*octadr, 8*octnum*sizeof(mjtNum));
   mju_copy(box_, m->oct_aabb + 6*octadr, 6);
 }
 

plugin/sdf/sdflib.h

@@ -23,7 +23,6 @@
 #include <mujoco/mjtnum.h>
 #include <mujoco/mjvisualize.h>
 #include "sdf.h"
-#include <TriangleMeshDistance/include/tmd/TriangleMeshDistance.h>
 
 namespace mujoco::plugin::sdf {
 class SdfLib {
@@ -44,7 +43,7 @@ class SdfLib {
   static void RegisterPlugin();
 
  private:
-  SdfLib(const tmd::TriangleMeshDistance& sdf, const mjModel* m, int meshid);
+  SdfLib(const mjModel* m, int meshid);
   SdfVisualizer visualizer_;
   std::vector<double> sdf_coeff_;
 

python/mujoco/introspect/structs.py

@@ -1589,6 +1589,14 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  doc='octree node bounding box (center, size)',
                  array_extent=('noct', 6),
              ),
+             StructFieldDecl(
+                 name='oct_coeff',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjtNum'),
+                 ),
+                 doc='octree interpolation coefficients',
+                 array_extent=('noct', 8),
+             ),
              StructFieldDecl(
                  name='jnt_type',
                  type=PointerType(

src/user/user_mesh.cc

@@ -32,6 +32,7 @@
 
 #include <mujoco/mjspec.h>
 #include "user/user_api.h"
+#include <TriangleMeshDistance/include/tmd/TriangleMeshDistance.h>
 
 #ifdef MUJOCO_TINYOBJLOADER_IMPL
 #define TINYOBJLOADER_IMPLEMENTATION
@@ -1548,6 +1549,25 @@ void mjCMesh::Process() {
     }
     octree_.SetFace(vert_, face_);
     octree_.CreateOctree(aamm);
+
+    // compute sdf coefficients
+    // TODO: only check !plugin.active once sdflib is removed
+    if (plugin.active && *plugin.name == "sdf") {
+      tmd::TriangleMeshDistance sdf(vert_.data(), nvert(), face_.data(), nface());
+
+      // TODO: do not evaluate the SDF multiple times at the same vertex
+      // TODO: the value at hanging vertices should be computed from the parent
+      const double* nodes = octree_.Nodes().data();
+      for (int i = 0; i < octree_.NumNodes(); ++i) {
+        for (int j = 0; j < 8; j++) {
+            mjtNum v[3];
+            v[0] = nodes[6*i+0] + (j&1 ? 1 : -1) * nodes[6*i+3];
+            v[1] = nodes[6*i+1] + (j&2 ? 1 : -1) * nodes[6*i+4];
+            v[2] = nodes[6*i+2] + (j&4 ? 1 : -1) * nodes[6*i+5];
+            octree_.AddCoeff(sdf.signed_distance(v).distance);
+        }
+      }
+    }
   }
 
   // transform CoM to origin

src/user/user_model.cc

@@ -2997,6 +2997,11 @@ void mjCModel::CopyObjects(mjModel* m) {
       memcpy(m->oct_aabb + 6*oct_adr, pme->octree().Nodes().data(), 6*n_oct*sizeof(mjtNum));
       memcpy(m->oct_child + 8*oct_adr, pme->octree().Child().data(), 8*n_oct*sizeof(int));
       memcpy(m->oct_depth + oct_adr, pme->octree().Level().data(), n_oct*sizeof(int));
+      if (!pme->octree().Coeff().empty()) {
+        memcpy(m->oct_coeff + 8*oct_adr, pme->octree().Coeff().data(), 8*n_oct*sizeof(mjtNum));
+      } else {
+        mjuu_zerovec(m->oct_coeff + 8*oct_adr, 8*n_oct);
+      }
     }
 
     // advance counters

src/user/user_objects.h

@@ -222,6 +222,7 @@ struct mjCOctree_ {
   std::vector<int> child_;           // children of each node     (nnode x 8)
   std::vector<double> node_;         // bounding boxes            (nnode x 6)
   std::vector<int> level_;           // levels of each node       (nnode x 1)
+  std::vector<double> coeff_;        // interpo coefficients      (nnode x 8)
   std::vector<Triangle> face_;       // mesh faces                (nface x 3)
   double ipos_[3] = {0, 0, 0};
   double iquat_[4] = {1, 0, 0, 0};
@@ -246,6 +247,8 @@ class mjCOctree : public mjCOctree_ {
     level_.clear();
     face_.clear();
   }
+  void AddCoeff(double coeff) { coeff_.push_back(coeff); }
+  const std::vector<double>& Coeff() const { return coeff_; }
 
  private:
   void Make(std::vector<Triangle>& elements);

unity/Runtime/Bindings/MjBindings.cs

@@ -5364,6 +5364,7 @@ public unsafe struct mjModel_ {
   public int* oct_depth;
   public int* oct_child;
   public double* oct_aabb;
+  public double* oct_coeff;
   public int* jnt_type;
   public int* jnt_qposadr;
   public int* jnt_dofadr;