Fix contact visualization in mujoco collision pipeline

[vidurv-nvidia]

Description

Fix update_contacts() not populating rigid_contact_point0 / rigid_contact_point1 when using SolverMuJoCo with use_mujoco_contacts=True.

The convert_mjw_contacts_to_newton_kernel declared both point arrays as outputs but never wrote to them. The kernel was missing mj_contact.pos (the world-frame contact position from MuJoCo) as an input, along with the body transforms (xpos, xquat) and geom-to-body mapping (geom_bodyid) needed to convert that world-frame position into Newton's body-local convention.

This caused:

• Contact visualization lines in the viewer to render at the world origin instead of at the actual contact surface
• Any downstream code reading contact point positions (sensors, custom logic) to see zeros
• The Newton-to-MuJoCo direction (convert_newton_contacts_to_mjw_kernel) was unaffected since it reads from Newton's collision pipeline which populates the points correctly

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)

Test plan

uv run --extra dev -m newton.tests --no-cache-clear -k test_contact_points_populated

Bug fix

Steps to reproduce:

1. Create a model with a ground plane and a dynamic box
2. Build SolverMuJoCo with use_mujoco_contacts=True
3. Simulate until the box hits the ground
4. Call solver.update_contacts(contacts, state)
5. Read contacts.rigid_contact_point0.numpy() and contacts.rigid_contact_point1.numpy()
6. Observe both arrays are all zeros despite nonzero rigid_contact_count

#2059

Summary by CodeRabbit

• Bug Fixes

  • Contact position data (rigid_contact_point0 and rigid_contact_point1) are now populated when MuJoCo contact simulation is enabled, providing accurate per-contact locations.

• Tests

  • Added a test that validates contact point population during MuJoCo contact scenarios.

• Documentation

  • Changelog updated to note the contact position population behavior.

[coderabbitai]

📝 Walkthrough

Walkthrough

The MuJoCo→Newton contact pipeline was extended so update_contacts now converts MuJoCo world-space contact positions into Newton body-local contact points. Kernel signature and solver wiring were updated to pass per-contact positions, geom→body mapping, and body poses; a test was added to validate populated rigid_contact_point0/rigid_contact_point1.

Changes

Cohort / File(s)	Summary
Changelog CHANGELOG.md	Added Unreleased entry noting update_contacts now populates rigid_contact_point0 and rigid_contact_point1 when use_mujoco_contacts is true.
Contact Conversion Kernel newton/_src/solvers/mujoco/kernels.py	Extended convert_mjw_contacts_to_newton_kernel signature to accept mj_contact_pos, mj_geom_bodyid, mj_xpos, mj_xquat; compute world-space contact points and transform them into Newton body-local frames, populating rigid_contact_point0/rigid_contact_point1.
Solver Wiring newton/_src/solvers/mujoco/solver_mujoco.py	Wired additional per-contact fields (pos, frame, dim, geom, efc_address, worldid) and body pose arrays (xpos, xquat) into the update_contacts kernel call and outputs.
Tests newton/tests/test_mujoco_solver.py	Added TestUpdateContactsPointPositions.test_contact_points_populated asserting contact point arrays are populated for a box-on-plane scenario; the diff contains a duplicated copy of this test class/method.

Sequence Diagram

sequenceDiagram
    participant MJ as MuJoCo Simulator
    participant SW as Solver Wiring
    participant K as Conversion Kernel
    participant NC as Newton Contacts

    MJ->>SW: Emit contacts (mj_contact.pos, geom, frame, worldid) and body poses (`xpos`,`xquat`)
    SW->>K: Pass per-contact pos, geom->body mapping, and body poses
    K->>K: Build body transforms X_wb from `xpos`/`xquat`
    K->>K: Compute world contact points and transform -> body-local points
    K->>NC: Write `rigid_contact_point0` and `rigid_contact_point1`

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~22 minutes

Possibly related PRs

• Refactor Newton<->MuJoCo data mapping #1198 #1199: Refactors MuJoCo↔Newton kernel signatures and mapping arrays used in contact conversion — overlaps kernel interface changes.
• Contact force sensors #429: Modifies MuJoCo→Newton contact conversion and solver update_contacts wiring — related to added per-contact fields.
• Fix collision pipeline mismatch in SolverMuJoCo examples #1927: Adjusts SolverMuJoCo contact update pipeline and sizing/calls — related to update_contacts usage.

Suggested labels

1.0-release

Suggested reviewers

• eric-heiden
• adenzler-nvidia
• vreutskyy

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The PR title claims to fix contact visualization, but the actual changes populate contact point fields that were previously zero; the visualization fix is a side effect of the core functionality.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

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

• X
• Mastodon
• Reddit
• LinkedIn

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ All tests successful. No failed tests found.

📢 Thoughts on this report? Let us know!

[coderabbitai]

🧹 Nitpick comments (1)

newton/tests/test_mujoco_solver.py (1)

8148-8160: Strengthen the plane-side assertion for point0.

Right now point0 only has a “not all zeros” check, so a bad transform that writes arbitrary non-zero values on the plane side would still pass. Since the plane is fixed at the origin, this test can cheaply assert that plane-side points stay on z ~= 0 and within the box footprint.

Suggested assertion tightening

                 self.assertFalse(
                     np.allclose(point0, 0.0),
                     "rigid_contact_point0 is all zeros — update_contacts did not populate contact positions",
                 )
