Add octree support to MJX API.

Also, enable the touch sensor to use the octree SDF.

PiperOrigin-RevId: 783313462
Change-Id: I0032c2a7164f4770aa10a3612db3fd70bea954db

src/engine/engine_collision_sdf.c

@@ -144,6 +144,11 @@ mjtNum oct_distance(const mjModel* m, const mjtNum p[3], int meshid) {
   mjtNum* oct_aabb = m->oct_aabb + 6*octadr;
   mjtNum* oct_coeff = m->oct_coeff + 8*octadr;
 
+  if (octadr == -1) {
+    mjERROR("Octree not found in mesh %d", meshid);
+    return 0;
+  }
+
   mjtNum w[8];
   mjtNum sdf = 0;
   mjtNum point[3] = {p[0], p[1], p[2]};
@@ -168,6 +173,10 @@ void oct_gradient(const mjModel* m, mjtNum grad[3], const mjtNum point[3], int m
   mjtNum* oct_aabb = m->oct_aabb + 6*octadr;
   mjtNum* oct_coeff = m->oct_coeff + 8*octadr;
 
+  if (octadr == -1) {
+    mjERROR("Octree not found in mesh %d", meshid);
+  }
+
   // analytic in the interior
   if (boxProjection(p, oct_aabb) <= 0) {
     mjtNum dw[8][3];
@@ -267,6 +276,8 @@ static mjtNum geomDistance(const mjModel* m, const mjData* d, const mjpPlugin* p
     } else {
       return oct_distance(m, x, i);
     }
+  case mjGEOM_MESH:
+    return oct_distance(m, x, i);
   default:
     mjERROR("sdf collisions not available for geom type %d", type);
     return 0;
@@ -371,6 +382,9 @@ static void geomGradient(mjtNum gradient[3], const mjModel* m, const mjData* d,
       oct_gradient(m, gradient, x, i);
     }
     break;
+  case mjGEOM_MESH:
+    oct_gradient(m, gradient, x, i);
+    break;
   default:
     mjERROR("sdf collisions not available for geom type %d", type);
   }

src/engine/engine_vis_visualize.c

@@ -763,7 +763,8 @@ void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* vopt,
   if (vopt->flags[mjVIS_MESHBVH]) {
     for (int geomid = 0; geomid < m->ngeom; geomid++) {
       int meshid = m->geom_dataid[geomid];
-      if (m->geom_type[geomid] == mjGEOM_SDF || meshid == -1) {
+      // skip if not a mesh or if there is an octree
+      if (m->geom_type[geomid] == mjGEOM_SDF || meshid == -1 || m->mesh_octadr[meshid] >= 0) {
         continue;
       }
 
@@ -812,7 +813,7 @@ void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* vopt,
   if (vopt->flags[mjVIS_MESHBVH]) {
     for (int geomid = 0; geomid < m->ngeom; geomid++) {
       int meshid = m->geom_dataid[geomid];
-      if (m->geom_type[geomid] != mjGEOM_SDF || meshid == -1) {
+      if (meshid == -1 || (m->geom_type[geomid] == mjGEOM_MESH && m->mesh_octadr[meshid] == -1)) {
         continue;
       }