Backtracking line search for SDF collisions.

Increased tolerance above zero in nutbolt.xml since now the two parts lock otherwise. Measured speedup of 2x in this example.

This solves the issue of scale-dependent step sizes. Tested on gears with 10cm diameter and 2cm thickness.

PiperOrigin-RevId: 584009449
Change-Id: Ied2f62dbbbc592470b91cf910f3553802a57c752

doc/changelog.rst

@@ -2,6 +2,16 @@
 Changelog
 =========
 
+Upcoming version (not yet released)
+-----------------------------------
+
+General
+^^^^^^^
+
+- Improved convergence of Signed Distance Function (SDF) collisions by using line search and a new objective function
+  for the optimization. This allows to decrease the number of initial points needed for finding the contacts and is more
+  robust for very small or large geom sizes.
+
 Version 3.0.1 (November 15, 2023)
 ---------------------------------
 

model/plugin/sdf/nutbolt.xml

@@ -8,7 +8,7 @@
     </plugin>
     <plugin plugin="mujoco.sdf.bolt">
       <instance name="bolt">
-        <config key="radius" value="0.26"/>
+        <config key="radius" value="0.255"/>
       </instance>
     </plugin>
   </extension>
@@ -30,7 +30,7 @@
     </mesh>
   </asset>
 
-  <option sdf_iterations="15" sdf_initpoints="60"/>
+  <option sdf_iterations="10" sdf_initpoints="20"/>
 
   <default>
     <geom solref="0.01 1" solimp=".95 .99 .0001" friction="0.01"/>

plugin/sdf/bolt.cc

