[PERF] Optional shared static-geometry raycast BVH (n_batches=1)

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 every
env's tree is identical and the cast already reads batch 0 when it detects
the trees match (#2867).

Add RigidOptions.shared_static_raycast_bvh (default False). When True, the
static (fully-fixed) collision BVH is allocated ONCE and shared by every env,
removing the n_envs-fold replication. update_aabbs now iterates the AABB
buffer's own batch dimension, and a single-batch tree is flagged shared so
the cast reads batch 0 for all envs.

It is opt-in rather than auto-detected because env-identity is a runtime
property: a per-env set_pos on a fixed body diverges the geometry after
build. The default path still handles that via the existing per-build
shared-across-envs comparison (which needs the per-env trees, so it cannot
itself be the basis for a 1-batch allocation). The flag is a caller guarantee
that the static collision geometry stays identical across envs (e.g. one
terrain shared by all envs).

Benchmark — raycast DepthCamera over an 8087-face static terrain (RTX 3080),
total GPU memory at 64x36:
   256 env :  775 ->  154 MB (5.0x)
  1024 env : 2884 ->  366 MB (7.9x)
  4096 env : OOM  -> runs (previously CUDA_ERROR_OUT_OF_MEMORY)
Depth is bit-identical to the per-env path; cast speed is unchanged (it
already read batch 0; the build is one-time).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: Kashu7100

genesis/engine/sensors/raycaster.py

@@ -120,7 +120,16 @@ def activate(self):
             maybe_static = all(link.is_fixed for link in solver.links)
             if isinstance(solver, RigidSolver):
                 n_faces = solver.faces_info.geom_idx.shape[0]
-                aabb = AABB(n_batches=n_envs, n_aabbs=n_faces)
+                # A static, env-identical collision mesh produces bit-identical per-env BVHs, so one shared tree
+                # serves every env (the cast reads batch 0) — dropping the n_envs-fold replication of nodes /
+                # aabbs / morton codes / radix-sort scratch, which dominates GPU memory for a high-poly terrain.
+                # This is opt-in (``shared_static_raycast_bvh``) rather than auto-detected because env-identity is
+                # a runtime property — a per-env ``set_pos`` on a fixed body diverges the geometry after build —
+                # so it cannot be proven at allocation time; the default path keeps per-env trees and the runtime
+                # shared-across-envs detection below. The flag is a caller guarantee the static geometry stays
+                # env-identical (e.g. a single terrain shared by all envs).
+                shared_static = maybe_static and solver._options.shared_static_raycast_bvh
+                aabb = AABB(n_batches=1 if shared_static else n_envs, n_aabbs=n_faces)
                 bvh = LBVH(aabb, max_n_query_result_per_aabb=0, n_radix_sort_groups=64)
                 self._bvh_contexts.append(BVHContext(solver, bvh, aabb, None, maybe_static))
             n_vfaces = solver.vfaces_info.vgeom_idx.shape[0]
@@ -201,7 +210,10 @@ def update(self):
                     and torch.equal(aabb_max, aabb_max[:1].expand_as(aabb_max))
                 )
             else:
-                entry.shared_across_envs = False
+                # A single-batch tree (n_batches==1, allocated for env-identical static geometry above) is shared
+                # by construction: the one tree must serve every env, so the cast reads batch 0 for all. The
+                # per-env dynamic case (n_batches==n_envs) is never shared.
+                entry.shared_across_envs = entry.aabb.n_batches == 1
 
     def reset(self, envs_idx):
         # A reset may change otherwise-static geometry (re-randomized terrain, teleported obstacles), so force every

genesis/options/solvers.py

@@ -480,6 +480,15 @@ class RigidOptions(Options):
         Broadphase traversal strategy. ``SAP`` (sweep-and-prune) or ``ALL_VS_ALL`` (parallel pair iteration). Defaults
         to ``None`` (auto: ``SAP`` on CPU or when hibernation/heterogeneous entities are enabled, ``ALL_VS_ALL`` on GPU
         otherwise). See ``gs.broadphase_traversal`` for details on each strategy.
