Add octree support to MJX API.

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

PiperOrigin-RevId: 783313462
Change-Id: I0032c2a7164f4770aa10a3612db3fd70bea954db

mjx/mujoco/mjx/_src/types.py

@@ -527,6 +527,9 @@ class ModelC(PyTreeNode):
   bvh_child: jax.Array
   bvh_nodeid: jax.Array
   bvh_aabb: jax.Array
+  oct_child: jax.Array
+  oct_aabb: jax.Array
+  oct_coeff: jax.Array
   geom_plugin: jax.Array
   light_bodyid: jax.Array
   light_targetbodyid: jax.Array
@@ -742,6 +745,8 @@ class Model(PyTreeNode):
   mesh_faceadr: np.ndarray
   mesh_bvhadr: np.ndarray
   mesh_bvhnum: np.ndarray
+  mesh_octadr: np.ndarray
+  mesh_octnum: np.ndarray
   mesh_graphadr: np.ndarray
   mesh_vert: np.ndarray
   mesh_face: np.ndarray

plugin/sensor/touch_stress.cc

@@ -309,8 +309,9 @@ void TouchStress::Compute(const mjModel* m, mjData* d, int instance) {
       geomtype[0] = (mjtGeom)m->geom_type[geom];
     }
 
-    // Skip mesh geoms.
-    if (geomtype[0] == mjGEOM_MESH) {
+    // Skip mesh geoms not having an octree.
+    if (geomtype[0] == mjGEOM_MESH &&
+        m->mesh_octadr[m->geom_dataid[geom]] == -1) {
       continue;
     }
 
@@ -345,9 +346,8 @@ void TouchStress::Compute(const mjModel* m, mjData* d, int instance) {
         mju_sub3(tmp, xpos, d->geom_xpos + 3*geom);
         mju_mulMatTVec3(lpos, d->geom_xmat + 9*geom, tmp);
 
-        // Add mesh position if needed.
-        if (m->geom_type[geom] == mjGEOM_MESH ||
-            m->geom_type[geom] == mjGEOM_SDF) {
+        // SDF plugins are in the original mesh frame.
+        if (sdf_ptr[0] != NULL) {
           mjtNum mesh_mat[9];
           mju_quat2Mat(mesh_mat, m->mesh_quat + 4 * m->geom_dataid[geom]);
           mju_mulMatVec3(lpos, mesh_mat, lpos);

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;
       }