Skip to content

[PERF] Split rigid collision BVH into static + dynamic subsets (RPL multi-depth) #2878

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Kashu7100 wants to merge 1 commit into from
Closed

[PERF] Split rigid collision BVH into static + dynamic subsets (RPL multi-depth) #2878

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 37 additions & 14 deletions genesis/engine/sensors/raycaster.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,10 @@ class BVHContext:
# True when the geometry is bit-identical across envs, so the cast reads one shared copy (batch 0) with coalesced
# node loads instead of scattering over n_env identical trees. Recomputed on every rebuild.
shared_across_envs: bool = False
# Compacted face subset this collision BVH covers: face_ids[k] is the global face index at leaf slot k (1-D int
# device tensor). Identity arange(n_faces) for a single full-mesh BVH; None for visual entries (their kernels take
# no face_ids). Splitting into static/dynamic subsets is what lets the static tree stay skipped - see activate().
face_ids: torch.Tensor | None = None


class RaycastContext(SharedSensorContext):
Expand All @@ -84,6 +88,25 @@ def bvh_contexts(self) -> list[BVHContext]:
raise gs.GenesisException("RaycastContext queried before activation; no sensor declared a raycast need.")
return self._bvh_contexts

@staticmethod
def _partition_collision_faces(solver: "RigidSolver") -> list[tuple[torch.Tensor, bool]]:
"""Partition the solver's collision faces into static (fixed-link) and dynamic (movable-link) subsets,
returning ``(face_ids, maybe_static)`` per non-empty subset. A pure static/dynamic solver yields one entry
(a single full-mesh BVH); a mixed scene (robot on terrain) yields two. See :meth:`activate`.
"""
face_geom = qd_to_numpy(solver.faces_info.geom_idx).reshape(-1) # (n_faces,) global geom per face
geom_link = qd_to_numpy(solver.geoms_info.link_idx).reshape(-1) # (n_geoms,) global link per geom
link_fixed = np.array([bool(link.is_fixed) for link in solver.links], dtype=bool) # (n_links,)
face_static = link_fixed[geom_link[face_geom]] # (n_faces,) is this face on a fixed link?

out: list[tuple[torch.Tensor, bool]] = []
for is_static in (True, False):
sel = np.nonzero(face_static == is_static)[0]
if sel.size == 0:
continue
out.append((torch.as_tensor(sel, dtype=gs.tc_int, device=gs.device), bool(is_static)))
return out