+                np.testing.assert_allclose(
+                    point0[:, 2],
+                    0.0,
+                    atol=0.02,
+                    err_msg="plane-side contact points should lie on the ground plane",
+                )
+                self.assertTrue(np.all(np.abs(point0[:, 0]) <= 0.27))
+                self.assertTrue(np.all(np.abs(point0[:, 1]) <= 0.27))
                 self.assertFalse(
                     np.allclose(point1, 0.0),
                     "rigid_contact_point1 is all zeros — update_contacts did not populate contact positions",
                 )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@newton/tests/test_mujoco_solver.py` around lines 8148 - 8160, The plane-side
checks for contact points are weak: instead of only asserting point0 is not all
zeros, update the test around update_contacts to assert each plane-side contact
in point0 has z approximately 0 (use assertAlmostEqual with a small delta) and
that the x,y components lie within the expected box footprint (check |x| and |y|
≤ the box half-width/half-depth used in the test) to ensure the plane transform
and contact projection are correct; reference point0, point1, update_contacts
and n when adding these assertions.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@newton/tests/test_mujoco_solver.py`:
- Around line 8148-8160: The plane-side checks for contact points are weak:
instead of only asserting point0 is not all zeros, update the test around
update_contacts to assert each plane-side contact in point0 has z approximately
0 (use assertAlmostEqual with a small delta) and that the x,y components lie
within the expected box footprint (check |x| and |y| ≤ the box
half-width/half-depth used in the test) to ensure the plane transform and
contact projection are correct; reference point0, point1, update_contacts and n
when adding these assertions.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

Run ID: 2a3722c0-89b5-40a0-a7c8-cb048bb8b484

📥 Commits

Reviewing files that changed from the base of the PR and between aa59270 and ab73813.

📒 Files selected for processing (4)

• CHANGELOG.md
• newton/_src/solvers/mujoco/kernels.py
• newton/_src/solvers/mujoco/solver_mujoco.py
• newton/tests/test_mujoco_solver.py

[coderabbitai]

♻️ Duplicate comments (1)

newton/_src/solvers/mujoco/kernels.py (1)

683-745: ⚠️ Potential issue | 🔴 Critical

Add mj_contact_dist to the kernel inputs.

Line 739 reads mj_contact_dist[contact_idx], but that array is neither declared in this kernel signature nor passed by the current launch site in newton/_src/solvers/mujoco/solver_mujoco.py. As written, this path cannot compile/run, so rigid_contact_point0 / rigid_contact_point1 will never be populated by the new logic. Cross-file evidence: the launch snippet still passes mj_contact.pos directly to mj_contact.frame without a mj_contact.dist argument in between.

🐛 Minimal fix

 def convert_mjw_contacts_to_newton_kernel(
     # inputs
     mjc_geom_to_newton_shape: wp.array2d(dtype=wp.int32),
     mjc_body_to_newton: wp.array(dtype=wp.int32, ndim=2),
     pyramidal_cone: bool,
     mj_nacon: wp.array(dtype=wp.int32),
     mj_contact_pos: wp.array(dtype=wp.vec3),
+    mj_contact_dist: wp.array(dtype=float),
     mj_contact_frame: wp.array(dtype=wp.mat33f),
     mj_contact_dim: wp.array(dtype=int),
     mj_contact_geom: wp.array(dtype=wp.vec2i),
     mj_contact_efc_address: wp.array2d(dtype=int),
     mj_contact_worldid: wp.array(dtype=wp.int32),

# newton/_src/solvers/mujoco/solver_mujoco.py
inputs=[
    self.mjc_geom_to_newton_shape,
    self.mjc_body_to_newton,
    self.mjw_model.opt.cone == int(self._mujoco.mjtCone.mjCONE_PYRAMIDAL),
    mj_data.nacon,
    mj_contact.pos,
    mj_contact.dist,
    mj_contact.frame,
    mj_contact.dim,
    ...
]

#!/bin/bash
set -euo pipefail

echo "Inspect kernel signature and use site:"
sed -n '677,745p' newton/_src/solvers/mujoco/kernels.py

echo
echo "Inspect launch inputs:"
sed -n '3534,3561p' newton/_src/solvers/mujoco/solver_mujoco.py

Expected result: mj_contact_dist is used inside the kernel body, but missing from the signature and missing from the launch inputs until the fix is applied.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@newton/_src/solvers/mujoco/kernels.py` around lines 683 - 745, The kernel
reads mj_contact_dist but it is missing from the kernel signature and the launch
call; add mj_contact_dist as an input (e.g. mj_contact_dist:
wp.array(dtype=float)) to the kernel parameter list (near mj_contact_pos /
mj_contact_frame) and update the launch inputs in solver_mujoco.py to pass
mj_contact.dist (insert mj_contact.dist between mj_contact.pos and
mj_contact.frame as shown in the review's minimal fix) so rigid_contact_point0/1
can be computed.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Duplicate comments:
In `@newton/_src/solvers/mujoco/kernels.py`:
- Around line 683-745: The kernel reads mj_contact_dist but it is missing from
the kernel signature and the launch call; add mj_contact_dist as an input (e.g.
mj_contact_dist: wp.array(dtype=float)) to the kernel parameter list (near
mj_contact_pos / mj_contact_frame) and update the launch inputs in
solver_mujoco.py to pass mj_contact.dist (insert mj_contact.dist between
mj_contact.pos and mj_contact.frame as shown in the review's minimal fix) so
rigid_contact_point0/1 can be computed.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

Run ID: 1b95ce75-ebf7-469d-bf39-35588a89f127

📥 Commits

Reviewing files that changed from the base of the PR and between 7bf34d0 and c1646a9.

📒 Files selected for processing (1)

• newton/_src/solvers/mujoco/kernels.py