Ensure that the Octree is balanced.

This means that two adjacent nodes can only have at most one level of refinement difference.

Note: This is a small breaking change for mjWarp since the order of the children was flipped.
PiperOrigin-RevId: 806299445
Change-Id: Ia326a31c7a8162ae02d70c868d87779198e805be

src/engine/engine_collision_sdf.c

@@ -124,9 +124,9 @@ static int findOct(mjtNum w[8], mjtNum dw[8][3], const mjtNum* oct_aabb,
     }
 
     // compute which of 8 children to visit next
-    int x = coord[0] < .5 ? 1 : 0;
-    int y = coord[1] < .5 ? 1 : 0;
-    int z = coord[2] < .5 ? 1 : 0;
+    int x = coord[0] < .5 ? 0 : 1;
+    int y = coord[1] < .5 ? 0 : 1;
+    int z = coord[2] < .5 ? 0 : 1;
     stack = oct_child[8 * node + 4*z + 2*y + x];
   }
 

src/user/user_objects.cc

@@ -692,16 +692,18 @@ static bool boxTriangle(const Triangle& v, const double aamm[6]) {
 void mjCOctree::TaskToNode(const OctreeTask& task, OctNode& node,
                            std::unordered_map<Point, int>& vert_map) {
   node.level = task.lev;
+  node.parent_index = task.parent_index;
+  node.child_slot = task.child_slot;
 
   if (task.parent_index != -1) {
     node_[task.parent_index].child[task.child_slot] = task.node_index;
     const auto parent_aamm = node_[task.parent_index].aamm;
-    node.aamm[0] = task.child_slot & 1 ? parent_aamm[0] : (parent_aamm[3] + parent_aamm[0]) / 2;
-    node.aamm[1] = task.child_slot & 2 ? parent_aamm[1] : (parent_aamm[4] + parent_aamm[1]) / 2;
-    node.aamm[2] = task.child_slot & 4 ? parent_aamm[2] : (parent_aamm[5] + parent_aamm[2]) / 2;
-    node.aamm[3] = task.child_slot & 1 ? (parent_aamm[0] + parent_aamm[3]) / 2 : parent_aamm[3];
-    node.aamm[4] = task.child_slot & 2 ? (parent_aamm[1] + parent_aamm[4]) / 2 : parent_aamm[4];
-    node.aamm[5] = task.child_slot & 4 ? (parent_aamm[2] + parent_aamm[5]) / 2 : parent_aamm[5];
+    node.aamm[0] = task.child_slot & 1 ? (parent_aamm[3] + parent_aamm[0]) / 2 : parent_aamm[0];
+    node.aamm[1] = task.child_slot & 2 ? (parent_aamm[4] + parent_aamm[1]) / 2 : parent_aamm[1];
+    node.aamm[2] = task.child_slot & 4 ? (parent_aamm[5] + parent_aamm[2]) / 2 : parent_aamm[2];
+    node.aamm[3] = task.child_slot & 1 ? parent_aamm[3] : (parent_aamm[0] + parent_aamm[3]) / 2;
+    node.aamm[4] = task.child_slot & 2 ? parent_aamm[4] : (parent_aamm[1] + parent_aamm[4]) / 2;
+    node.aamm[5] = task.child_slot & 4 ? parent_aamm[5] : (parent_aamm[2] + parent_aamm[5]) / 2;
   }
 
   for (int i = 0; i < 8; i++) {
@@ -721,8 +723,8 @@ void mjCOctree::TaskToNode(const OctreeTask& task, OctNode& node,
 }
 
 
-void mjCOctree::Subdivide(std::deque<OctreeTask>& queue, const std::vector<Triangle*>& colliding,
-                          const OctreeTask& task, std::unordered_map<Point, int>& vert_map) {
+void mjCOctree::Subdivide(const OctreeTask& task, std::unordered_map<Point, int>& vert_map,
+                          std::deque<OctreeTask>* queue, const std::vector<Triangle*>& colliding) {
   for (int i = 0; i < 8; i++) {
     OctreeTask new_task;
     new_task.elements = colliding;
@@ -733,7 +735,138 @@ void mjCOctree::Subdivide(std::deque<OctreeTask>& queue, const std::vector<Trian
 
     node_.push_back(OctNode());
     TaskToNode(new_task, node_.back(), vert_map);
-    queue.push_back(std::move(new_task));
+    if (queue) {
+      queue->push_back(std::move(new_task));
+    }
+  }
+}
+
+
+// recursively finds the adjacent ancestor neighbor region
+int mjCOctree::FindCoarseNeighbor(int node_idx, int dir) {
+  if (node_idx == -1) {
+    return -1;
+  }
+
+  int parent_idx = node_[node_idx].parent_index;
+
+  // if we are at the root, we have no parent and thus no siblings or external neighbors
+  if (parent_idx == -1) {
+    return -1;
+  }
+
+  int child_slot = node_[node_idx].child_slot;
+  int dim = dir / 2;
+  int side = dir % 2;
+  int bit = 1 << dim;
+
+  if (side != ((child_slot & bit) != 0)) {
+    // internal neighbor case: This is the successful termination of the climb
+    // return the adjacent sibling node
+    return node_[parent_idx].child[child_slot ^ bit];
+  } else {
+    // external neighbor case: Recurse up the tree
+    // ask our parent to find its neighbor in the same direction
+    return FindCoarseNeighbor(parent_idx, dir);
+  }
+}
+
+
+int mjCOctree::FindNeighbor(int node_idx, int dir) {
+  if (node_idx == -1) {
+    return -1;
+  }
+
+  // call the helper to find the adjacent to the coarse neighbor.
+  // this might be an internal node (e.g., our parent's sibling)
+  int result = FindCoarseNeighbor(node_idx, dir);
+
+  if (result == -1) {
+    // no neighbor found (either at tree boundary or some other error)
+    return -1;
+  }
+
+  // leaf descent
+  double node_center[3] = {
+      (node_[node_idx].aamm[0] + node_[node_idx].aamm[3]) / 2,
+      (node_[node_idx].aamm[1] + node_[node_idx].aamm[4]) / 2,
+      (node_[node_idx].aamm[2] + node_[node_idx].aamm[5]) / 2,
+  };
+
+  while (node_[result].child[0] != -1) {
+    double result_center[3] = {
+        (node_[result].aamm[0] + node_[result].aamm[3]) / 2,
+        (node_[result].aamm[1] + node_[result].aamm[4]) / 2,
+        (node_[result].aamm[2] + node_[result].aamm[5]) / 2,
+    };
+
+    // find relative octant of our node w.r.t. the neighbor's center
+    int next_child_slot = 0;
+    if (node_center[0] > result_center[0]) next_child_slot |= 1;
+    if (node_center[1] > result_center[1]) next_child_slot |= 2;
+    if (node_center[2] > result_center[2]) next_child_slot |= 4;
+
+    int dim = dir / 2;
+    int side = dir % 2;
+    int bit = 1 << dim;
+
+    // we need the child on the opposite side (adjacent to this node)
+    int op_side = (side != 1);
+    next_child_slot = (next_child_slot & ~bit) | (op_side * bit);
+    result = node_[result].child[next_child_slot];
+  }
+
+  return result;
+}
+
+
+// refine the octree by subdividing nodes that are too coarse such that the
+// maximum level difference between adjacent nodes is at most 1.
+void mjCOctree::BalanceOctree(std::unordered_map<Point, int>& vert_map) {
+  bool changed = true;
+  while (changed) {
+    changed = false;
+    std::vector<int> leaves;
+    for (int i = 0; i < nnode_; ++i) {
+      if (node_[i].child[0] == -1) {
+        leaves.push_back(i);
+      }
+    }
+
+    // find the nodes that are too coarse, only leaves need to be checked
+    std::vector<int> leaves_to_subdivide;
+    for (int leaf_idx : leaves) {
+      if (node_[leaf_idx].child[0] != -1) {
+        continue;
+      }
+
+      for (int dir = 0; dir < 6; ++dir) {
+        int neighbor_idx = FindNeighbor(leaf_idx, dir);
+        if (neighbor_idx == -1) {
+          continue;
+        }
+        int neighbor_level = node_[neighbor_idx].level;
+        if (neighbor_level > node_[leaf_idx].level + 1) {
+          leaves_to_subdivide.push_back(leaf_idx);
+        }
+        if (node_[leaf_idx].level > neighbor_level + 1) {
+          leaves_to_subdivide.push_back(neighbor_idx);
+        }
+      }
+    }
+
+    // subdivide the nodes that are too coarse
+    if (!leaves_to_subdivide.empty()) {
+      changed = true;
+      for (int node_idx : leaves_to_subdivide) {
+        if (node_[node_idx].child[0] == -1) {  // check if not already subdivided
+          OctreeTask task;
+          task.node_index = node_idx;
+          task.lev = node_[node_idx].level;
+          Subdivide(task, vert_map);
+        }
+      }
+    }
   }
 }
 
@@ -773,8 +906,11 @@ void mjCOctree::MakeOctree(const std::vector<Triangle*>& elements, const double
     }
 
     // subdivide the node
-    Subdivide(queue, colliding, task, vert_map);
+    Subdivide(task, vert_map, &queue, colliding);
   }
+
+  // store the neighbors of each node
+  BalanceOctree(vert_map);
 }
 
 //------------------------- class mjCDef implementation --------------------------------------------

src/user/user_objects.h

@@ -250,6 +250,8 @@ typedef std::array<std::array<double, 3>, 3> Triangle;
 
 struct OctNode {
   int level = 0;                       // level of the node
+  int parent_index = -1;               // index of the parent node
+  int child_slot = -1;                 // slot of the child node in the parent node
   std::array<int, 8> child = {-1};     // children nodes
   std::array<int, 8> vertid = {-1};    // vertex id's
   std::array<double, 6> aamm = {0};    // bounding box
@@ -302,8 +304,12 @@ class mjCOctree : public mjCOctree_ {
   void MakeOctree(const std::vector<Triangle*>& elements, const double aamm[6],
                   std::unordered_map<Point, int>& vert_map);
   void TaskToNode(const OctreeTask& task, OctNode& node, std::unordered_map<Point, int>& vert_map);
-  void Subdivide(std::deque<OctreeTask>& queue, const std::vector<Triangle*>& colliding,
-                 const OctreeTask& task, std::unordered_map<Point, int>& vert_map);
+  void Subdivide(const OctreeTask& task, std::unordered_map<Point, int>& vert_map,
+                 std::deque<OctreeTask>* queue = nullptr,
+                 const std::vector<Triangle*>& colliding = {});
+  int FindNeighbor(int node_idx, int dir);
+  int FindCoarseNeighbor(int node_idx, int dir);
+  void BalanceOctree(std::unordered_map<Point, int>& vert_map);
 };
 
 

test/user/user_mesh_test.cc

@@ -1279,6 +1279,93 @@ TEST_F(MjCMeshTest, Octree) {
   mj_deleteModel(model);
 }
 
+namespace {
+bool AreAabbsAdjacent(const mjtNum* aabb1, const mjtNum* aabb2) {
+  const double kEps = 1e-6;
+  int touching_dims = 0;
+  int overlapping_dims = 0;
+
+  for (int dim = 0; dim < 3; ++dim) {
+    const mjtNum center1 = aabb1[dim];
+    const mjtNum half_size1 = aabb1[dim + 3];
+    const mjtNum center2 = aabb2[dim];
+    const mjtNum half_size2 = aabb2[dim + 3];
+    const mjtNum gap =
+        std::abs(center1 - center2) - (half_size1 + half_size2);
+    if (std::abs(gap) < kEps) {
+      touching_dims++;
+    } else if (gap < -kEps) {
+      overlapping_dims++;
+    }
+  }
+
+  return touching_dims == 1 && overlapping_dims == 2;
+}
+}  // namespace
+
+TEST_F(MjCMeshTest, OctreeIsBalanced) {
+  const std::string xml_path = GetTestDataFilePath(kTorusPath);
+  std::array<char, 1024> error;
+  mjSpec* spec = mj_parseXML(xml_path.c_str(), 0, error.data(), error.size());
+  mjsGeom* geom = mjs_asGeom(mjs_firstElement(spec, mjOBJ_GEOM));
+  geom->type = mjGEOM_SDF;
+  mjModel* model = mj_compile(spec, 0);
+  ASSERT_THAT(model, NotNull()) << error.data();
+  EXPECT_GT(model->mesh_octnum[0], 0);
+
+  const int octree_adr = model->mesh_octadr[0];
+  const int noct = model->mesh_octnum[0];
+
+  std::vector<int> leaves;
+  for (int i = 0; i < noct; ++i) {
+    bool is_leaf = true;
+    for (int j = 0; j < 8; ++j) {
+      if (model->oct_child[(octree_adr + i) * 8 + j] != -1) {
+        is_leaf = false;
+        break;
+      }
+    }
+    if (is_leaf) {
+      leaves.push_back(i);
+    }
+  }
+
+  int unbalanced_pairs = 0;
+  for (int i = 0; i < leaves.size(); ++i) {
+    for (int j = i + 1; j < leaves.size(); ++j) {
+      const int node1_idx = leaves[i];
+      const int node2_idx = leaves[j];
+      const mjtNum* aabb1 = &model->oct_aabb[(octree_adr + node1_idx) * 6];
+      const mjtNum* aabb2 = &model->oct_aabb[(octree_adr + node2_idx) * 6];
+
+      if (AreAabbsAdjacent(aabb1, aabb2)) {
+        const int level1 = model->oct_depth[octree_adr + node1_idx];
+        const int level2 = model->oct_depth[octree_adr + node2_idx];
+        if (std::abs(level1 - level2) > 1) {
+          if (unbalanced_pairs < 10) {
+            ADD_FAILURE()
+                << "Nodes " << node1_idx << " (level " << level1 << ") and "
+                << node2_idx << " (level " << level2 << ") are not balanced."
+                << "\nAABB1: center=(" << aabb1[0] << ", " << aabb1[1] << ", "
+                << aabb1[2] << "), half_size=(" << aabb1[3] << ", "
+                << aabb1[4] << ", " << aabb1[5] << ")"
+                << "\nAABB2: center=(" << aabb2[0] << ", " << aabb2[1] << ", "
+                << aabb2[2] << "), half_size=(" << aabb2[3] << ", "
+                << aabb2[4] << ", " << aabb2[5] << ")";
+          }
+          unbalanced_pairs++;
+        }
+      }
+    }
+  }
+
+  EXPECT_EQ(unbalanced_pairs, 0)
+      << "Found " << unbalanced_pairs << " unbalanced adjacent leaf pairs.";
+
+  mj_deleteSpec(spec);
+  mj_deleteModel(model);
+}
+
 TEST_F(MjCMeshTest, OctreeNotComputedForNonSDF) {
   const std::string xml_path = GetTestDataFilePath(kTorusPath);
   std::array<char, 1024> error;
@@ -1313,7 +1400,7 @@ TEST_F(MjCMeshTest, OctreeCube) {
   std::array<char, 1024> error;
   mjModel* m = LoadModelFromString(xml, error.data(), error.size());
   ASSERT_THAT(m, NotNull()) << error.data();
-  EXPECT_EQ(m->noct, 54089);
+  EXPECT_EQ(m->noct, 63497);
   mjData* d = mj_makeData(m);
   ASSERT_THAT(d, NotNull());
   mj_forward(m, d);