[PERF] Split rigid collision BVH into static + dynamic subsets (RPL multi-depth)#2878
Conversation
…ulti-depth) Follow-up to Genesis-Embodied-AI#2867. A scene with one moving robot on a large static terrain still rebuilds a single combined collision BVH over every face each step, because the rebuild-skip keys off "all links in the solver are fixed" -- false as soon as the robot moves. The static terrain (the bulk of the faces) is re-fit every step for nothing. This decomposes the rigid solver's collision faces into two compacted BVHs by owning-link fixedness (RaycasterSensor._partition_collision_faces): - static subset (faces on fixed links: terrain / walls): maybe_static, built once, then skipped + shared across envs (the dominant per-step cost for one robot on a big static terrain). - dynamic subset (faces on movable links: the robot): rebuilt each step, but the rebuild + radix sort now scale with the robot's face count, not the whole scene. The two are cast separately and merged via the existing is_merge path (closest hit wins), so the result is identical to one combined BVH. This is the "multi-depth" decomposition from RPL (arXiv:2602.03002): cast the dynamic robot and static terrain meshes separately and reuse the static acceleration structure across timesteps and environments. Implementation - Each BVH is built over a compacted face subset. A `face_ids` array maps a BVH leaf slot to the global face index; bvh_ray_cast remaps after reading the morton-code primitive id, and update_aabbs iterates the subset. `n_triangles` in bvh_ray_cast now derives from the morton-codes shape (the BVH's own leaf count) instead of the solver-global face count. - The existing maybe_static/needs_rebuild skip and the AABB-derived shared_across_envs test are already per-entry, so they apply to each subset unchanged: the static subset's GEOMETRY subscriber only fires on an explicit set_pos/set_quat (e.g. re-randomized terrain), never on physics-driven robot motion, so it stays skipped + shared while the robot subset rebuilds. - A pure-static or pure-dynamic solver yields a single subset with identity face_ids, i.e. the previous single-BVH behavior -- bit-identical. - kernel_cast_ray (viewer pick) and the viewer plugin thread an identity face_ids over the full mesh. Perf (perceptive depth camera, 64x36 rays, 18.5k-face terrain + G1, RTX 3080): 1024 envs: 1037 -> 706 ms/step (1.47x) measured on the pre-refactor branch; win grows with env count x terrain faces. Re-validated functionally on the current main: a moving-robot-on-terrain scene now builds 1 static (skipped + shared) + 1 dynamic BVH, depth output unchanged. Tests - New tests/test_sensors.py::test_raycaster_static_dynamic_bvh_split asserts the split structure (one static + one dynamic collision BVH, static shared across envs), the merge reporting the closer of static/dynamic as a movable box enters/leaves a ray, and that the static BVH stays skipped across a dynamic move. Passes for n_envs in {0, 2}. - Existing raycaster/lidar suite unchanged (single full-mesh path is the identity-face_ids case; mixed scenes now exercise the two-BVH merge). Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 48736869ad
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| for face_ids, subset_static in self._partition_collision_faces(solver): | ||
| n_sub = int(face_ids.shape[0]) | ||
| aabb = AABB(n_batches=n_envs, n_aabbs=n_sub) | ||
| bvh = LBVH(aabb, max_n_query_result_per_aabb=0, n_radix_sort_groups=min(64, n_sub)) |
There was a problem hiding this comment.
Avoid creating one-face BVHs after the split
In mixed static/dynamic solvers, this now builds an LBVH for each compacted subset independently. If either subset contains exactly one collision face (for example a single-triangle custom movable mesh on a static mesh with other faces, or vice versa), LBVH.__init__ rejects aabb.n_aabbs < 2, so scene.build() raises even though the old combined BVH had enough faces to build. Please keep singleton subsets in the combined path, merge them with another subset, or add a non-BVH fallback for the single-triangle case.
Useful? React with 👍 / 👎.
Description
Follow-up to #2867. That PR added the static-rebuild-skip + shared-BVH-across-envs fast paths, but they only engage when every link in the solver is fixed (
maybe_static = all(link.is_fixed for link in solver.links)). The common RL case — one moving robot on a large static terrain in a single rigid solver — never qualifies: the robot's movable links make the whole solver non-static, so the single combined collision BVH (terrain + robot) is rebuilt over every face each step. The static terrain, which is the bulk of the faces, is re-fit every step for nothing.This PR decomposes the rigid solver's collision faces into two compacted BVHs by owning-link fixedness (
RaycasterSensor._partition_collision_faces):maybe_static, built once, then skipped + shared across envs (the dominant per-step cost for one robot on a big static terrain).The two are cast separately and merged via the existing
is_mergepath (closest hit wins), so the result is identical to one combined BVH. This is the "multi-depth" decomposition from RPL (arXiv:2602.03002): cast the dynamic robot and static terrain meshes separately and reuse the static acceleration structure across timesteps and environments.Implementation
face_idsarray maps a BVH leaf slot → global face index;bvh_ray_castremaps after reading the morton-code primitive id, andupdate_aabbsiterates the subset.n_trianglesinbvh_ray_castnow derives from the morton-codes shape (the BVH's own leaf count) instead of the solver-global face count.maybe_static/needs_rebuildskip and the AABB-derivedshared_across_envstest are already per-entry, so they apply to each subset unchanged: the static subset's GEOMETRY subscriber fires only on an explicitset_pos/set_quat(e.g. re-randomized terrain), never on physics-driven robot motion — so it stays skipped + shared while the robot subset rebuilds.face_ids, i.e. the previous single-BVH behavior — bit-identical.kernel_cast_ray(viewer pick) and the viewer plugin thread an identityface_idsover the full mesh.Motivation and Context
Profiling perceptive whole-body imitation (torso depth camera, 64×36 rays, 18.5k-face terrain + G1-29DoF, RTX 3080): the combined-BVH rebuild dominated
env.stepand #2867's static-skip never engaged because the robot shares the solver. With the split, a moving-robot-on-terrain scene builds 1 static (skipped + shared) + 1 dynamic BVH; depth output is unchanged.Per-step (measured on the pre-refactor branch): 1024 envs 1037 → 706 ms/step (1.47×); the win grows with env count × terrain face count (the eliminated static rebuild scales with both).
How Has This Been Tested?
Environment: single RTX 3080, CUDA backend.
tests/test_sensors.py::test_raycaster_static_dynamic_bvh_split(n_envs ∈ {0, 2}): asserts the split structure (one static + one dynamic collision BVH; staticshared_across_envs), that the merge reports the closer of static/dynamic as a movable box enters/leaves a ray, and that the static BVH stays skipped (needs_rebuildFalse) across a dynamic-only move.face_idscase (bit-identical); mixed scenes now exercise the two-BVH merge.Checklist
test_raycaster_static_dynamic_bvh_split).🤖 Generated with Claude Code