@staticmethod
def _compute_visual_raycast_mask(solver: "KinematicSolver") -> np.ndarray:
"""Build a per-vface mask (int8, shape (n_vfaces,)) selecting vfaces opted into visual raycasting.
Expand Down Expand Up @@ -119,10 +142,14 @@ def activate(self):
# Applies to both the collision and the visual BVH.
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)
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))
# Static (fixed-link) faces get a BVH built once + skipped + shared across envs; dynamic (movable-link)
# faces get one that rebuilds each step. Cast separately, merged (is_merge) into one identical result.
# RPL "multi-depth" decomposition (arXiv:2602.03002); a pure static/dynamic solver is a single subset.
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=64)
self._bvh_contexts.append(BVHContext(solver, bvh, aabb, None, subset_static, face_ids=face_ids))
n_vfaces = solver.vfaces_info.vgeom_idx.shape[0]
if n_vfaces > 0:
mask = self._compute_visual_raycast_mask(solver)
Expand Down Expand Up @@ -168,6 +195,7 @@ def update(self):
free_verts_state=entry.solver.free_verts_state,
fixed_verts_state=entry.solver.fixed_verts_state,
links_info=entry.solver.links_info,
face_ids=entry.face_ids,
static_rigid_sim_config=entry.solver._static_rigid_sim_config,
aabb_state=entry.aabb,
)
Expand Down Expand Up @@ -323,14 +351,6 @@ def build(self):
self._shared_metadata.no_hit_values, self._options.no_hit_value
)

# Multi-BVH merge passes use raw distance comparison to pick the closer hit; this only works if no_hit_value >=
# max_range. The negated form also rejects NaN (every IEEE 754 comparison with NaN is False).
if len(self._shared_context.bvh_contexts) > 1 and not (self._options.no_hit_value >= self._options.max_range):
gs.raise_exception(
f"no_hit_value ({self._options.no_hit_value}) must be >= max_range ({self._options.max_range}) "
f"when multiple BVHs are active (the merge step compares raw distances)."
)

def _get_return_format(self) -> tuple[tuple[int, ...], ...]:
shape = self._options.pattern.return_shape
return ((*shape, 3), shape)
Expand Down Expand Up @@ -372,8 +392,9 @@ def _update_raw_data(
links_pos[:, group.sensor_cols, :] = pos
links_quat[:, group.sensor_cols, :] = quat

# First entry initializes the cache (is_merge=False, writes a hit or no_hit_value into every slot). Each
# subsequent entry merges in place (is_merge=True, writes only where it found a closer hit).
# Chain the entries into one output buffer: first initializes (is_merge=False), the rest merge closer hits
# (is_merge=True), the last (is_last) finalizes misses to no_hit_value - see write_ray_hit.
n_entries = len(bvh_contexts)
for i, entry in enumerate(bvh_contexts):
solver = entry.solver
args_common = (
Expand All @@ -393,6 +414,7 @@ def _update_raw_data(
raw_data_T,
gs.EPS,
i > 0,
i == n_entries - 1,
entry.shared_across_envs,
)
if entry.raycast_mask is None:
Expand All @@ -401,6 +423,7 @@ def _update_raw_data(
solver.free_verts_state,
solver.verts_info,
solver.faces_info,
entry.face_ids,
*args_common,
)
else:
Expand Down
62 changes: 46 additions & 16 deletions genesis/utils/raycast_qd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ def bvh_ray_cast(
verts_info: array_class.VertsInfo,
fixed_verts_state: array_class.VertsState,
free_verts_state: array_class.VertsState,
face_ids: qd.types.ndarray(ndim=1),
eps: float,
):
"""
Expand All @@ -70,7 +71,9 @@ def bvh_ray_cast(
hit_normal : qd.math.vec3
normal vector at hit point (zero vector if no hit)
"""
n_triangles = faces_info.verts_idx.shape[0]
# This BVH's leaf count, which is its morton-code count - NOT the solver's global face count: the BVH may cover a
# compacted face subset, so face_ids[leaf] remaps a subset-local leaf back to the global face below.
n_triangles = bvh_morton_codes.shape[1]

hit_face = -1
closest_distance = gs.qd_float(max_range)
Expand All @@ -92,9 +95,9 @@ def bvh_ray_cast(

if aabb_t >= 0.0 and aabb_t < closest_distance:
if node.left == -1: # Leaf node
# Get original triangle/face index
# The morton code carries the subset-local leaf; face_ids remaps it to the global face.
sorted_leaf_idx = node_idx - (n_triangles - 1)
i_f = qd.cast(bvh_morton_codes[i_b, sorted_leaf_idx][1], gs.qd_int)
i_f = qd.cast(face_ids[qd.cast(bvh_morton_codes[i_b, sorted_leaf_idx][1], gs.qd_int)], gs.qd_int)

# Get triangle vertices
tri_vertices = get_triangle_vertices(
Expand Down Expand Up @@ -237,20 +240,25 @@ def update_aabbs(
faces_info: array_class.FacesInfo,
geoms_info: array_class.GeomsInfo,
links_info: array_class.LinksInfo,
face_ids: qd.types.ndarray(ndim=1),
static_rigid_sim_config: qd.template(),
aabb_state: qd.template(),
):
"""Update per-face collision AABBs from current vertex positions.

AABB slot k holds the bounding box of the global face face_ids[k]; iterating this compacted subset scales the
rebuild + radix sort with the subset size, not every face in the solver (identity face_ids for a full-mesh BVH).

A face contributes to env i_b only if its geom lies in that env's active geom range (links_info.geom_start /
geom_end); otherwise its AABB is left inverted (unhittable) and skipped by ray queries. For a homogeneous solver
every geom is always in range, so this never excludes anything. For a heterogeneous solver, where all envs share
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]):
aabb_state.aabbs[i_b, i_f].min.fill(qd.math.inf)
aabb_state.aabbs[i_b, i_f].max.fill(-qd.math.inf)
for i_b, k in qd.ndrange(free_verts_state.pos.shape[1], face_ids.shape[0]):
i_f = face_ids[k]
aabb_state.aabbs[i_b, k].min.fill(qd.math.inf)
aabb_state.aabbs[i_b, k].max.fill(-qd.math.inf)

