Fix external contact cache lifetime under CUDA graphs

[ooctipus]

Description

SolverMuJoCo already creates the contact fast-path generation and count buffers before the first step, but the thread-to-contact map was still created lazily in _convert_contacts_to_mjwarp(). When the first step was captured in a CUDA graph, the solver kept a pointer to that captured allocation after capture ended.

This allocates the map after model conversion, when mjw_data.naconmax is known. The allocation is limited to the GPU path with external Newton contacts. The existing capacity check remains in place in case the MJWarp data capacity changes later.

The fast-path buffer test now covers the map's type, size, device, and initial contents.

Closes #3767.

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 python -m unittest newton.tests.test_mujoco_solver.TestMuJoCoSolverNewtonContacts
uvx pre-commit run -a

The eight external-contact MuJoCo solver tests pass. I also ran the reproduction from #3767 directly:

• main: prints the cache pointer, then exits with code 139 on cache.numpy()
• this branch: reads all 200 entries successfully and exits with code 0

I verified the regression test separately with the source change removed; it fails because _contact_tid_to_cid is None.

Bug fix

Steps to reproduce:

1. Construct SolverMuJoCo with use_mujoco_contacts=False on CUDA.
2. Run its first contact step inside wp.ScopedCapture.
3. Read back solver._contact_tid_to_cid after capture.
4. The process segfaults while copying the cache to the host.

Minimal reproduction:

import newton
import warp as wp

from newton.solvers import SolverMuJoCo


device = wp.get_device("cuda:0")
builder = newton.ModelBuilder()
builder.add_ground_plane()
body = builder.add_body(xform=wp.transform(wp.vec3(0.0, 0.0, 0.18), wp.quat_identity()))
builder.add_shape_box(body, hx=0.1, hy=0.1, hz=0.1)
model = builder.finalize(device=device)

solver = SolverMuJoCo(model, use_mujoco_contacts=False, njmax=200, nconmax=200)
collision_pipeline = newton.CollisionPipeline(model)
contacts = collision_pipeline.contacts()
state_in, state_out, control = model.state(), model.state(), model.control()
newton.eval_fk(model, model.joint_q, model.joint_qd, state_in)

with wp.ScopedCapture(device=device):
    state_in.clear_forces()
    collision_pipeline.collide(state_in, contacts)
    solver.step(state_in, state_out, control, contacts, 0.002)

cache = solver._contact_tid_to_cid
print(f"cache ptr={cache.ptr}, shape={cache.shape}, device={cache.device}", flush=True)
print(cache.numpy(), flush=True)

Summary by CodeRabbit

• Bug Fixes
  • Fixed an issue where invalid external-contact data could persist when the first simulation step was captured in a CUDA graph.
  • Improved contact mapping allocation and cache invalidation to keep contact data accurate as simulation capacity changes.
  • Ensured contact buffers are initialized with the correct size, device, type, and default values.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro Plus

Run ID: d1f7dc65-e4d8-4016-a951-d2056b2ccae4

📥 Commits

Reviewing files that changed from the base of the PR and between ee84505 and 348deba.

📒 Files selected for processing (3)

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

---

📝 Walkthrough

Walkthrough

SolverMuJoCo now eagerly allocates its contact mapping with finalized MJWarp capacity, grows it when capacity changes, and tests its device, type, size, and initialization state.

Changes

MuJoCo contact cache lifecycle

Layer / File(s)	Summary
Eager contact mapping allocation newton/_src/solvers/mujoco/solver_mujoco.py, newton/tests/test_mujoco_solver.py	Contact tracking buffers initialize before model conversion. The contact mapping allocates after conversion with naconmax. Tests verify its type, device, shape, and -1 initialization.
Contact capacity change handling newton/_src/solvers/mujoco/solver_mujoco.py, CHANGELOG.md	Contact conversion grows the mapping when MJWarp capacity changes. The changelog records the CUDA-graph capture fix.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

• newton-physics/newton#2678: Earlier MuJoCo contact fast-path initialization and cache invalidation changes.
• newton-physics/newton#2393: Related MuJoCo contact-conversion fast-path buffer changes.
• newton-physics/newton#2208: Related MuJoCo contact conversion and mapping handling.

Suggested labels: release-blocking

Suggested reviewers: adenzler-nvidia

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly identifies the fix for external contact cache lifetime during CUDA graph capture.
Linked Issues check	✅ Passed	The changes allocate and validate _contact_tid_to_cid before the first step, addressing the CUDA graph capture regression in issue #3767.
Out of Scope Changes check	✅ Passed	The changelog, solver changes, and regression test directly support the linked issue and stated CUDA contact-cache objectives.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.}

[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!