From 4ea778f2d3d1815a8d8acaec2aa92946c4de0cb5 Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Tue, 15 Jul 2025 06:43:58 -0700
Subject: [PATCH] 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      |  5 +++++
 plugin/sensor/touch_stress.cc     | 10 +++++-----
 src/engine/engine_collision_sdf.c | 14 ++++++++++++++
 src/engine/engine_vis_visualize.c |  5 +++--
 4 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py
index 57d68208..6bed72c2 100644
--- a/mjx/mujoco/mjx/_src/types.py
+++ b/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
diff --git a/plugin/sensor/touch_stress.cc b/plugin/sensor/touch_stress.cc
index 9eaabd1e..50c264b1 100644
--- a/plugin/sensor/touch_stress.cc
+++ b/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);
diff --git a/src/engine/engine_collision_sdf.c b/src/engine/engine_collision_sdf.c
index 53ec258c..71040e60 100644
--- a/src/engine/engine_collision_sdf.c
+++ b/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);
   }
diff --git a/src/engine/engine_vis_visualize.c b/src/engine/engine_vis_visualize.c
index 3973ecb0..e6c3e5f7 100644
--- a/src/engine/engine_vis_visualize.c
+++ b/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;
       }