i_g = faces_info.geom_idx[i_f]
i_l = geoms_info.link_idx[i_g]
Expand All @@ -261,12 +269,12 @@ def update_aabbs(
i_fv = verts_info.verts_state_idx[i_v]
if verts_info.is_fixed[i_v]:
pos_v = fixed_verts_state.pos[i_fv]
aabb_state.aabbs[i_b, i_f].min = qd.min(aabb_state.aabbs[i_b, i_f].min, pos_v)
aabb_state.aabbs[i_b, i_f].max = qd.max(aabb_state.aabbs[i_b, i_f].max, pos_v)
aabb_state.aabbs[i_b, k].min = qd.min(aabb_state.aabbs[i_b, k].min, pos_v)
aabb_state.aabbs[i_b, k].max = qd.max(aabb_state.aabbs[i_b, k].max, pos_v)
else:
pos_v = free_verts_state.pos[i_fv, i_b]
aabb_state.aabbs[i_b, i_f].min = qd.min(aabb_state.aabbs[i_b, i_f].min, pos_v)
aabb_state.aabbs[i_b, i_f].max = qd.max(aabb_state.aabbs[i_b, i_f].max, pos_v)
aabb_state.aabbs[i_b, k].min = qd.min(aabb_state.aabbs[i_b, k].min, pos_v)
aabb_state.aabbs[i_b, k].max = qd.max(aabb_state.aabbs[i_b, k].max, pos_v)


@qd.kernel
Expand All @@ -278,6 +286,7 @@ def kernel_update_verts_and_aabbs(
free_verts_state: array_class.VertsState,
fixed_verts_state: array_class.VertsState,
links_info: array_class.LinksInfo,
face_ids: qd.types.ndarray(ndim=1),
static_rigid_sim_config: qd.template(),
aabb_state: qd.template(),
):
Expand All @@ -291,6 +300,7 @@ def kernel_update_verts_and_aabbs(
faces_info,
geoms_info,
links_info,
face_ids,
static_rigid_sim_config,
aabb_state,
)
Expand Down Expand Up @@ -445,6 +455,7 @@ def kernel_cast_ray(
free_verts_state: array_class.VertsState,
verts_info: array_class.VertsInfo,
faces_info: array_class.FacesInfo,
face_ids: qd.types.ndarray(ndim=1), # maps BVH leaf slot -> global face index (identity for a full-mesh BVH)
bvh_nodes: qd.template(),
bvh_morton_codes: qd.template(),
ray_start: qd.types.ndarray(ndim=1), # (3,)
Expand Down Expand Up @@ -485,6 +496,7 @@ def kernel_cast_ray(
verts_info=verts_info,
fixed_verts_state=fixed_verts_state,
free_verts_state=free_verts_state,
face_ids=face_ids,
eps=eps,
)
if cur_hit_face >= 0:
Expand All @@ -507,17 +519,20 @@ def write_ray_hit(
i_p_offset: int,
i_p_dist: int,
is_world_frame: qd.types.ndarray(ndim=1),
max_ranges: qd.types.ndarray(ndim=1),
no_hit_values: qd.types.ndarray(ndim=1),
output_hits: qd.types.ndarray(ndim=2),
eps: float,
is_merge: qd.template(),
is_last: qd.template(),
):
"""Common post-BVH write block for both collision and visual cast kernels.

`is_merge` is a compile-time flag. When False the function writes a value into every output slot (hit or
no_hit_value), initializing the cache. When True the function only writes when it found a closer hit than what
is already in the cache, so multiple BVH casts can be composed by chaining calls (first with is_merge=False,
subsequent with is_merge=True) into the same output buffer with no scratch storage.
`is_merge` / `is_last` mark a cast's position in a chain of BVH passes sharing one output buffer (first
is_merge=False, rest is_merge=True, final is_last). A miss must not beat a real hit from another pass for any
``no_hit_value`` (which may be < max_range): an intermediate miss seeds the slot with ``max_range`` - a hit is
always strictly below it, so it loses the distance comparison - and ``no_hit_value`` is stamped only on the final
pass, over any slot still at the sentinel.
"""
if hit_face >= 0 and (not is_merge or hit_distance < output_hits[i_p_dist, i_b]):
# Store distance at: cache_offset + (num_points_in_sensor * 3) + point_idx_in_sensor
Expand All @@ -534,11 +549,18 @@ def write_ray_hit(
output_hits[i_p_offset + i_p_sensor * 3 + 1, i_b] = hit_point.y
output_hits[i_p_offset + i_p_sensor * 3 + 2, i_b] = hit_point.z
elif not is_merge:
# No hit
# First-pass miss: zero the point; seed no_hit_value if single-BVH (also last), else the max_range sentinel.
output_hits[i_p_offset + i_p_sensor * 3 + 0, i_b] = 0.0
output_hits[i_p_offset + i_p_sensor * 3 + 1, i_b] = 0.0
output_hits[i_p_offset + i_p_sensor * 3 + 2, i_b] = 0.0
output_hits[i_p_dist, i_b] = no_hit_values[i_s]
if is_last:
output_hits[i_p_dist, i_b] = no_hit_values[i_s]
else:
output_hits[i_p_dist, i_b] = max_ranges[i_s]
elif is_last:
# Final-pass miss: a slot still at the sentinel means every pass missed -> stamp no_hit_value.
if output_hits[i_p_dist, i_b] >= max_ranges[i_s]:
output_hits[i_p_dist, i_b] = no_hit_values[i_s]


@qd.kernel
Expand All @@ -547,6 +569,7 @@ def kernel_cast_rays(
free_verts_state: array_class.VertsState,
verts_info: array_class.VertsInfo,
faces_info: array_class.FacesInfo,
face_ids: qd.types.ndarray(ndim=1), # maps BVH leaf slot -> global face index (identity for a full-mesh BVH)
bvh_nodes: qd.template(),
bvh_morton_codes: qd.template(), # maps sorted leaves to original triangle indices
links_pos: qd.types.ndarray(ndim=3), # [n_env, n_sensors, 3]
Expand All @@ -563,6 +586,7 @@ def kernel_cast_rays(
output_hits: qd.types.ndarray(ndim=2), # [total_cache_size, n_env]
eps: float,
is_merge: qd.template(),
is_last: qd.template(),
shared_bvh: qd.template(),
):
"""Cast rays against a collision-mesh BVH, accelerated by a BVH. See write_ray_hit for `is_merge` semantics.
Expand Down Expand Up @@ -613,6 +637,7 @@ def kernel_cast_rays(
verts_info=verts_info,
fixed_verts_state=fixed_verts_state,
free_verts_state=free_verts_state,
face_ids=face_ids,
eps=eps,
)

Expand All @@ -631,10 +656,12 @@ def kernel_cast_rays(
i_p_offset,
i_p_dist,
is_world_frame,
max_ranges,
no_hit_values,
output_hits,
eps,
is_merge,
is_last,
)


Expand All @@ -660,6 +687,7 @@ def kernel_cast_rays_visual(
output_hits: qd.types.ndarray(ndim=2),
eps: float,
is_merge: qd.template(),
is_last: qd.template(),
shared_bvh: qd.template(),
):
"""Visual-mesh variant of kernel_cast_rays. See kernel_cast_rays for shared_bvh and the thread mapping."""
Expand Down Expand Up @@ -715,8 +743,10 @@ def kernel_cast_rays_visual(
i_p_offset,
i_p_dist,
is_world_frame,
max_ranges,
no_hit_values,
output_hits,
eps,
is_merge,
is_last,
)
6 changes: 6 additions & 0 deletions genesis/vis/viewer_plugins/raycast.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from typing import TYPE_CHECKING

import numpy as np
import torch
from typing_extensions import override

import genesis as gs
Expand Down Expand Up @@ -51,6 +52,9 @@ def __init__(self, scene: "Scene"):
max_n_query_result_per_aabb=0, # Not used for ray queries
n_radix_sort_groups=min(64, n_faces),
)
# The viewer casts one BVH over the full mesh, so the leaf-slot -> global-face map the kernels take is the
# identity. (Sensors build compacted per-subset BVHs with a real map; see RaycastContext.)
self.face_ids = torch.arange(n_faces, dtype=gs.tc_int, device=gs.device)
self.result = array_class.get_raycast_result(n_envs_max)

self.update()
Expand All @@ -72,6 +76,7 @@ def update(self) -> None:
free_verts_state=self.solver.free_verts_state,
fixed_verts_state=self.solver.fixed_verts_state,
links_info=self.solver.links_info,
face_ids=self.face_ids,
static_rigid_sim_config=self.solver._static_rigid_sim_config,
aabb_state=self.aabb,
)
Expand Down Expand Up @@ -105,6 +110,7 @@ def cast(
self.solver.free_verts_state,
self.solver.verts_info,
self.solver.faces_info,
self.face_ids,
self.bvh.nodes,
self.bvh.morton_codes,
np.ascontiguousarray(ray_origin, dtype=gs.np_float),
Expand Down
Loading
Loading