From 4fc69fa64d0a623ccdc4f811a7b9bdb626ae5159 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Fri, 12 Sep 2025 08:33:25 -0700 Subject: [PATCH] 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 | 6 +- src/user/user_objects.cc | 156 ++++++++++++++++++++++++++++-- src/user/user_objects.h | 10 +- test/user/user_mesh_test.cc | 89 ++++++++++++++++- 4 files changed, 245 insertions(+), 16 deletions(-) diff --git a/src/engine/engine_collision_sdf.c b/src/engine/engine_collision_sdf.c index ee30adbc..9cc1a8f2 100644 --- a/src/engine/engine_collision_sdf.c +++ b/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]; } diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc index b3781975..9e843333 100644 --- a/src/user/user_objects.cc +++ b/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& 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& queue, const std::vector& colliding, - const OctreeTask& task, std::unordered_map& vert_map) { +void mjCOctree::Subdivide(const OctreeTask& task, std::unordered_map& vert_map, + std::deque* queue, const std::vector& colliding) { for (int i = 0; i < 8; i++) { OctreeTask new_task; new_task.elements = colliding; @@ -733,7 +735,138 @@ void mjCOctree::Subdivide(std::deque& queue, const std::vectorpush_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& vert_map) { + bool changed = true; + while (changed) { + changed = false; + std::vector 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 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& 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 -------------------------------------------- diff --git a/src/user/user_objects.h b/src/user/user_objects.h index 9d00d11e..9ab8f589 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -250,6 +250,8 @@ typedef std::array, 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 child = {-1}; // children nodes std::array vertid = {-1}; // vertex id's std::array aamm = {0}; // bounding box @@ -302,8 +304,12 @@ class mjCOctree : public mjCOctree_ { void MakeOctree(const std::vector& elements, const double aamm[6], std::unordered_map& vert_map); void TaskToNode(const OctreeTask& task, OctNode& node, std::unordered_map& vert_map); - void Subdivide(std::deque& queue, const std::vector& colliding, - const OctreeTask& task, std::unordered_map& vert_map); + void Subdivide(const OctreeTask& task, std::unordered_map& vert_map, + std::deque* queue = nullptr, + const std::vector& colliding = {}); + int FindNeighbor(int node_idx, int dir); + int FindCoarseNeighbor(int node_idx, int dir); + void BalanceOctree(std::unordered_map& vert_map); }; diff --git a/test/user/user_mesh_test.cc b/test/user/user_mesh_test.cc index 3713bc56..679975fa 100644 --- a/test/user/user_mesh_test.cc +++ b/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 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 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 error; @@ -1313,7 +1400,7 @@ TEST_F(MjCMeshTest, OctreeCube) { std::array 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);