Skip to content

Optimize SensorContact and simplify API#2135

Merged
camevor merged 10 commits into
newton-physics:mainfrom
camevor:optimize-sensor-contact
Mar 19, 2026
Merged

Optimize SensorContact and simplify API#2135
camevor merged 10 commits into
newton-physics:mainfrom
camevor:optimize-sensor-contact

Conversation

@camevor

@camevor camevor commented Mar 17, 2026

Copy link
Copy Markdown
Member

Description

Redesign SensorContact for clarity and performance. Replace the monolithic net_force array (where column 0 was optionally a total) with two explicit outputs: total_force (1D, per sensing object) and force_matrix (2D, per
counterpart), each nullable. Add flat metadata attributes (sensing_obj_idx, counterpart_indices, sensing_obj_type, counterpart_type) replacing the nested per-world tuple lists.

Internally, replace sorted pair-table binary search with row/col index maps, yielding 50-90x faster init and 1.4-2x faster update at 16k worlds.

All old attributes are preserved as deprecated shims with DeprecationWarning.

Test plan

uv run --extra dev -m newton.tests -k test_sensor_contact
uv run --extra dev -m newton.tests -k test_basic.example_sensor_contact

Checklist

  • New or existing tests cover these changes
  • The documentation is up to date with these changes
  • CHANGELOG.md has been updated (if user-facing change)

New feature / API change

# Before (deprecated)
sensor = SensorContact(model, sensing_obj_bodies=["a", "b"], counterpart_bodies="*", include_total=True)
sensor.net_force              # wp.array2d — column 0 is total, rest are counterparts
sensor.shape                  # (n_rows, n_cols)
sensor.sensing_objs           # list[list[tuple[int, ObjectType]]] — nested per world
sensor.counterparts           # list[list[tuple[int, ObjectType]]] — nested per world

# After
sensor = SensorContact(model, sensing_obj_bodies=["a", "b"], counterpart_bodies="*", measure_total=True)
sensor.total_force            # wp.array (n_sensing_objs,) vec3, or None if measure_total=False
sensor.force_matrix           # wp.array2d (n_sensing_objs, n_counterparts) vec3, or None if no counterparts
sensor.sensing_obj_idx        # [0, 1] — flat list, caller's order preserved when passing pre-resolved indices
sensor.sensing_obj_type       # literal, "body" or "shape"
sensor.counterpart_indices[0] # counterpart body indices for row 0
sensor.counterpart_type       # literal, "body" or "shape"

Benchmarks

SensorContact.__init__

These init benchmarks use lists of body indices instead of relying on label matching.

Benchmark before (ms) after (ms) speedup
Ant (4 bodies/world), 16384 worlds, total only 1253 16.3 77x
Ant (4 sensing vs 9 counterpart bodies/world), 16384 worlds 3854 43.1 89x
Humanoid (10 bodies/world), 64 worlds, total only 1.68 0.31 5.4x
Humanoid (10 sensing vs 13 counterpart bodies/world), 64 worlds 9.37 0.43 22x
Humanoid (10 bodies/world), 16384 worlds, total only 2127 40.1 53x
Humanoid (10 sensing vs 13 counterpart bodies/world), 16384 worlds 5467 97.1 56x

SensorContact.update

Per-call cost averaged over 1000 CUDA-graph iterations.

Benchmark before (μs) after (μs) speedup
Ant (4 bodies/world), 16384 worlds, total only 14.3 6.90 2.1x
Ant (4 sensing vs 9 counterpart bodies/world), 16384 worlds 17.2 9.92 1.7x
Humanoid (10 bodies/world), 64 worlds, total only 6.94 4.26 1.6x
Humanoid (10 sensing vs 13 counterpart bodies/world), 64 worlds 7.11 4.69 1.5x
Humanoid (10 bodies/world), 16384 worlds, total only 20.6 11.8 1.7x
Humanoid (10 sensing vs 13 counterpart bodies/world), 16384 worlds 26.0 18.3 1.4x

