From 96118da08bbd4ede267817be111a932feec71be9 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Wed, 9 Jul 2025 08:56:20 -0700 Subject: [PATCH] Move SDF precomputation to the mujoco compiler. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 1 + include/mujoco/mjmodel.h | 1 + include/mujoco/mjxmacro.h | 1 + plugin/sdf/sdflib.cc | 34 ++++------------------------ plugin/sdf/sdflib.h | 3 +-- python/mujoco/introspect/structs.py | 8 +++++++ src/user/user_mesh.cc | 20 ++++++++++++++++ src/user/user_model.cc | 5 ++++ src/user/user_objects.h | 3 +++ unity/Runtime/Bindings/MjBindings.cs | 1 + 10 files changed, 45 insertions(+), 32 deletions(-) diff --git a/doc/includes/references.h b/doc/includes/references.h index 894cf25b..0f27ccf0 100644 --- a/doc/includes/references.h +++ b/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) diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index c92edd75..5a62d521 100644 --- a/include/mujoco/mjmodel.h +++ b/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) diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 68032375..4c7272e3 100644 --- a/include/mujoco/mjxmacro.h +++ b/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 ) \ diff --git a/plugin/sdf/sdflib.cc b/plugin/sdf/sdflib.cc index 84da6553..9fae83bd 100644 --- a/plugin/sdf/sdflib.cc +++ b/plugin/sdf/sdflib.cc @@ -18,7 +18,6 @@ #include #include -#include #include #include #include "sdf.h" @@ -140,44 +139,19 @@ std::optional 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 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); } diff --git a/plugin/sdf/sdflib.h b/plugin/sdf/sdflib.h index 17246448..65def3aa 100644 --- a/plugin/sdf/sdflib.h +++ b/plugin/sdf/sdflib.h @@ -23,7 +23,6 @@ #include #include #include "sdf.h" -#include 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 sdf_coeff_; diff --git a/python/mujoco/introspect/structs.py b/python/mujoco/introspect/structs.py index a969052b..82137c3b 100644 --- a/python/mujoco/introspect/structs.py +++ b/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( diff --git a/src/user/user_mesh.cc b/src/user/user_mesh.cc index 53155309..facb78d2 100644 --- a/src/user/user_mesh.cc +++ b/src/user/user_mesh.cc @@ -32,6 +32,7 @@ #include #include "user/user_api.h" +#include #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 diff --git a/src/user/user_model.cc b/src/user/user_model.cc index 5a9ae0ba..257b6524 100644 --- a/src/user/user_model.cc +++ b/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 diff --git a/src/user/user_objects.h b/src/user/user_objects.h index cc111ae7..98bc5d2d 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -222,6 +222,7 @@ struct mjCOctree_ { std::vector child_; // children of each node (nnode x 8) std::vector node_; // bounding boxes (nnode x 6) std::vector level_; // levels of each node (nnode x 1) + std::vector coeff_; // interpo coefficients (nnode x 8) std::vector 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& Coeff() const { return coeff_; } private: void Make(std::vector& elements); diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 2c540d15..f95fc924 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/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;