Add octree to meshes of SDF geoms.

PiperOrigin-RevId: 777634862
Change-Id: I3210c2f8d73fd30ceecb0a11abd1d82e086b3d5d

src/engine/engine_io.c

@@ -460,7 +460,7 @@ static void freeModelBuffers(mjModel* m) {
 // allocate and initialize mjModel structure
 void mj_makeModel(mjModel** dest,
     int nq, int nv, int nu, int na, int nbody, int nbvh,
-    int nbvhstatic, int nbvhdynamic, int njnt, int ngeom, int nsite, int ncam,
+    int nbvhstatic, int nbvhdynamic, int noct, int njnt, int ngeom, int nsite, int ncam,
     int nlight, int nflex, int nflexnode, int nflexvert, int nflexedge, int nflexelem,
     int nflexelemdata, int nflexelemedge, int nflexshelldata, int nflexevpair, int nflextexcoord,
     int nmesh, int nmeshvert, int nmeshnormal, int nmeshtexcoord, int nmeshface,
@@ -499,6 +499,7 @@ void mj_makeModel(mjModel** dest,
   m->nbvh = nbvh;
   m->nbvhstatic = nbvhstatic;
   m->nbvhdynamic = nbvhdynamic;
+  m->noct = noct;
   m->njnt = njnt;
   m->ngeom = ngeom;
   m->nsite = nsite;
@@ -639,7 +640,7 @@ mjModel* mj_copyModel(mjModel* dest, const mjModel* src) {
   if (!dest) {
     mj_makeModel(&dest,
       src->nq, src->nv, src->nu, src->na, src->nbody, src->nbvh,
-      src->nbvhstatic, src->nbvhdynamic, src->njnt, src->ngeom, src->nsite,
+      src->nbvhstatic, src->nbvhdynamic, src->noct, src->njnt, src->ngeom, src->nsite,
       src->ncam, src->nlight, src->nflex, src->nflexnode, src->nflexvert, src->nflexedge,
       src->nflexelem, src->nflexelemdata, src->nflexelemedge, src->nflexshelldata,
       src->nflexevpair, src->nflextexcoord, src->nmesh, src->nmeshvert,
@@ -835,7 +836,7 @@ mjModel* mj_loadModelBuffer(const void* buffer, int buffer_sz) {
                ints[42], ints[43], ints[44], ints[45], ints[46], ints[47], ints[48],
                ints[49], ints[50], ints[51], ints[52], ints[53], ints[54], ints[55],
                ints[56], ints[57], ints[58], ints[59], ints[60], ints[61], ints[62],
-               ints[63], ints[64], ints[65], ints[66], ints[67], ints[68]);
+               ints[63], ints[64], ints[65], ints[66], ints[67], ints[68], ints[69]);
   if (!m || m->nbuffer != sizes[getnsize()-1]) {
     mju_warning("Corrupted model, wrong size parameters");
     mj_deleteModel(m);

src/engine/engine_io.h

@@ -52,7 +52,7 @@ void mj_defaultStatistic(mjStatistic* stat);
 
 // allocate mjModel
 void mj_makeModel(mjModel** dest,
-    int nq, int nv, int nu, int na, int nbody, int nbvh, int nbvhstatic, int nbvhdynamic,
+    int nq, int nv, int nu, int na, int nbody, int nbvh, int nbvhstatic, int nbvhdynamic, int noct,
     int njnt, int ngeom, int nsite, int ncam, int nlight, int nflex, int nflexnode, int nflexvert,
     int nflexedge, int nflexelem, int nflexelemdata, int nflexelemedge, int nflexshelldata,
     int nflexevpair, int nflextexcoord, int nmesh, int nmeshvert, int nmeshnormal,

src/engine/engine_vis_visualize.c

@@ -763,7 +763,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 (meshid == -1) {
+      if (m->geom_type[geomid] == mjGEOM_SDF || meshid == -1) {
         continue;
       }
 
@@ -806,6 +806,43 @@ void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* vopt,
     }
   }
 
+  // mesh octree
+  category = mjCAT_DECOR;
+  objtype = mjOBJ_UNKNOWN;
+  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) {
+        continue;
+      }
+
+      for (int b = 0; b < m->mesh_octnum[meshid]; b++) {
+        int i = b + m->mesh_octadr[meshid];
+
+        START
+        if (m->oct_depth[i] != vopt->bvh_depth) {
+          continue;
+        }
+
+        // box color
+        const float* rgba = m->vis.rgba.bv;
+
+        // get xpos, xmat, size
+        const mjtNum* xpos = d->geom_xpos + 3 * geomid;
+        const mjtNum* xmat = d->geom_xmat + 9 * geomid;
+        const mjtNum *size = m->oct_aabb + 6*i + 3;
+
+        // offset xpos with aabb center (not always at geom origin)
+        const mjtNum *center = m->oct_aabb + 6*i;
+        mjtNum pos[3];
+        mju_mulMatVec3(pos, xmat, center);
+        mju_addTo3(pos, xpos);
+        mjv_initGeom(thisgeom, mjGEOM_LINEBOX, size, pos, xmat, rgba);
+        FINISH
+      }
+    }
+  }
+
   // inertia
   objtype = mjOBJ_BODY;
   if (vopt->flags[mjVIS_INERTIA]) {