Summary by CodeRabbit

  • New Features

    • SensorContact now exposes total_force (per-sensor aggregate), force_matrix (per-counterpart breakdown), and metadata fields (sensing_obj_idx, counterpart_indices, sensing_obj_type, counterpart_type).
  • Deprecated

    • include_total renamed to measure_total; legacy parameter/fields (net_force, sensing_objs, counterparts/reading_indices, shape) deprecated with warnings and shim accessors.
  • Benchmarks

    • Added ASV benchmarks measuring SensorContact init/update performance (CUDA graph paths included).
  • Tests

    • Updated unit/integration tests to assert force_matrix/total_force and added new validation tests.
  • Bug Fixes / UX

    • Improved label-matching input validation and duplicate-index checks (raises ValueError); warmup added to benchmark timing.

@coderabbitai

coderabbitai Bot commented Mar 17, 2026

Copy link
Copy Markdown
Contributor
📝 Walkthrough

Walkthrough

Updates SensorContact: replaces net_force with total_force and force_matrix, renames include_totalmeasure_total, adds public sensing/counterpart index/type fields, rewrites accumulation kernels and multi-world index normalization, adds deprecated shims, and updates benchmarks, example, utils, and tests.

Changes

Cohort / File(s) Summary
Changelog
CHANGELOG.md
Documents API rename (include_totalmeasure_total), new outputs (total_force, force_matrix), new public metadata (sensing_obj_idx, counterpart_indices, sensing_obj_type, counterpart_type) and deprecations (net_force, sensing_objs, counterparts, reading_indices, shape).
SensorContact core
newton/_src/sensors/sensor_contact.py
Major refactor: new kernels (accumulate_contact_forces_kernel, expand_body_to_shape_kernel), world-aware index normalization helpers, strict sorted-unique index validation, device checks, new nullable outputs (total_force, force_matrix) and public index/type fields, constructor flag renamed to measure_total with deprecated include_total shim, and backward-compatible deprecated accessors.
Benchmarks
asv/benchmarks/simulation/bench_sensor_contact.py
Adds ASV benchmarks for SensorContact construction and update (with/without counterparts), model caching, warmup, and optional CUDA graph capture/execution.
Examples
newton/examples/sensors/example_sensor_contact.py
Switches example to explicit per-plate labels and per-counterpart force_matrix usage (measure_total=False), adds shape maps/counterpart column lookups, and updates plotting to use total_force where appropriate.
Tests
newton/tests/test_sensor_contact.py
Replaces MockModel with ModelBuilder fixtures, adapts assertions to force_matrix/total_force, adds tests for duplicate/unmatched indices and ordering, and verifies counterpart column mappings.
Utilities
newton/_src/utils/benchmark.py, newton/_src/utils/selection.py
Benchmark util: adds unmeasured warmup call before timing loops. Selection.match_labels: stricter type validation, fast path for list[int], early return for empty lists, and clearer TypeError messages.

Sequence Diagram(s)

sequenceDiagram
    participant User as User
    participant SensorContact as SensorContact
    participant Contacts as Contacts
    participant DeviceKernel as DeviceKernel

    User->>SensorContact: construct(measure_total, sensing selectors, counterpart selectors)
    User->>Contacts: provide contacts (device, pairs, forces)
    User->>SensorContact: update(Contacts)
    SensorContact->>Contacts: validate device matches
    SensorContact->>DeviceKernel: launch accumulate_contact_forces_kernel(contacts, mappings)
    DeviceKernel-->>SensorContact: write into force_matrix and total_force
    SensorContact-->>User: expose total_force / force_matrix (or None based on measure_total)
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~55 minutes

Possibly related PRs

Suggested labels

1.0-release

Suggested reviewers

  • adenzler-nvidia
  • nvlukasz
  • mmacklin
  • eric-heiden
🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 67.31% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (2 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title 'Optimize SensorContact and simplify API' directly addresses the two main objectives of this PR: optimizing the SensorContact implementation with performance improvements and simplifying the API by replacing net_force with total_force/force_matrix outputs and using flat metadata attributes instead of nested structures.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
📝 Coding Plan
  • Generate coding plan for human review comments

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@codecov

codecov Bot commented Mar 17, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 73.89163% with 53 lines in your changes missing coverage. Please review.
✅ All tests successful. No failed tests found.

Files with missing lines Patch % Lines
newton/_src/sensors/sensor_contact.py 74.60% 48 Missing ⚠️
newton/_src/utils/selection.py 75.00% 3 Missing ⚠️
newton/_src/utils/benchmark.py 0.00% 2 Missing ⚠️

📢 Thoughts on this report? Let us know!

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (1)
newton/_src/sensors/sensor_contact.py (1)

155-166: Consider simplifying the sorting logic.

The recursive call to sort is correct but could be simplified:

def _ensure_sorted_unique(indices: list[int], param_name: str) -> list[int]:
    indices = sorted(indices)  # always sort
    for i in range(1, len(indices)):
        if indices[i] == indices[i - 1]:
            raise ValueError(f"{param_name} contains duplicate index {indices[i]}")
    return indices

This avoids the recursive call while maintaining the same behavior. However, the current implementation is functionally correct.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@newton/_src/sensors/sensor_contact.py` around lines 155 - 166, The
_ensure_sorted_unique function currently uses a recursive sort when it
encounters an out-of-order pair; replace that with a simpler always-sort
approach: sort the indices up front (use sorted(indices)), then iterate once to
detect duplicates and raise ValueError with the same message using param_name
and the duplicate value, and finally return the sorted list; update the body of
_ensure_sorted_unique to remove recursion and always return the sorted,
duplicate-checked indices.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@CHANGELOG.md`:
- Line 30: Update the Changelog entry for deprecating SensorContact.shape and
SensorContact.reading_indices to include explicit migration guidance: state
which new field(s) replace each deprecated property (referencing
SensorContact.shape and SensorContact.reading_indices), provide a short example
or pseudo-code showing how to map old values to the new fields, and call out any
serialization or API compatibility steps users must take (e.g., updating stored
records or client code). Ensure the text is concise and actionable so users can
follow a clear migration path.

---

Nitpick comments:
In `@newton/_src/sensors/sensor_contact.py`:
- Around line 155-166: The _ensure_sorted_unique function currently uses a
recursive sort when it encounters an out-of-order pair; replace that with a
simpler always-sort approach: sort the indices up front (use sorted(indices)),
then iterate once to detect duplicates and raise ValueError with the same
message using param_name and the duplicate value, and finally return the sorted
list; update the body of _ensure_sorted_unique to remove recursion and always
return the sorted, duplicate-checked indices.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

Run ID: 2270ca68-c49f-43fa-abd4-5c41de4cc69f

📥 Commits

Reviewing files that changed from the base of the PR and between 7d8c3f0 and b605c02.

📒 Files selected for processing (7)
  • CHANGELOG.md
  • asv/benchmarks/simulation/bench_sensor_contact.py
  • newton/_src/sensors/sensor_contact.py
  • newton/_src/utils/benchmark.py
  • newton/_src/utils/selection.py
  • newton/examples/sensors/example_sensor_contact.py
  • newton/tests/test_sensor_contact.py

Comment thread CHANGELOG.md Outdated

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR refactors SensorContact to provide clearer force outputs (separating aggregate totals from per-counterpart breakdowns), updates downstream tests/examples to the new API, and adds benchmarking support to evaluate construction/update cost.

Changes:

  • Redesign SensorContact outputs: introduce total_force and optional force_matrix, plus new indexing metadata (sensing_obj_idx, counterpart_indices, type fields); deprecate legacy properties/params.
  • Update unit tests and the contact-sensor example to use the new API and ordering semantics.
  • Tighten match_labels() input validation, add a warmup call to the local benchmark runner, and add an ASV benchmark for SensorContact.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

Show a summary per file
File Description
newton/_src/sensors/sensor_contact.py Replaces legacy “net_force matrix” design with total_force + optional force_matrix, adds multi-world column mapping, and introduces deprecation shims.
newton/tests/test_sensor_contact.py Reworks tests to use real ModelBuilder models and validate the new SensorContact outputs and ordering behavior.
newton/examples/sensors/example_sensor_contact.py Updates the example to use total_force and per-counterpart force_matrix, and clarifies ordering assumptions.
newton/_src/utils/selection.py Adds stricter runtime validation to match_labels() for clearer error behavior.
newton/_src/utils/benchmark.py Adds an unmeasured warmup call before timing loops in run_benchmark().
asv/benchmarks/simulation/bench_sensor_contact.py Adds ASV benchmarks for SensorContact init and update (including CUDA graph capture path).
CHANGELOG.md Documents the SensorContact API change and deprecations.
Comments suppressed due to low confidence (1)

newton/_src/utils/benchmark.py:210

  • The warmup call runs method(*params) immediately before the timed loop. If method enqueues asynchronous Warp/CUDA work and only synchronizes during/after the timed loop, the first measured iteration may end up paying for outstanding warmup work. To keep timings isolated, consider synchronizing after the warmup (and/or before starting the timer) when Warp is initialized and the active device is CUDA.
                if attr.startswith("time_"):
                    # Warmup run (not measured).
                    method(*params)
                    # Run timing benchmarks multiple times and measure elapsed time.
                    for _ in range(number):
                        start = time.perf_counter()
                        method(*params)
                        t = time.perf_counter() - start

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.

Comment thread newton/_src/sensors/sensor_contact.py Outdated
@camevor camevor force-pushed the optimize-sensor-contact branch from b605c02 to 04c7fc2 Compare March 18, 2026 09:13
@camevor camevor enabled auto-merge March 18, 2026 09:13
@camevor camevor disabled auto-merge March 18, 2026 09:13

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@newton/_src/sensors/sensor_contact.py`:
- Around line 495-496: The metadata sets counterpart_type to "body"/"shape"
unconditionally based on counterpart_is_body even when counterpart_indices is
empty or force_matrix is None, making it inconsistent; update the assignment for
counterpart_type (where sensing_obj_type and counterpart_type are set) to
produce None when there are zero-column configs—i.e., only set "shape" or "body"
if counterpart_indices contains at least one non-empty entry (or force_matrix is
not None/has columns); otherwise set counterpart_type = None so it stays in sync
with counterpart_indices and force_matrix.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

Run ID: 814d7309-7fc3-4fc8-b59b-8b9abd05cfe0

📥 Commits

Reviewing files that changed from the base of the PR and between b605c02 and 04c7fc2.

📒 Files selected for processing (7)
  • CHANGELOG.md
  • asv/benchmarks/simulation/bench_sensor_contact.py
  • newton/_src/sensors/sensor_contact.py
  • newton/_src/utils/benchmark.py
  • newton/_src/utils/selection.py
  • newton/examples/sensors/example_sensor_contact.py
  • newton/tests/test_sensor_contact.py
🚧 Files skipped from review as they are similar to previous changes (3)
  • newton/_src/utils/benchmark.py
  • newton/examples/sensors/example_sensor_contact.py
  • newton/_src/utils/selection.py

Comment thread newton/_src/sensors/sensor_contact.py
@camevor camevor self-assigned this Mar 19, 2026

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@newton/_src/sensors/sensor_contact.py`:
- Around line 445-446: The current allocation of counterpart columns uses
max_readings computed before filtering to worlds that actually have sensing
rows, causing unused columns (see _assign_counterpart_columns,
counterpart_world_start, world_count, max_readings, counterpart_indices). Update
the logic so max_readings is derived only from counterparts whose worlds
intersect the sensing rows' worlds: first compute the set/list of worlds that
have sensing rows from counterpart_indices (or reject counterparts mapped to
worlds with no sensing rows), then call or modify _assign_counterpart_columns
(or its caller) to calculate max_readings and build
counterparts_by_world/_net_force based on that filtered set so columns are only
allocated for worlds that expose sensing rows.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

Run ID: 04e23f81-aa07-4db2-8674-88e7e4c95318

📥 Commits

Reviewing files that changed from the base of the PR and between 04c7fc2 and 4acc083.

📒 Files selected for processing (3)
  • CHANGELOG.md
  • newton/_src/sensors/sensor_contact.py
  • newton/tests/test_sensor_contact.py
🚧 Files skipped from review as they are similar to previous changes (1)
  • CHANGELOG.md

Comment thread newton/_src/sensors/sensor_contact.py
@camevor camevor requested a review from jvonmuralt March 19, 2026 10:41
Comment thread asv/benchmarks/simulation/bench_sensor_contact.py Outdated
jvonmuralt
jvonmuralt previously approved these changes Mar 19, 2026

@jvonmuralt jvonmuralt left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

looks good to me!

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
asv/benchmarks/simulation/bench_sensor_contact.py (1)

57-58: Use immutable tuples instead of mutable lists for ASV benchmark parameters to prevent accidental mutation and comply with RUF012.

Lines 57–58, 80–81, and 102–103 use mutable lists at class scope, which violates the RUF012 linting rule and risks accidental mutation across benchmark iterations. Refactor to use immutable tuples instead.

Proposed refactor
 class InitSensorContact:
@@
-    params = (["ant", "humanoid"], [64, 8192])
-    param_names = ["robot", "world_count"]
+    params = (("ant", "humanoid"), (64, 8192))
+    param_names = ("robot", "world_count")
@@
 class InitSensorContactWithMatching:
@@
-    params = (["ant", "humanoid"], [8192])
-    param_names = ["robot", "world_count"]
+    params = (("ant", "humanoid"), (8192,))
+    param_names = ("robot", "world_count")
@@
 class UpdateSensorContact:
@@
-    params = (["ant", "humanoid"], [64, 8192])
-    param_names = ["robot", "world_count"]
+    params = (("ant", "humanoid"), (64, 8192))
+    param_names = ("robot", "world_count")

No downstream code mutates these class-level attributes, making the refactoring safe.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@asv/benchmarks/simulation/bench_sensor_contact.py` around lines 57 - 58,
Class-level benchmark parameter attributes use mutable lists (params and
param_names) which violates RUF012 and risks accidental mutation; change them to
immutable tuples by replacing ["ant", "humanoid"] and [64, 8192] with ("ant",
"humanoid") and (64, 8192), and replace ["robot", "world_count"] with ("robot",
"world_count"); also audit and convert any other class-scope parameter lists in
this file (the other occurrences of params/param_names at the same class scope)
to tuples to ensure immutability.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@asv/benchmarks/simulation/bench_sensor_contact.py`:
- Around line 57-58: Class-level benchmark parameter attributes use mutable
lists (params and param_names) which violates RUF012 and risks accidental
mutation; change them to immutable tuples by replacing ["ant", "humanoid"] and
[64, 8192] with ("ant", "humanoid") and (64, 8192), and replace ["robot",
"world_count"] with ("robot", "world_count"); also audit and convert any other
class-scope parameter lists in this file (the other occurrences of
params/param_names at the same class scope) to tuples to ensure immutability.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

Run ID: b0a4dbf2-0fa5-46af-b970-0b5bd9ab8048

📥 Commits

Reviewing files that changed from the base of the PR and between 4acc083 and 3565f3f.

📒 Files selected for processing (1)
  • asv/benchmarks/simulation/bench_sensor_contact.py

@camevor camevor added this pull request to the merge queue Mar 19, 2026
Merged via the queue into newton-physics:main with commit 08c2c2f Mar 19, 2026
25 checks passed
mmacklin pushed a commit to mmacklin/newton that referenced this pull request Apr 7, 2026
@camevor camevor deleted the optimize-sensor-contact branch April 21, 2026 13:27
@coderabbitai coderabbitai Bot mentioned this pull request Apr 28, 2026
3 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants