[MISC] Reduce memory footprint of ray casting-based sensors via cross-env memory sharing.#2914
Merged
duburcqa merged 13 commits intoJul 21, 2026
Conversation
Kashu7100
force-pushed
the
kashu/raycast-shared-static-bvh
branch
from
June 8, 2026 10:23
27c3160 to
89a1d7e
Compare
Kashu7100
changed the title
Kashu7100
force-pushed
the
kashu/raycast-shared-static-bvh
branch
2 times, most recently
from
June 10, 2026 09:18
9304fb8 to
a530b77
Compare
Kashu7100
marked this pull request as ready for review
June 11, 2026 11:22
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a530b7776c
ℹ️ 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".
|
|
The raycast BVH (Raycaster / DepthCamera) is allocated per env (n_batches=n_envs): nodes, AABBs, morton codes and radix-sort scratch are all replicated across envs. For a high-poly static terrain this is the dominant GPU-memory cost — it OOMs at a few thousand envs even though the trees are identical (or fall into a handful of distinct variants) and the cast already reads batch 0 when it detects the trees match (Genesis-Embodied-AI#2867). Add RigidOptions.shared_static_raycast_bvh (default False). When True, the static (fully-fixed) collision BVH is allocated per *distinct geometry* rather than per env: envs that share geometry share one tree. - homogeneous scene -> 1 shared tree - N variants across n_envs -> N trees (N << n_envs, e.g. a terrain curriculum: distinct geometries are read from the per-env active-geom ranges Genesis already stores for heterogeneous envs) Mechanism: two small routing arrays unify all three layouts — env_bvh_idx [n_env] (which tree each env casts against) and batch_repr_env[n_batches] (which env's geometry builds each tree). The cast kernels index env_bvh_idx; update_aabbs builds each tree slot from its representative env. The per-env path is env_bvh_idx=arange (unchanged); the runtime shared-across-envs detection (Genesis-Embodied-AI#2867) becomes env_bvh_idx=0. Opt-in rather than auto-detected because env-identity is a runtime property (a per-env set_pos on a fixed body diverges geometry after build); the flag is a caller guarantee the static geometry stays grouped as built. Default False keeps per-env trees and the runtime detection unchanged. Benchmark — raycast DepthCamera over static eden_dwbp terrain (RTX 3080), total GPU memory: single terrain, 64x36: 256 env 775 -> 154 MB (5.0x) 1024 env 2884 -> 366 MB (7.9x) 4096 env OOM -> runs (was CUDA_ERROR_OUT_OF_MEMORY) 4 terrains across 1024 env: 7300 -> 530 MB (13.8x) Depth is bit-identical to the per-env path (homogeneous and N-variant); cast speed is unchanged (it already read its tree; the build is one-time). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
kernel_update_verts_and_aabbs gained a required batch_repr_env parameter, but the interactive viewer caller (RaycasterViewerPlugin -> Raycaster) was left unchanged, so any scene that opens the viewer raycaster would fail with a missing-argument error before the BVH could build. The viewer BVH is always per-env (n_batches=n_envs), so pass the identity mapping arange(n_envs) for batch_repr_env, reproducing the pre-change per-env semantics (i_env == i_b) exactly. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Milotrince
force-pushed
the
kashu/raycast-shared-static-bvh
branch
from
July 7, 2026 21:15
4b1e870 to
1348ed8
Compare
… opt-in flag) Replace RigidOptions.shared_static_raycast_bvh (opt-in, caller-guaranteed) with automatic, always-correct detection of env-identical static collision geometry. Why: the flag had two failure modes. It silently did nothing on any solver with a movable link (the common robot-on-terrain case), and with the flag on a legal per-env set_pos on a fixed body silently produced wrong results -- the range-only grouping could not see per-env vertex divergence, so every env read env 0's geometry. How: - Group envs by an exact per-env signature: the active geom ranges (when link info is batched) plus the per-env free-vert positions. fixed_verts_state is env-shared, so it is omitted. Two envs share a tree iff every input to update_aabbs matches -> a wrong merge is impossible; failure modes reduce to tree count, never correctness. - Refresh the free verts (existing kernel_update_all_verts) before grouping, since they are B-sized and independent of tree count; then size the BVH to the group count and build. Split the AABB build out of the fused kernel via a new kernel_update_aabbs so the static path can group in between. - Re-group on the existing rebuild trigger (the GEOMETRY subscriber fires on set_pos/set_quat/ set_vverts; reset forces a rebuild), reallocating the AABB+LBVH pair only when the group count changes. A divergent set_pos now splits to per-env trees and stays correct; an identical reset merges back to one. - Movable solvers and the visual BVH stay per-env; the visual path keeps the runtime all-identical shared-cast detection. Mixed static/dynamic scenes fall back to per-env automatically (documented; extending the win to the static portion needs a static/dynamic BVH split, left as future work). Tests: shared (n_envs 0 and 4) -> 1 tree; N variants -> N trees with block routing; per-env set_pos divergence -> per-env trees with correct per-env distances (guards the silent-corruption bug) and a reset that merges back to one; mixed static/dynamic -> per-env fallback. Existing lidar_bvh and heterogeneous_object tests updated for the automatic behavior and now assert the tree count.
Milotrince
force-pushed
the
kashu/raycast-shared-static-bvh
branch
from
July 7, 2026 21:52
1348ed8 to
7c343ed
Compare
Member
|
I removed the |
Milotrince
added
the
ready for review
Port the automatic shared static collision BVH onto the per-component rigid data repack (Genesis-Embodied-AI#3060) and the relocated sensor tests (Genesis-Embodied-AI#3037): - Kernels take the DynState / DynInfo aggregates; the per-face AABB update is factored into a shared func so the new grouped build kernel (kernel_update_grouped_aabbs) reuses it, and the fused per-env kernel keeps its signature (the viewer plugin caller is untouched). - The BVH traversal funcs now take a tree slot index for nodes / morton codes and the env index for leaf verts. Routing the single batch index through the group id read the wrong env's verts whenever the env -> group mapping was a non-monotonic permutation (covered by the non-monotonic set_pos phase of test_shared_static_bvh_regroup). - Env grouping compares the active geom ranges and free-vert positions exactly, with a zero-copy identical-envs short-circuit and integer / float parts grouped separately (a mixed concatenation would promote the ranges to float). The representative env is derived with a stable argsort, portable across CUDA / CPU / Metal. - Static collision BVH allocations are pooled per tree count: quadrants fields live for the whole process and template kernels specialize per field instance, so per-regroup reallocation would leak GPU memory and recompile the cast kernels on every diverge / reset cycle. - Tests live in tests/sensors/test_raycaster.py; the sharing, regroup and merge-back scenarios are packed into one parametrized test, and test_heterogeneous_object now runs 6 envs so the 3-variants < n_envs collapse is observable.
…e mapping. The tactile sensors (contact depth probe, kinematic taxel, elastomer taxel) traverse the raycaster's collision BVH with the env index as the tree batch index. Once a fully-fixed solver shares that BVH across envs, the tree buffer holds one slot per distinct geometry, so the per-env indexing reads out of bounds. The BVH traversal funcs now take the tree slot for nodes / morton codes, resolved through the collision entry's env_bvh_idx, while leaf triangles and per-env masks keep the env index.
…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 (RaycastContext.activate): - 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 - The split lives in the shared RaycastContext (activate / update), so a Raycaster and a DepthCamera on the same solver share the one static + dynamic BVH pair. Each subset BVH is built over compacted AABB slots: a `face_ids` tensor maps a leaf slot to the global face index; bvh_ray_cast remaps after reading the morton-code primitive id and update_aabbs iterates the slots. `n_triangles` in bvh_ray_cast derives from the morton-codes shape (the BVH's own leaf count) instead of the solver-global face count. - The remap is gated by a compile-time `is_remapped` flag: a BVH covering every face in order (pure-static / pure-dynamic solver, viewer pick, combined consumers below) skips both the map allocation and the per-leaf indirection, compiling to the exact pre-split kernels. - An entry is static when every face it covers sits on a fixed link. This per-face criterion also upgrades the single-BVH case: a movable link without collision geoms no longer disables the static-rebuild skip. - 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. - Tactile sensors in contact_depth_query="raycast" mode walk one tree over every collision face (collision_bvh_context), which the split would tear apart. They declare require_combined_collision_bvh() at construction time (before any consumer activates the context), keeping today's single combined BVH in such scenes; collision_bvh_context now raises unless a consumer declared that need. Merge correctness for no_hit_value < max_range - Splitting a robot-on-terrain collision scene turns a former single-BVH cast into a two-BVH merge. The merge overwrites a slot only when it finds a closer hit, so a first-pass miss that stamped `no_hit_value` directly would beat a real hit from a later pass whenever no_hit_value < max_range (e.g. 0/-1) -- previously guarded by a hard `no_hit_value >= max_range` check that would now newly reject those (valid, formerly single-BVH) configs. - write_ray_hit instead seeds an intermediate miss with the `max_range` sentinel (a real hit is always strictly below it, so it loses every distance comparison) and stamps `no_hit_value` only on the final pass (is_last), over any slot still at the sentinel. A single-BVH cast is both first and last, so it writes no_hit_value directly -- unchanged. The `no_hit_value >= max_range` check is removed; this also lifts the same restriction from the pre-existing visual+collision merge. The debug-draw hit reconstruction for distances-only sensors now discriminates misses by equality with no_hit_value instead of an ordering test, which the lifted restriction would break. Tests - tests/sensors/test_raycaster.py::test_static_dynamic_bvh_split_merge (n_envs in {0, 2}): asserts the split structure (one static + one dynamic collision entry, static shared across envs), the three-pass merge (static + dynamic + visual) reporting the closest hit as boxes enter/leave a ray with no_hit_value=-1 < max_range, the rebuild-on-explicit-set_pos of the static subset, the settle-to-no_hit_value when every pass misses, and the static skip across physics-driven motion; a distances-only sensor shares the chain. - tests/sensors/test_raycaster.py:: test_tactile_raycast_consumer_forces_combined_collision_bvh: a tactile probe in raycast mode keeps the mixed scene on one combined collision BVH and both consumers stay functional. - tests/sensors/test_api.py::test_shared_context updated to expect the static+dynamic split (2 entries). 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. Measured on this branch (A/B/A, default Terrain with 41.7k static faces + Franka, 2.3k rays/env, distances-only): CUDA 1024 envs 76.8 -> 9.4 ms/step (8.2x); CPU single env 8.5 -> 2.8 ms/step (3.1x). Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…shared-static-bvh Compose the static/dynamic collision BVH split (Genesis-Embodied-AI#2878) with the shared static BVH grouping: the split's static face subset is now the groupable entry, so its trees collapse across envs even when movable links share the solver. The grouped build kernel and the traversal remap compose through the shared per-face AABB func (face_ids / is_remapped), the grouped allocation pool sizes slots to the subset, and the cast chain carries is_merge / is_last / is_remapped / is_env_major together. The split test asserts the static subset collapses to one tree while the dynamic subset stays per env.
…ks and free verts.
Re-express the shared static BVH grouping on the landed static/dynamic split (Genesis-Embodied-AI#3078): - The grouped build kernels adopt the landed subset architecture: they fit compacted AABB slots per tree and rely on the post-build leaf payload rewrite (kernel_remap_leaf_faces), so the traversals stay subset-agnostic and the cast kernels only gain the env routing (env_bvh_idx, env_cast_order, is_env_major replacing shared_bvh). - An is_env_uniform entry collapses to a single tree with the grouping signature left unread; other static entries group by exact equality of their subset's build inputs, restricted by the signature masks so the movable subset's divergence under physics cannot split them. - The tactile probes route each of the two folded trees through its entry's env_bvh_idx (the tree slot for nodes and morton codes, the env for leaf triangles and candidate masks). - The link-filtered rebuild subscriptions and the refresh-once flow of the landed update() are kept; the standalone vert refresh serves the grouped static entries. - Tests follow suit: tree-count assertions replace the retired shared_across_envs flag, the split test keeps its phases with the static subset asserted at one tree throughout, the heterogeneous test runs 6 envs to expose the 3-variants collapse, and the regroup test covers a mixed solver (static-subset masks) with a kinematic visual entity (grouped visual splits with interleaved membership). The monkeypatch-based spies are removed from tests/sensors: integration tests assert observable behavior.
duburcqa
changed the title
duburcqa
changed the title
duburcqa
changed the title
duburcqa
approved these changes
Jul 21, 2026
|
🔴 Benchmark Regression Detected ➡️ Report |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
The raycast BVHs (
Raycaster/DepthCamera/Lidar, also traversed by the tactile probes) allocate one tree per env: nodes, AABBs, morton codes and the radix-sort scratch are all(n_envs, ...). For a high-poly static terrain this replication dominates GPU memory and OOMs at a few thousand envs, even though the per-env trees are identical or fall into a handful of variants.Building on the static/dynamic split of #3078, this PR makes every static BVH automatically allocate one tree per distinct per-env geometry instead of one per env: identical envs share a single tree, N heterogeneous variants get N trees. It applies to the static collision subset of a mixed robot-on-terrain solver, to fully static solvers, and to the static visual BVH.
Envs are grouped by exact equality of every per-env input of the tree build - active geom ranges and free-vert positions for collision, vgeom poses and custom vverts for visual, each restricted to the entry's own subset - so a wrong merge is impossible and the grouping is blind to the movable geometry diverging under physics. An
is_env_uniformentry (#3078) collapses to one tree with the signature left unread. A per-envset_poson fixed geometry regroups on rebuild (up to one tree per env) and an identical reset merges back; tree allocations are pooled per tree count so diverge/reset cycles reuse them. Two routing arrays drive the cast:env_bvh_idx(tree slot per env) andenv_cast_order(envs sorted by tree, keeping the env-major thread mapping coalesced even when group membership interleaves). Tactile probes traverse both folded trees through the same routing.Related Issue
None.
Motivation and Context
Massively parallel RL on static terrain: static-BVH memory drops by the ratio of env count to distinct-geometry count, on top of #3078's rebuild-cost win.
Measured on the #3078 benchmark scene (default Terrain, 41.7k static faces + Franka, 2.3k rays/env, distances-only), interleaved A/B/A/B against main:
scene.reset()93.8 -> 29.5 ms (3.2x: one static tree rebuilt instead of 1024, and the identical-envs short-circuit replaces the per-env AABB comparison).On a fully static 4.6k-face terrain scene against per-env trees, 1284 -> 497 MB; per-env trees OOM at 4096 envs while the shared tree runs.
How Has This Been / Can This Be Tested?
tests/sensors/test_raycaster.py::test_shared_static_bvh_regroup(new, n_envs 0 and 4, mixed solver): collision static subset and kinematic visual BVH collapse to one tree each with analytic distances; a non-monotonic per-envset_possplits the static subset to per-env trees with per-env distances; an interleaved per-envset_poson the visual entity splits the visual BVH into two trees; a per-env divergence of the movable entity leaves the restored static subset on one shared tree; reset merges everything back.tests/sensors/test_raycaster.py::test_static_dynamic_bvh_split_mergeasserts the static subset holds a single tree while the dynamic subset stays per env, through every merge/rebuild phase of [MISC] Speed up raycast sensors by splitting the rigid collision BVH into static and dynamic subsets. #3078 (tactile fold included).tests/sensors/test_raycaster.py::test_heterogeneous_objectnow runs 6 envs and asserts the collapse to 3 trees with per-variant distances.tests/sensors/test_tactile.py::test_raycast_probe_on_fully_fixed_solver(new): tactile probes traverse the shared tree through the env routing.tests/sensors/test_raycaster.py+test_tactile.py+test_api.py+tests/core/test_misc.pypass on CPU, CUDA and Metal backends.Checklist:
Submitting Code Changessection of CONTRIBUTING document.