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
This commit is contained in:
Alessio Quaglino
2025-07-09 06:20:36 -07:00
committed by Copybara-Service
parent e6c5715903
commit b81f1db8af
8 changed files with 194 additions and 121 deletions
+1 -1
View File
@@ -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}
+5 -5
View File
@@ -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
+138 -46
View File
@@ -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;
}
+8 -4
View File
@@ -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