Fix external contact cache lifetime under CUDA graphs

CHANGELOG.md

@@ -102,6 +102,7 @@
 ### Fixed
 
 - Make deterministic collision pipelines cover hydroelastic contact generation and reduction, including unique reduced-contact sort keys and overflow-safe fixed-point pressure accumulation.
+- Fix `SolverMuJoCo` retaining an invalid external-contact cache when its first step is captured in a CUDA graph. (#3767)
 - Convert `newton:mimicCoef0` from degrees to radians when the mimic follower joint is angular. Assets authored against the old behavior need the value rescaled to degrees.
 - Complete Kamino RCM traversal for large and disconnected systems and reuse the resulting permutation by default; set `reuse_permutation=False` to recompute it for changing matrix topology.
 - Bound Kamino DVI contact allocation with a per-world geometry heuristic instead of sizing every contact pair simultaneously; set `collision_detector.max_contacts_per_world` to override the inferred capacity.

newton/_src/solvers/mujoco/solver_mujoco.py

@@ -3692,13 +3692,9 @@ def __init__(
         # Initialised before _convert_to_mjc because notify_model_changed (called
         # during conversion) may call _invalidate_contact_fast_path.
         #
-        # Eagerly pre-allocate the device tracking buffers here (rather than
-        # lazily inside _convert_contacts_to_mjwarp).  Lazy wp.full(...) calls
-        # that happen on the first step often run while a CUDA graph is being
-        # captured; the resulting buffers can have a tangled lifetime and
-        # _invalidate_contact_fast_path() — which is called from outside the
-        # captured graph (e.g. notify_model_changed) — would then touch stale
-        # captured memory and trigger CUDA 700 (illegal memory access).
+        # Allocate the generation and count buffers before conversion because
+        # notify_model_changed() may invalidate them during conversion. The
+        # contact map is allocated below once the converted capacity is known.
         self._contact_tid_to_cid: wp.array[wp.int32] | None = None
         self._last_contact_generation = wp.full(1, _GENERATION_SENTINEL, dtype=wp.int32, device=self.device)
         self._last_nacon_count = wp.zeros(1, dtype=wp.int32, device=self.device)
@@ -3751,6 +3747,8 @@ def __init__(
                 include_sites=include_sites,
                 skip_visual_only_geoms=skip_visual_only_geoms,
             )
+        if not use_mujoco_cpu and not use_mujoco_contacts:
+            self._contact_tid_to_cid = wp.full(self.mjw_data.naconmax, -1, dtype=wp.int32, device=self.device)
         self._initial_model_sync = False
         self.update_data_interval = update_data_interval
         self._step = 0
@@ -4216,8 +4214,8 @@ def _convert_contacts_to_mjwarp(self, model: Model, state_in: State, contacts: C
         naconmax = self.mjw_data.naconmax
         launch_dim = min(contacts.rigid_contact_max, naconmax)
 
-        # Lazy-allocate the tid_to_cid buffer; reallocate if launch_dim grew
-        # (e.g. a different Contacts object with a larger rigid_contact_max).
+        # Grow the tid_to_cid buffer if the MJWarp data capacity changed after
+        # construction.
         # Invalidate the cached tid_to_cid mapping whenever any of the
         # invariants it depends on change:
         #

newton/tests/test_mujoco_solver.py

@@ -4477,14 +4477,23 @@ def test_fast_path_buffers_eagerly_allocated(self):
             solver._last_nacon_count,
             "_last_nacon_count must be eagerly pre-allocated in __init__",
         )
+        self.assertIsNotNone(
+            solver._contact_tid_to_cid,
+            "_contact_tid_to_cid must be eagerly pre-allocated in __init__",
+        )
         self.assertIsInstance(solver._last_contact_generation, wp.array)
         self.assertIsInstance(solver._last_nacon_count, wp.array)
+        self.assertIsInstance(solver._contact_tid_to_cid, wp.array)
         self.assertEqual(solver._last_contact_generation.dtype, wp.int32)
         self.assertEqual(solver._last_nacon_count.dtype, wp.int32)
+        self.assertEqual(solver._contact_tid_to_cid.dtype, wp.int32)
         self.assertEqual(solver._last_contact_generation.shape, (1,))
         self.assertEqual(solver._last_nacon_count.shape, (1,))
+        self.assertEqual(solver._contact_tid_to_cid.shape, (solver.mjw_data.naconmax,))
         self.assertEqual(solver._last_contact_generation.device, model.device)
         self.assertEqual(solver._last_nacon_count.device, model.device)
+        self.assertEqual(solver._contact_tid_to_cid.device, model.device)
+        self.assertTrue(np.all(solver._contact_tid_to_cid.numpy() == -1))
 
         # Calling _invalidate_contact_fast_path() before any step must succeed
         # cleanly — this is the exact path that previously hit stale captured