newton/_src/utils/contact_sensor.py (7)

120-133: Nice: populate_contacts docstring adds missing context.

The function is self-explanatory now. Consider adding a one-liner in ContactSensor.eval referencing this helper for discoverability, but optional.

---

138-164: ContactSensor docstring is solid; minor nuance on matrix shape when pruning is disabled.

Great high-level overview. One nuance: when prune_noncolliding=False, the code currently skips counterparts that yield no shape pairs (e.g., identical shape vs shape), potentially making the number of columns smaller than “n_counterparts”. Either (a) include a zero column for such counterparts (my preferred), or (b) amend the docstring to say columns equal the number of counterparts that yield at least one shape pair. See code-level suggestion below on how to keep zero columns.

Would you like me to propagate the chosen convention to the docs and add a small example?

---

208-210: Unify vec3 dtype across allocation, attribute type, and return type.

You allocate with dtype=wp.vec3 but annotate/return as wp.vec3f. While Warp aliases may hide this at runtime, keeping a single dtype avoids confusion.

Suggested diff to standardize on vec3f (change both type hints and allocation):

-        self.net_force: wp.array2d(dtype=wp.vec3)
+        self.net_force: wp.array2d(dtype=wp.vec3f)
@@
-        self._net_force = wp.zeros(self.shape[0] * self.shape[1], dtype=wp.vec3, device=self.device)
+        self._net_force = wp.zeros(self.shape[0] * self.shape[1], dtype=wp.vec3f, device=self.device)
@@
-    def get_total_force(self) -> wp.array2d(dtype=wp.vec3f):
+    def get_total_force(self) -> wp.array2d(dtype=wp.vec3f):

Alternatively, standardize on wp.vec3 everywhere—just be consistent.

If you’re unsure about aliasing guarantees in your installed Warp version, I can check the exact behavior and recommend one canonical dtype.

Also applies to: 261-263, 276-283

---

109-109: Nit: avoid bool(i) inside kernel loop.

Using bool(i) is clever but cryptic. Make the condition explicit for readability.

-            wp.atomic_add(net_force, force_acc, wp.where(bool(i), -force, force))
+            wp.atomic_add(net_force, force_acc, wp.where(i == 1, -force, force))

---

265-274: eval docstring is accurate; consider validating contacts freshness (optional).

As a nicety, you could assert contacts.n_contacts is set or recommend calling populate_contacts before eval when verbose, but not mandatory.

---

277-283: Method name get_total_force is misleading; it returns the full matrix.

Two options:

• Rename to get_force_matrix() or get_net_force().
• Or adjust docstring to clarify it returns the net force matrix and, when include_total=True, the first column contains totals.

Minimal doc tweak:

-        Get the total net force measured by the contact sensor.
+        Get the net force matrix measured by the contact sensor.
+        When ``include_total`` is True, column 0 contains the per-sensor total.

I can update references across the repo if you choose to rename the method.

---

401-414: _lol_to_arrays returns None on empty input; ensure callers never pass Nones to kernels.

Currently, _assemble_sensor_mappings builds non-empty sp_reading as long as at least one reading exists (MatchAny ensures this). That said, adding a defensive fallback avoids hard-to-debug Warp errors.

Optional safeguard:

-        self.sp_reading, self.sp_ep_offset, self.sp_ep_count = _lol_to_arrays(sp_reading, wp.vec2i, device=self.device)
+        self.sp_reading, self.sp_ep_offset, self.sp_ep_count = _lol_to_arrays(sp_reading, wp.vec2i, device=self.device)
+        if self.sp_reading is None:
+            # Ensure valid empty arrays to satisfy kernel signatures
+            self.sp_reading = wp.empty(0, dtype=wp.vec2i, device=self.device)
+            self.sp_ep_offset = wp.empty(0, dtype=wp.int32, device=self.device)
+            self.sp_ep_count = wp.empty(0, dtype=wp.int32, device=self.device)

If you prefer assertions instead of fallbacks, I can add asserts documenting invariants established by _assemble_sensor_mappings.

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

newton/_src/utils/contact_sensor.py (5)

31-36: Docstring for MatchKind is clear and precise.

The updated description concisely explains the enum semantics. No action needed.

---

112-114: Good: MatchAny sentinel clarified.

Keeping a simple sentinel class with a brief docstring is fine and matches how you use it as a marker (not instances).

---

177-194: Init-time API validation reads well.

The exclusivity rules and argument descriptions are clear and aligned with the new bodies/shapes split.

---

308-315: get_colliding_sps correctness looks good; complexity OK for expected sizes.

The intersection-based pruning is clear and efficient for a precomputed set. No changes needed.

---

358-376: Kernel launch arguments consistent with solver_mujoco.update_contacts.

The shapes and dtypes of pair, normal, and force match the kernel expectations based on the converter kernel. Good.

newton/_src/utils/contact_sensor.py (5)

138-156: Clarify “sparse” wording to avoid confusion with storage format.

The force matrix is conceptually sparse under pruning, but the storage is still a dense wp.array2d with width max_active_counterparts. Consider tightening the wording to reflect conceptual sparsity to avoid users expecting a sparse container.

Proposed doc tweak:

-    If ``prune_noncolliding`` is True, the force matrix will be sparse, containing only readings for shape pairs that
-    can collide. In this case, force matrix will have as many columns as the maximum number of active counterparts
-    for any sensing object, and the ``reading_indices`` attribute can be used to recover the active counterparts
-    for each sensing object.
+    If ``prune_noncolliding`` is True, the force matrix is conceptually sparse: it contains readings only for shape
+    pairs that can collide. The stored array remains dense with width equal to the maximum number of active
+    counterparts across sensing objects; use ``reading_indices`` to recover which counterparts are active per row.

---

202-203: Good fix: reading_indices now correctly typed as list[list[int]].

This matches what _assemble_sensor_mappings returns and the intended “indices into counterparts” usage. Consider adding “indices into self.counterparts” for extra clarity.

-        """List of active counterpart indices per sensing object."""
+        """Indices into ``self.counterparts`` for the active counterparts of each sensing object."""

---

210-213: Avoid callable-style types in annotations; use wp.array instead of wp.array2d(...).

wp.array2d(dtype=wp.vec3) is an invocation, not a type. It can confuse type checkers and isn’t valid as a return annotation. Prefer annotating with wp.array (or Any) and expressing dtype/shape in the docstring.

-        self.net_force: wp.array2d(dtype=wp.vec3)
+        self.net_force: wp.array

And for get_total_force (see comment on Lines 279–285).

---

246-255: Edge case: empty sensing selection leads to a runtime error in _assemble_sensor_mappings.

If the patterns select zero sensing objects, counterpart_indices is empty and max(map(len, counterpart_indices)) raises ValueError. Consider validating early and raising a clear error when no sensing objects match.

Example guard:

         sp_sorted, sp_reading, self.shape, self.reading_indices, self.sensing_objs, self.counterparts = (
             self._assemble_sensor_mappings(
                 ...
             )
         )
+        if self.shape[0] == 0:
+            raise ValueError("No sensing objects matched the provided patterns.")

---

279-285: Return annotation should not call wp.array2d; annotate with wp.array.

Similar to the instance attribute, using a call expression in a return type is not valid. Keep shape/dtype details in the docstring.

-    def get_total_force(self) -> wp.array2d(dtype=wp.vec3):
+    def get_total_force(self) -> wp.array:
         """Get the total net force measured by the contact sensor.
 
         Returns:
-            The net force array, shaped according to the sensor configuration.
+            A 2-D warp array (dtype=wp.vec3) shaped according to the sensor configuration.
         """

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.