Add ParticleMeshCounter Warp utility for particle-in-mesh counting#6090
Add ParticleMeshCounter Warp utility for particle-in-mesh counting#6090maxkra15 wants to merge 2 commits into
Conversation
![isaaclab-review-bot[bot]](https://avatars.githubusercontent.com/u/157846462?s=60&v=4){kind=small}
There was a problem hiding this comment.
🤖 Isaac Lab Review Bot — PR #6090
Summary: Well-structured Warp utility with comprehensive tests. The architecture cleanly follows existing isaaclab.utils.warp patterns (lazy export via __init__.pyi, kernel + class in a dedicated module file). The winding-number approach is the right choice for robust containment testing.
🔵 Suggestion: Validate mesh factory parameters for degenerate inputs
File: source/isaaclab/isaaclab/utils/warp/particle_mesh.py (lines ~346-396)
make_frustum_region_mesh validates num_segments >= 3 but does not validate:
radius_bottom > 0andradius_top > 0(zero radius produces a degenerate cone tip)z_bottom < z_top(inverted range produces an inside-out or zero-volume mesh)
Similarly make_box_region_mesh does not check for non-positive half-extents.
Degenerate meshes will silently produce incorrect winding numbers. A brief guard with ValueError would match the existing num_segments validation pattern:
if radius_bottom <= 0 or radius_top <= 0:
raise ValueError(f"Radii must be positive, got bottom={radius_bottom}, top={radius_top}.")
if z_bottom >= z_top:
raise ValueError(f"`z_bottom` must be < `z_top`, got {z_bottom} >= {z_top}.")🟡 Warning: Returned mask buffer is reused — easy to introduce subtle bugs
File: source/isaaclab/isaaclab/utils/warp/particle_mesh.py (line ~213)
compute_inside_mask returns the internal _inside_torch buffer directly (documented in the docstring). However, count(..., return_mask=True) returns inside > 0.5 which does allocate a new tensor. The asymmetry between the public-facing count (safe, allocates) and semi-public compute_inside_mask (reuses buffer) is a potential footgun.
Consider either:
- Making
compute_inside_maskprivate (_compute_inside_mask) sincecount()is the primary API, or - Always returning a clone/detached view from public methods.
This is documented behavior so it's not a bug — just a usability concern for downstream consumers.
🔵 Suggestion: Axis convention note for region transforms
File: source/isaaclab/isaaclab/utils/warp/particle_mesh.py (line ~155)
The API uses shape (num_regions, num_envs, 3) for region transforms but (num_envs, num_particles, 3) for particle positions. Both are valid choices (region-major vs. env-major indexing for different access patterns in the kernel), but a brief note in the class docstring explaining the rationale for the transposed dimension ordering would help users avoid shape confusion.
🔵 Suggestion: Consider documenting the sign < 0.0 convention
File: source/isaaclab/isaaclab/utils/warp/particle_mesh.py (line ~82)
if query.result and query.sign < 0.0:
flag = wp.float32(1.0)The sign < 0.0 → inside convention for wp.mesh_query_point_sign_winding_number is correct per Warp's API, but a brief inline comment would help reviewers and future maintainers unfamiliar with Warp's sign convention. Something like:
# Warp convention: negative sign = point is inside the mesh
if query.result and query.sign < 0.0:✅ Strengths
- Excellent test coverage: Analytic ground-truth validation for both box and frustum, multi-env transforms, rotation tests, NaN/Inf sanitization, buffer reuse, prebuilt mesh support, and input validation — all parametrized across CPU and CUDA.
- Clean BVH reuse: Building the mesh BVH once and only updating transforms per-call is the right performance pattern.
- Non-finite sanitization: Graceful handling of NaN/Inf particles prevents silent corruption of counts.
- Proper lazy export integration: Only
__init__.pyiis modified, consistent with thelazy_export()pattern.
Verdict: Solid addition. The suggestions above are quality-of-life improvements, not blockers. No critical issues found.
Update (24c2644): All four suggestions from the initial review have been addressed in this push:
- ✅ Degenerate input validation —
make_box_region_meshnow raisesValueErrorfor non-positive half-extents;make_frustum_region_meshvalidates positive radii andz_bottom < z_top. Corresponding test cases added (test_box_rejects_non_positive_half_extents,test_frustum_rejects_non_positive_radius,test_frustum_rejects_inverted_z). - ✅
_compute_inside_maskmade private — Renamed fromcompute_inside_maskto_compute_inside_mask, clarifying thatcount()is the sole public API. - ✅ Dimension ordering note — The class docstring now includes a "Note on input layouts" paragraph explaining the region-major vs. env-major transposition.
- ✅
sign < 0.0convention comment — Inline comment added:# Warp convention: a negative winding-number sign means the point is inside the mesh.
Additionally, this push includes a separate Newton visualizer fix (set_camera_view now persists the eye/lookat to cfg and works without a prior initialize() call), with its own changelog entry and tests.
No new concerns. All previous suggestions resolved. 🟢
Adds a solver-agnostic Warp utility (isaaclab.utils.warp.ParticleMeshCounter) that counts particles inside closed region meshes via winding-number point queries, plus make_box_region_mesh / make_frustum_region_mesh helpers.
Greptile SummaryThis PR adds
Confidence Score: 3/5The utility is correct for the common case (tuple inputs built by the helper factories), but a pre-built wp.Mesh without winding-number support will silently return all-zero counts for every call with no error or warning. The winding-number sign convention, triangle orientations, and buffer-reuse logic are all correct and well-tested. The gap is in _as_winding_mesh: a pre-built wp.Mesh that lacks support_winding_number=True is accepted without complaint, causing mesh_query_point_sign_winding_number to produce wrong signs throughout the kernel — counts stay zero regardless of where particles are. This is an invisible failure mode for any downstream reward signal that relies on particle counts. particle_mesh.py — specifically _as_winding_mesh and _prepare_region_transforms Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A["counter.count(...)"] --> B["compute_inside_mask()"]
B --> C{"Validate particle_positions shape"}
C -- invalid --> ERR1["ValueError"]
C -- valid --> D["_prepare_region_transforms()"]
D --> E{"region_pos 2-D? Broadcast to (R,E,3)"}
E --> F{"Validate final shape"}
F -- invalid --> ERR2["ValueError"]
F -- valid --> G["_ensure_buffers(num_particles)"]
G --> H{"Same num_particles?"}
H -- yes --> I["Reuse buffer"]
H -- no --> J["Allocate new torch.zeros buffer"]
I --> K["wp.launch kernel (num_envs, num_particles, num_regions)"]
J --> K
K --> L["Transform particle to region local frame"]
L --> M["mesh_query_point_sign_winding_number"]
M --> N{"result AND sign < 0?"}
N -- yes --> O["inside = 1.0"]
N -- no --> P["inside = 0.0"]
O --> Q["Return _inside_torch (shared buffer)"]
P --> Q
Q --> R["sum(dim=1) → counts (num_envs, num_regions)"]
R --> S{"return_mask?"}
S -- yes --> T["Return counts, inside > 0.5"]
S -- no --> U["Return counts"]
Reviews (1): Last reviewed commit: "Add ParticleMeshCounter Warp utility for..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| def _as_winding_mesh(self, mesh: wp.Mesh | tuple[np.ndarray, np.ndarray]) -> wp.Mesh: | ||
| """Return a winding-number-capable :class:`warp.Mesh` for ``mesh``.""" | ||
| if isinstance(mesh, wp.Mesh): | ||
| return mesh |
There was a problem hiding this comment.
_as_winding_mesh passes a pre-built wp.Mesh straight through without checking that it was constructed with support_winding_number=True. Calling mesh_query_point_sign_winding_number on a mesh without that BVH will silently return an incorrect sign for every particle — every containment test returns False with no error or warning, making counts always zero and hard to debug.
| def _as_winding_mesh(self, mesh: wp.Mesh | tuple[np.ndarray, np.ndarray]) -> wp.Mesh: | |
| """Return a winding-number-capable :class:`warp.Mesh` for ``mesh``.""" | |
| if isinstance(mesh, wp.Mesh): | |
| return mesh | |
| def _as_winding_mesh(self, mesh: wp.Mesh | tuple[np.ndarray, np.ndarray]) -> wp.Mesh: | |
| """Return a winding-number-capable :class:`warp.Mesh` for ``mesh``.""" | |
| if isinstance(mesh, wp.Mesh): | |
| if not mesh.support_winding_number: | |
| raise ValueError( | |
| "Pre-built wp.Mesh must be created with `support_winding_number=True`;" | |
| " mesh_query_point_sign_winding_number requires the winding-number BVH." | |
| ) | |
| return mesh |
| region_pos = region_positions.to(device=self._device, dtype=torch.float32) | ||
| if region_pos.dim() == 2: | ||
| region_pos = region_pos.unsqueeze(1).expand(self._num_regions, self._num_envs, 3) | ||
| if tuple(region_pos.shape) != (self._num_regions, self._num_envs, 3): | ||
| raise ValueError( | ||
| f"`region_positions` must broadcast to (num_regions={self._num_regions}," | ||
| f" num_envs={self._num_envs}, 3), got {tuple(region_positions.shape)}." | ||
| ) | ||
|
|
||
| if region_orientations is None: | ||
| region_quat = torch.zeros((self._num_regions, self._num_envs, 4), device=self._device, dtype=torch.float32) | ||
| region_quat[..., 3] = 1.0 | ||
| else: | ||
| region_quat = region_orientations.to(device=self._device, dtype=torch.float32) | ||
| if region_quat.dim() == 2: | ||
| region_quat = region_quat.unsqueeze(1).expand(self._num_regions, self._num_envs, 4) | ||
| if tuple(region_quat.shape) != (self._num_regions, self._num_envs, 4): |
There was a problem hiding this comment.
RuntimeError
When region_pos.dim() == 2 but the first dimension doesn't equal num_regions, the expand() call throws a RuntimeError from PyTorch's internals before the friendly ValueError check can fire. The same gap applies to region_quat. Validating the shape before expanding keeps the error message consistent with the one already produced for the 3-D case.
| region_pos = region_positions.to(device=self._device, dtype=torch.float32) | |
| if region_pos.dim() == 2: | |
| region_pos = region_pos.unsqueeze(1).expand(self._num_regions, self._num_envs, 3) | |
| if tuple(region_pos.shape) != (self._num_regions, self._num_envs, 3): | |
| raise ValueError( | |
| f"`region_positions` must broadcast to (num_regions={self._num_regions}," | |
| f" num_envs={self._num_envs}, 3), got {tuple(region_positions.shape)}." | |
| ) | |
| if region_orientations is None: | |
| region_quat = torch.zeros((self._num_regions, self._num_envs, 4), device=self._device, dtype=torch.float32) | |
| region_quat[..., 3] = 1.0 | |
| else: | |
| region_quat = region_orientations.to(device=self._device, dtype=torch.float32) | |
| if region_quat.dim() == 2: | |
| region_quat = region_quat.unsqueeze(1).expand(self._num_regions, self._num_envs, 4) | |
| if tuple(region_quat.shape) != (self._num_regions, self._num_envs, 4): | |
| region_pos = region_positions.to(device=self._device, dtype=torch.float32) | |
| if region_pos.dim() == 2: | |
| if tuple(region_pos.shape) != (self._num_regions, 3): | |
| raise ValueError( | |
| f"`region_positions` must broadcast to (num_regions={self._num_regions}," | |
| f" num_envs={self._num_envs}, 3), got {tuple(region_positions.shape)}." | |
| ) | |
| region_pos = region_pos.unsqueeze(1).expand(self._num_regions, self._num_envs, 3) | |
| if tuple(region_pos.shape) != (self._num_regions, self._num_envs, 3): | |
| raise ValueError( | |
| f"`region_positions` must broadcast to (num_regions={self._num_regions}," | |
| f" num_envs={self._num_envs}, 3), got {tuple(region_positions.shape)}." | |
| ) | |
| if region_orientations is None: | |
| region_quat = torch.zeros((self._num_regions, self._num_envs, 4), device=self._device, dtype=torch.float32) | |
| region_quat[..., 3] = 1.0 | |
| else: | |
| region_quat = region_orientations.to(device=self._device, dtype=torch.float32) | |
| if region_quat.dim() == 2: | |
| if tuple(region_quat.shape) != (self._num_regions, 4): | |
| raise ValueError( | |
| f"`region_orientations` must broadcast to (num_regions={self._num_regions}," | |
| f" num_envs={self._num_envs}, 4), got {tuple(region_orientations.shape)}." | |
| ) | |
| region_quat = region_quat.unsqueeze(1).expand(self._num_regions, self._num_envs, 4) | |
| if tuple(region_quat.shape) != (self._num_regions, self._num_envs, 4): |
| def compute_inside_mask( | ||
| self, | ||
| particle_positions: torch.Tensor, | ||
| region_positions: torch.Tensor, | ||
| region_orientations: torch.Tensor | None = None, | ||
| ) -> torch.Tensor: | ||
| """Compute the per-particle containment mask for every region. | ||
|
|
||
| Args: | ||
| particle_positions: Particle positions in a common frame, shape | ||
| ``(num_envs, num_particles, 3)`` [m]. | ||
| region_positions: Region origins in the same frame, shape ``(num_regions, num_envs, 3)`` | ||
| or ``(num_regions, 3)`` (broadcast) [m]. | ||
| region_orientations: Region orientations as ``(x, y, z, w)`` quaternions, shape | ||
| ``(num_regions, num_envs, 4)`` or ``(num_regions, 4)`` (broadcast), or ``None`` for | ||
| identity. | ||
|
|
||
| Returns: | ||
| The containment mask (``1.0`` inside, ``0.0`` outside), shape | ||
| ``(num_envs, num_particles, num_regions)``, float. The returned tensor is an internal, | ||
| reused buffer; clone it if you need to retain it across calls. | ||
| """ | ||
| points = particle_positions.to(device=self._device, dtype=torch.float32) | ||
| if points.dim() != 3 or points.shape[0] != self._num_envs or points.shape[2] != 3: | ||
| raise ValueError( | ||
| f"`particle_positions` must have shape (num_envs={self._num_envs}, num_particles, 3)," | ||
| f" got {tuple(particle_positions.shape)}." | ||
| ) | ||
| points = points.contiguous() | ||
| num_particles = points.shape[1] | ||
| region_pos, region_quat = self._prepare_region_transforms(region_positions, region_orientations) | ||
| self._ensure_buffers(num_particles) | ||
| wp.launch( | ||
| count_particles_in_meshes_kernel, | ||
| dim=(self._num_envs, num_particles, self._num_regions), | ||
| inputs=[ | ||
| wp.from_torch(points, dtype=wp.vec3), | ||
| self._mesh_ids, | ||
| wp.from_torch(region_pos, dtype=wp.vec3), | ||
| wp.from_torch(region_quat, dtype=wp.quat), | ||
| self._max_query_dist, | ||
| self._inside_wp, | ||
| ], | ||
| device=self._wp_device, | ||
| ) | ||
| return self._inside_torch |
There was a problem hiding this comment.
compute_inside_mask documents that callers must clone the returned tensor, but the count() method — the primary public API — transparently shields users from this. Any code path that calls compute_inside_mask directly and stores the result without cloning will see its tensor silently mutated on the next call. Consider naming the method with a leading underscore (e.g., _compute_inside_mask) to signal that it is an internal helper, or always returning self._inside_torch.clone() from the public path.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
- Rebuild region meshes with winding-number support so a pre-built wp.Mesh without it can no longer silently produce zero counts (Warp does not expose the flag for validation). - Make compute_inside_mask private; count(return_mask=True) is the public, allocation-safe mask API. - Validate box/frustum factory inputs and malformed 2-D region transforms. - Document the winding-sign convention and the region/particle layout.
maxkra15
force-pushed
the
max/warp-particle-mesh-counter
branch
from
50980aa to
24c2644
Compare
| # Disable Warp's init banner and initialize the module so the kernel below is registered at import | ||
| # time (matches the pattern used in :mod:`isaaclab.utils.warp.ops`). | ||
| wp.config.quiet = True | ||
| wp.init() |
There was a problem hiding this comment.
the init here might not be necessary
| if len(region_meshes) == 0: | ||
| raise ValueError("`region_meshes` must contain at least one region mesh.") | ||
| self._device = str(device) | ||
| self._wp_device = wp.device_from_torch(torch.device(self._device)) |
There was a problem hiding this comment.
some of the easily derivable field like _wp_device, I usually don't add fields whose semantic is not tightly related to class and is also easily derivable, if user needs it, let user derive themselves
| self._num_envs = int(num_envs) | ||
| self._max_query_dist = float(max_query_dist) | ||
| self._meshes: list[wp.Mesh] = [self._as_winding_mesh(mesh) for mesh in region_meshes] | ||
| self._num_regions = len(self._meshes) |
There was a problem hiding this comment.
similarity for num_regions
| ) | ||
| return region_pos.contiguous(), region_quat.contiguous() | ||
|
|
||
| def _ensure_buffers(self, num_particles: int) -> None: |
There was a problem hiding this comment.
I see there is only one use-case for _ensure_buffers, usually this type of naming tends to smell because it conveys the intent that I don't know if buffer is there or not, just ensure its exist If the behavior is deterministic, change to a more deterministic naming like initalize_buffers if the behavior is unknown lets understand the lifecycle better
| self._mesh_ids = wp.array([mesh.id for mesh in self._meshes], dtype=wp.uint64, device=self._wp_device) | ||
| # buffers sized on first use (depend on the per-call particle count) | ||
| self._num_particles = 0 | ||
| self._inside_torch: torch.Tensor | None = None |
There was a problem hiding this comment.
inside_torch and inside_warp is a bit weird, are you looking for something like ProxyArray?
ooctipus
left a comment
ooctipus
left a comment
There was a problem hiding this comment.
maybe we can take out of torch from this file all together.
2 participants
Description
This PR adds
isaaclab.utils.warp.ParticleMeshCounter, a fast, solver-agnostic Warp utility for countinghow many particles fall inside one or more closed (watertight) region meshes.
What's added
ParticleMeshCounter— counts particles inside region meshes using Warp's BVH-acceleratedwinding-number point query (
wp.mesh_query_point_sign_winding_number). It supports multiple,independently posed region meshes per environment (region geometry is fixed in its local frame and the
BVH is built once and reused; only per-environment transforms change between calls), sanitizes
non-finite particle positions, and returns per-region counts and/or the per-particle containment mask.
make_box_region_mesh/make_frustum_region_mesh— helpers that build watertight, outward-orientedregion meshes (axis-aligned boxes and capped circular frusta / "cup cavity" regions).
It lives next to the existing Warp helpers (
raycast_mesh,convert_to_warp_mesh) inisaaclab.utils.warp, as a single cohesiveparticle_mesh.pymodule (mirroring the package's existingfeature modules such as
proxy_array.py).Motivation / context
This enables training-time / privileged measurements such as "how many media particles are inside the
scoop bowl / source container / target container" without relying on hand-tuned analytic regions. It is
pure Warp + Torch with no physics-backend dependency, so it works for any particle source (MPM, PhysX
particles, point clouds).
Dependencies
None beyond what Isaac Lab already ships (Warp + Torch).
Fixes # (N/A — no tracking issue)
Type of change
Screenshots
N/A — non-visual utility (validated by unit tests against analytic ground truth).
Checklist
pre-commitchecks with./isaaclab.sh --formatconfig/extension.tomlfileCONTRIBUTORS.mdor my name already exists there