@@ -166,12 +166,12 @@ void Bolt::RegisterPlugin() {
   plugin.sdf_distance =
       +[](const mjtNum point[3], const mjData* d, int instance) {
         auto* sdf = reinterpret_cast<Bolt*>(d->plugin_data[instance]);
-        sdf->visualizer_.AddPoint(point);
         return sdf->Distance(point);
       };
   plugin.sdf_gradient = +[](mjtNum gradient[3], const mjtNum point[3],
                         const mjData* d, int instance) {
     auto* sdf = reinterpret_cast<Bolt*>(d->plugin_data[instance]);
+    sdf->visualizer_.AddPoint(point);
     sdf->Gradient(gradient, point);
   };
   plugin.sdf_staticdistance =

plugin/sdf/bowl.cc

@@ -164,12 +164,12 @@ void Bowl::RegisterPlugin() {
   plugin.sdf_distance =
       +[](const mjtNum point[3], const mjData* d, int instance) {
         auto* sdf = reinterpret_cast<Bowl*>(d->plugin_data[instance]);
-        sdf->visualizer_.AddPoint(point);
         return sdf->Distance(point);
       };
   plugin.sdf_gradient = +[](mjtNum gradient[3], const mjtNum point[3],
                         const mjData* d, int instance) {
     auto* sdf = reinterpret_cast<Bowl*>(d->plugin_data[instance]);
+    sdf->visualizer_.AddPoint(point);
     sdf->Gradient(gradient, point);
   };
   plugin.sdf_staticdistance =

plugin/sdf/gear.cc

@@ -250,12 +250,12 @@ void Gear::RegisterPlugin() {
   plugin.sdf_distance =
       +[](const mjtNum point[3], const mjData* d, int instance) {
         auto* sdf = reinterpret_cast<Gear*>(d->plugin_data[instance]);
-        sdf->visualizer_.AddPoint(point);
         return sdf->Distance(point);
       };
   plugin.sdf_gradient = +[](mjtNum gradient[3], const mjtNum point[3],
                         const mjData* d, int instance) {
     auto* sdf = reinterpret_cast<Gear*>(d->plugin_data[instance]);
+    sdf->visualizer_.AddPoint(point);
     sdf->Gradient(gradient, point);
   };
   plugin.sdf_staticdistance =

plugin/sdf/nut.cc

@@ -166,12 +166,12 @@ void Nut::RegisterPlugin() {
   plugin.sdf_distance =
       +[](const mjtNum point[3], const mjData* d, int instance) {
         auto* sdf = reinterpret_cast<Nut*>(d->plugin_data[instance]);
-        sdf->visualizer_.AddPoint(point);
         return sdf->Distance(point);
       };
   plugin.sdf_gradient = +[](mjtNum gradient[3], const mjtNum point[3],
                         const mjData* d, int instance) {
     auto* sdf = reinterpret_cast<Nut*>(d->plugin_data[instance]);
+    sdf->visualizer_.AddPoint(point);
     sdf->Gradient(gradient, point);
   };
   plugin.sdf_staticdistance =

plugin/sdf/sdflib.cc

@@ -146,12 +146,12 @@ void SdfLib::RegisterPlugin() {
   plugin.sdf_distance =
       +[](const mjtNum point[3], const mjData* d, int instance) {
         auto* sdf = reinterpret_cast<SdfLib*>(d->plugin_data[instance]);
-        sdf->visualizer_.AddPoint(point);
         return sdf->Distance(point);
       };
   plugin.sdf_gradient = +[](mjtNum gradient[3], const mjtNum point[3],
                         const mjData* d, int instance) {
     auto* sdf = reinterpret_cast<SdfLib*>(d->plugin_data[instance]);
+    sdf->visualizer_.AddPoint(point);
     sdf->Gradient(gradient, point);
   };
 

src/engine/engine_collision_sdf.c

@@ -249,7 +249,6 @@ void mjc_gradient(const mjModel* m, const mjData* d, const mjSDF* s,
       gradient[1] = - grad1[1] * B - grad2[1] * A;
       gradient[2] = - grad1[2] * B - grad2[2] * A;
     }
-    mju_normalize3(gradient);
     break;
   case mjSDFTYPE_SINGLE:
     geomGradient(gradient, m, d, s->plugin[0], s->id[0], point[0], s->geomtype[0]);
@@ -388,10 +387,15 @@ static mjtNum stepFrankWolfe(mjtNum x[3], const mjtNum* corners, int ncorners,
 // finds minimum using gradient descent
 static mjtNum stepGradient(mjtNum x[3], const mjModel* m, const mjSDF* s,
                            mjData* d) {
-  mjtNum alpha = 0.2;    // step along the gradient direction
+  const mjtNum c = .1;       // reduction factor for the target decrease in the objective function
+  const mjtNum rho = .5;     // reduction factor for the gradient scaling (alpha)
+  const mjtNum amin = 1e-4;  // minimum value for alpha
+  mjtNum dist = mjMAXVAL;
 
   for (int step=0; step < m->opt.sdf_iterations; step++) {
     mjtNum grad[3];
+    mjtNum alpha = 2.;  // initial line search factor scaling the gradient
+                        // the units of the gradient depend on s->type
 
     // evaluate gradient
     mjc_gradient(m, d, s, grad, x);
@@ -403,12 +407,29 @@ static mjtNum stepGradient(mjtNum x[3], const mjModel* m, const mjSDF* s,
       return mjMAXVAL;
     }
 
-    // update solution
-    mju_addToScl3(x, grad, -alpha/(mjtNum)(step+1));
+    // save current solution
+    mjtNum x0[] = {x[0], x[1], x[2]};
+
+    // evaluate distance
+    mjtNum dist0 = mjc_distance(m, d, s, x0);
+    mjtNum wolfe = - c * alpha * mju_dot3(grad, grad);
+
+    // backtracking line search
+    do {
+      alpha *= rho;
+      wolfe *= rho;
+      mju_addScl3(x, x0, grad, -alpha);
+      dist = mjc_distance(m, d, s, x);
+    } while (alpha > amin && dist - dist0 > wolfe);
+
+    // if no improvement, early stop
+    if (dist0 < dist) {
+      return dist;
+    }
   }
 
-  // compute distance
-  return mjc_distance(m, d, s, x);
+  // the distance will be used for the contact creation
+  return dist;
 }
 
 //---------------------------- bounding box vs sdf -------------------------------------------------