[BUG] Lazy contact cache does not survive CUDA graph capture

[ooctipus]

Bug Description

SolverMuJoCo leaves _contact_tid_to_cid unallocated until the first external-contact step. If that step runs during CUDA graph capture, the array does not remain valid after capture. It still has a pointer and the expected shape, but reading it back immediately segfaults.

I ran into this through a long Isaac Lab training job. One environment's base velocity, joint state, and height scan all became NaN at the same time while the policy actions were still finite. The failure was nondeterministic: the same job failed at very different iterations. Reducing it to the script below removes Isaac Lab and training entirely.

This looks like a remaining case from #2675 / #2678. That fix moved _last_contact_generation and _last_nacon_count out of the captured step, but _contact_tid_to_cid is still initialized to None and allocated in _convert_contacts_to_mjwarp().

I expected all persistent fast-path buffers owned by the solver to be allocated before the first step. Instead, the script prints the cache pointer and then exits with code 139 on cache.numpy().

Run with:

CUDA_LAUNCH_BLOCKING=1 uv run --extra dev python repro.py

Reproduction Script

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)

Output before the process exits:

cache ptr=84557169152, shape=(200,), device=cuda:0
Segmentation fault (exit code 139)

System Information

• Newton: 1.5.0.dev0, main at ee84505ae44f1d2c04a38ac513bf34ab93aa81aa
• Warp: 1.16.0.dev20260716
• MuJoCo Warp: 3.11.0
• Python: 3.12.12
• OS: Ubuntu 24.04.3 LTS, kernel 7.0.0-28-generic
• GPU: NVIDIA GeForce RTX 5090
• Driver: 580.173.02