Add mj_raySdfNormal (not exposed in public header)

PiperOrigin-RevId: 842178540
Change-Id: I5d2db555e9b6c0fa0daaa8bc14c7781663e4c94e

src/engine/engine_ray.c

@@ -846,10 +846,12 @@ mjtNum mju_rayTree(const mjModel* m, const mjData* d, int id, const mjtNum pnt[3
   return x;
 }
 
+// intersect ray with signed distance field, compute normal if given
+mjtNum mj_raySdfNormal(const mjModel* m, const mjData* d, int g,
+                       const mjtNum pnt[3], const mjtNum vec[3],
+                       mjtNum normal[3]) {
+  if (normal) mju_zero3(normal);
 
-// intersect ray with signed distance field
-mjtNum ray_sdf(const mjModel* m, const mjData* d, int g,
-               const mjtNum pnt[3], const mjtNum vec[3]) {
   mjtNum distance_total = 0;
   mjtNum p[3];
   mjtNum kMinDist = 1e-7;
@@ -887,9 +889,15 @@ mjtNum ray_sdf(const mjModel* m, const mjData* d, int g,
   // ray marching, see e.g. https://en.wikipedia.org/wiki/Ray_marching
   for (int i=0; i < 40; i++) {
     mju_addScl3(p, lpnt, lvec, distance_total);
-    mjtNum distance = mjc_distance(m, d, &sdf, p);
+    mjtNum distance = mju_abs(mjc_distance(m, d, &sdf, p));
     distance_total += distance;
     if (mju_abs(distance) < kMinDist) {
+      if (normal) {
+        mju_addScl3(p, lpnt, lvec, distance_total);
+        mjc_gradient(m, d, &sdf, normal, p);
+        mju_normalize3(normal);
+        mju_mulMatVec3(normal, d->geom_xmat + 9*g, normal);
+      }
       return distance_total;
     }
     if (distance > 1e6) {
@@ -906,6 +914,12 @@ mjtNum ray_sdf(const mjModel* m, const mjData* d, int g,
   return -1;
 }
 
+// intersect ray with signed distance field
+static mjtNum ray_sdf(const mjModel* m, const mjData* d, int g,
+                      const mjtNum pnt[3], const mjtNum vec[3]) {
+  return mj_raySdfNormal(m, d, g, pnt, vec, NULL);
+}
+
 // intersect ray with mesh, compute normal if given
 mjtNum mj_rayMeshNormal(const mjModel* m, const mjData* d, int id, const mjtNum pnt[3],
                         const mjtNum vec[3], mjtNum normal[3]) {

src/engine/engine_ray.h

@@ -58,6 +58,10 @@ MJAPI mjtNum mj_rayMesh(const mjModel* m, const mjData* d, int geomid,
 MJAPI mjtNum mj_rayMeshNormal(const mjModel* m, const mjData* d, int geomid,
                               const mjtNum pnt[3], const mjtNum vec[3], mjtNum normal[3]);
 
+// intersect ray with signed distance field, compute normal if given
+MJAPI mjtNum mj_raySdfNormal(const mjModel* m, const mjData* d, int geomid,
+                             const mjtNum pnt[3], const mjtNum vec[3], mjtNum normal[3]);
+
 // intersect ray with primitive geom, no meshes or hfields
 MJAPI mjtNum mju_rayGeom(const mjtNum pos[3], const mjtNum mat[9], const mjtNum size[3],
                          const mjtNum pnt[3], const mjtNum vec[3], int geomtype);

test/engine/CMakeLists.txt

@@ -47,7 +47,11 @@ mujoco_test(engine_passive_test)
 
 mujoco_test(engine_print_test)
 
-mujoco_test(engine_ray_test)
+mujoco_test(
+  engine_ray_test
+  PROPERTIES
+  ENVIRONMENT
+  "MUJOCO_PLUGIN_DIR=$<TARGET_FILE_DIR:sdf_plugin>")
 
 mujoco_test(engine_sensor_test)
 

test/engine/engine_ray_test.cc

@@ -82,7 +82,7 @@ using ::testing::DoubleNear;
 using ::testing::ElementsAre;
 using ::testing::NotNull;
 using ::testing::Pointwise;
-using RayTest = MujocoTest;
+using RayTest = PluginTest;
 
 TEST_F(RayTest, NoExclusions) {
   char error[1024];
@@ -475,11 +475,12 @@ static const char* const kEllipsoidModel = "engine/testdata/ray/ellipsoid.xml";
 static const char* const kCylinderModel = "engine/testdata/ray/cylinder.xml";
 static const char* const kBoxModel = "engine/testdata/ray/box.xml";
 static const char* const kMeshModel = "engine/testdata/ray/mesh.xml";
+static const char* const kSdfModel = "engine/testdata/ray/sdf.xml";
 
 TEST_F(RayTest, GeomNormal) {
   for (const char* path :
        {kPlaneModel, kSphereModel, kCapsuleModel, kEllipsoidModel,
-        kCylinderModel, kBoxModel, kMeshModel}) {
+        kCylinderModel, kBoxModel, kMeshModel, kSdfModel}) {
     const std::string xml_path = GetTestDataFilePath(path);
     char error[1024];
     mjModel* m = mj_loadXML(xml_path.c_str(), 0, error, sizeof(error));
@@ -524,6 +525,9 @@ TEST_F(RayTest, GeomNormal) {
         case mjGEOM_MESH:
           r = mj_rayMeshNormal(m, d, 0, pnt, vec, normal);
           break;
+        case mjGEOM_SDF:
+          r = mj_raySdfNormal(m, d, 0, pnt, vec, normal);
+          break;
         default:
           r = mju_rayGeomNormal(pos, mat, size, pnt, vec, type, normal);
       }
@@ -533,7 +537,8 @@ TEST_F(RayTest, GeomNormal) {
 
       // if no intersection, skip
       if (r < 0) {
-        EXPECT_THAT(normal, ElementsAre(0, 0, 0));
+        EXPECT_THAT(normal, ElementsAre(0, 0, 0))
+            << path << ", time " << d->time;
         continue;
       }
 
@@ -554,6 +559,9 @@ TEST_F(RayTest, GeomNormal) {
           case mjGEOM_MESH:
             dr = mj_rayMeshNormal(m, d, 0, dpnt, vec, nullptr);
             break;
+          case mjGEOM_SDF:
+            dr = mj_raySdfNormal(m, d, 0, dpnt, vec, nullptr);
+            break;
           default:
             dr = mju_rayGeomNormal(pos, mat, size, dpnt, vec, type, nullptr);
         }

test/engine/testdata/ray/sdf.xml

@@ -0,0 +1,28 @@
+<mujoco>
+  <extension>
+    <plugin plugin="mujoco.sdf.torus">
+      <instance name="torus">
+        <config key="radius1" value="1"/>
+        <config key="radius2" value="0.3"/>
+      </instance>
+    </plugin>
+  </extension>
+
+  <asset>
+    <mesh name="torus">
+      <plugin instance="torus"/>
+    </mesh>
+  </asset>
+
+  <option sdf_iterations="10" sdf_initpoints="40"/>
+
+  <include file="ray_swing.xml"/>
+
+  <worldbody>
+    <body pos="-.5 .002 -.2" euler="90 0 -40" mocap="true">
+      <geom type="sdf" mesh="torus">
+        <plugin instance="torus"/>
+      </geom>
+    </body>
+  </worldbody>
+</mujoco>