+    shared_static_raycast_bvh : bool, optional
+        Optimization for raycast sensors (``Raycaster`` / ``DepthCamera``) over large static scenes. When True, the
+        raycast BVH built over the solver's static (fully fixed) collision geometry is allocated **once** and shared
+        by every env (each env reads the same tree), instead of replicating one tree per env. For a high-poly static
+        terrain this is the dominant raycast memory cost, so sharing it cuts total GPU memory several-fold and lifts
+        the env-count ceiling. It is a caller guarantee that the static collision geometry is **identical across
+        envs**: do not enable it if you give fixed entities per-env poses (e.g. per-env ``set_pos`` on a fixed body)
+        or per-env geometry, as every env would then incorrectly cast against env 0's tree. Defaults to False, which
+        keeps per-env trees and the runtime "shared across envs" auto-detection. Only affects raycasting, not physics.
 
     Warning
     -------
@@ -504,6 +513,10 @@ class RigidOptions(Options):
     batch_joints_info: StrictBool = False
     batch_dofs_info: StrictBool = False
 
+    # raycast: share one static-geometry BVH across envs instead of one per env (caller guarantees env-identical
+    # static collision geometry). See the class docstring.
+    shared_static_raycast_bvh: StrictBool = False
+
     # constraint solver
     constraint_solver: gs.constraint_solver = gs.constraint_solver.Newton
     iterations: PositiveInt = 50

genesis/utils/raycast_qd.py

@@ -248,7 +248,10 @@ def update_aabbs(
     one vertex buffer but activate different per-env geom ranges, it makes each env cast against only its own variant
     instead of the union of every variant.
     """
-    for i_b, i_f in qd.ndrange(free_verts_state.pos.shape[1], faces_info.verts_idx.shape[0]):
+    # Iterate over the AABB buffer's own batch dimension, not the solver's env count: a shared static BVH is
+    # allocated with n_batches=1 (one tree built from env 0's verts and read by every env), while a per-env BVH
+    # has n_batches=n_envs. For the latter the two are equal, so this is a no-op there.
+    for i_b, i_f in qd.ndrange(aabb_state.aabbs.shape[0], faces_info.verts_idx.shape[0]):
         aabb_state.aabbs[i_b, i_f].min.fill(qd.math.inf)
         aabb_state.aabbs[i_b, i_f].max.fill(-qd.math.inf)
 

tests/test_sensors.py

@@ -1174,6 +1174,41 @@ def test_raycaster_hits(show_viewer, n_envs):
     assert_allclose(grid_distances, grid_distances_ref, tol=1e-3)
 
 
+@pytest.mark.required
+def test_raycaster_shared_static_bvh(show_viewer):
+    # With RigidOptions.shared_static_raycast_bvh=True, env-identical static collision geometry is allocated as ONE
+    # shared BVH (n_batches=1) read by every env, instead of a per-env tree. Verify the allocation collapses, the
+    # cast is flagged shared, and the distances are finite and identical across envs (they must be, since every env
+    # sees the same geometry from the same sensor pose). The opt-in default (False) and its per-env-set_pos
+    # divergence path are covered by test_lidar_bvh_parallel_env.
+    scene = gs.Scene(
+        rigid_options=gs.options.RigidOptions(shared_static_raycast_bvh=True),
+        show_viewer=show_viewer,
+    )
+    scene.add_entity(gs.morphs.Plane())
+    scene.add_entity(gs.morphs.Box(size=(0.4, 0.4, 0.4), pos=(1.0, 0.0, 0.2), fixed=True))
+    depth_camera = scene.add_sensor(
+        gs.sensors.DepthCamera(
+            pattern=gs.sensors.raycaster.DepthCameraPattern(res=(8, 8)),
+            pos_offset=(0.0, 0.0, 1.0),
+            euler_offset=(0.0, 40.0, 0.0),
+        )
+    )
+    scene.build(n_envs=4)
+    scene.step()
+
+    # The single rigid collision BVH collapsed to one shared tree and the cast reads batch 0 for every env.
+    (collision_bvh,) = depth_camera._shared_context.bvh_contexts
+    assert collision_bvh.aabb.n_batches == 1
+    assert collision_bvh.shared_across_envs is True
+
+    distances = depth_camera.read().distances  # (n_envs, H, W)
+    assert torch.isfinite(distances).all()
+    assert (distances < depth_camera._options.max_range).any()  # some rays hit the box/plane
+    # Identical geometry + identical sensor pose -> every env's depth image matches env 0's.
+    assert torch.equal(distances, distances[:1].expand_as(distances))
+
+
 @pytest.mark.required
 @pytest.mark.parametrize("n_envs", [0, 2])
 @pytest.mark.parametrize("kin_raycastable", [True, False])