@mar-yan24 thank you for contributing this feature to mujoco warp!

@mar-yan24 fyi there is a warp draft pr for introducing determinism in warp NVIDIA/warp#1355

We just discussed this - it could be worth pursuing this approach in parallel to Warp's low-level support for determinism as they two different approaches may have different performance tradeoffs.

@thowell, thanks for the input! Actually I haven't kept up with Warp as closely recently so I'll take a look at the PR brought up there and see if there are similar ideas compared to what I have in my current plan.

Regarding @erikfrey's comment, I don't mind working on the rest of the determinism implementation for this PR and comparing the performance once finished. I'll probably continue working on this for the week and I'll try to finish by around a week from now for the full end-to-end implementation.

Thank you both for the info/updates!

Thanks for the review @thowell! The changes should be good to go. I am planning the next determinism steps after this PR like constraint row allocation and actuator moment allocation. Before I continue building, would you prefer that I keep extending this branch/PR so the work is all on this PR or split it up into separate requests for review. Either works for me.

@mar-yan24 lets create separate prs for the next deterministic features. thanks!

@mar-yan24 can we benchmark the performance of this pr with the built-in determinism from warp NVIDIA/warp#1355? probably makes sense to confirm that writing custom deterministic kernels is more performant compared to the general purpose warp determinism functionality. thanks!

@thowell I just tried running the NVIDIA/warp#1355 build on my machine and it seems there are several issues with it that currently make it incompatible with mujoco_warp. When I first tried running, it just crashed, so I disabled graph capture. I think the Warp PR crashes inside wp.ScopedCapture when mujoco_warp launches kernels, discussed below.

The PR's codegen looks up the destination array for each atomic_add by matching the kernel's top-level argument names (context.py:7706). That works for wp.atomic_add(arr, idx, value) but not for the sliced pattern wp.atomic_add(arr[outer_idx], inner_idx, value) that we use in other kernels.

I don't mind helping try and fix this/look into this deeper, but I suspect there may be other blockers as well. In the meantime, I can draft a minimal repro kernel for the PR? Let me know your thoughts.

hello guys,
how is it going this?

Hey @johnnynunez I took a bit of a break on some of the projects I was working on as I started a new job recently. Alongside this, I am also waiting in part to take a look at the final implementation of determinism on the Warp side before doing too much more as I want to wait and see what the reach of the implementation there will be.

If you would like to contribute and help out in what I have here, be my guest. But I might take another week or two before I have enough time to pick things back up :)

Hey @johnnynunez I took a bit of a break on some of the projects I was working on as I started a new job recently. Alongside this, I am also waiting in part to take a look at the final implementation of determinism on the Warp side before doing too much more as I want to wait and see what the reach of the implementation there will be.

If you would like to contribute and help out in what I have here, be my guest. But I might take another week or two before I have enough time to pick things back up :)

thank you... let me push from Nvidia side (warp)

Following up here after working through the warp side (NVIDIA/warp#1355, fix PR mmacklin/warp#3) — continuing this work as @mar-yan24 invited contributors.

Rebased branch available: I rebased this PR onto current main (96 commits forward, conflicts in io.py override fields and types.py ls_parallel deprecation): https://github.com/johnnynunez/mujoco_warp/tree/det/mjwarp-determinism1 — full test suite passes (1069 passed, 23 skipped) on RTX PRO 6000 Blackwell (sm_120, CUDA 13.3). @mar-yan24 happy to open a PR into your branch or for the team to take it directly, whichever is easier.

Benchmark: custom deterministic kernels vs warp automatic determinism (@thowell asked about this). 20-body contact-rich pile, CUDA graph captured, 1000 steps:

The custom-kernel approach in this PR costs ~9-12%; warp's automatic interception costs 3x at 1 world and collapses at high world counts (the per-launch sort-reduce dominates, and scatter buffers hit the int32 allocation limit at 1024 worlds × njmax=2048). This matches @eric-heiden's observation on the warp PR that automatic mode is too slow for MuJoCo and custom kernels are the right path — i.e., this PR's approach is the right architecture.

Determinism measurements on the rebased branch (same-process, bit-identical full-Data resets, 200 steps): contact/constraint assembly is bitwise stable, and single-step replay is fully bitwise. Long rollouts still split at ~step 115 (also ~108 on main with and without any determinism flags; first divergent arrays are solver-internal efc.Ma/qacc with bit-identical inputs and equal solver_niter). So the residual nondeterminism lives in non-atomic solver internals (likely the blocked Cholesky / mathdx path) — consistent with this PR's stated scope of fixing contact ordering first, with solver reductions as follow-up (#1300 and beyond).

Root cause of the residual long-rollout drift found — correcting my earlier guess (it is not mathdx/blocked-Cholesky; it's two racing float atomics in the sparse solver path).

Method: bit-identical full-Data resets in one process, replaying 200-500 steps, replacing solver sub-phases one at a time with order-deterministic equivalents and checking the first divergent step (~108 on main for a 20-body pile).

Eliminations (each alone does NOT fix it): serial sparse LDL factorization, serial sparse LDL backsubstitution, forced full H rebuild, tolerance=0, CG vs Newton (CG drifts identically → Cholesky/mathdx exonerated), warp deterministic_max_records, zeroed workspaces.

The minimal fixing pair (drift gone over 500 steps, 3/3 runs):

1. _update_constraint_init_qfrc_constraint_sparse — wp.atomic_add(qfrc_constraint[worldid], colind, J*force) races across efc rows hitting the same dof; replaced with a per-(world,dof) gather in efc index order.
2. _JTDAJ_sparse (+ _update_gradient_h_incremental_sparse) — wp.atomic_add(h[worldid,row], col, ...) races across efc rows touching the same (row,col); replaced with a per-upper-tri-element gather.

A third, smaller source remains in linesearch-adjacent atomics (Jaref multi-thread init / search_dot): with only the pair above fixed, onset moves from ~108 → ~197; adding warp's RUN_TO_RUN interception for the rest gives bitwise determinism over 500 steps. Consistent with this: dense-jacobian mode is already fully deterministic (its qfrc_constraint and JTDAJ are gathers/tiles), only sparse drifts.

This cleanly defines determinism 3/N scope: deterministic sparse qfrc_constraint + sparse H assembly (+ optionally the linesearch Jaref/search_dot atomics). The gather versions I used for falsification are naive (O(nefc) per dof/element) — a production version would want the usual count/scan or segmented-reduce treatment, same architecture as 2/N. Happy to draft it.

Follow-up to the root-cause analysis above: the fix is now implemented as determinism 3/N — mar-yan24#8. End result across the stack (1/N rebase + 2/N port + 3/N): bitwise qacc/qpos/qvel over 1000 steps, graph-capture-safe, 1108 tests passing. Details and honest perf table in that PR.

@mar-yan24 could you update your fork and all branch with main? then my PRs will be clean

Wow @johnnynunez! Your development speed is quite fast; I cannot keep up. Are you still looking for me to update my branches? They are somewhat deprecated, so I need to take a look at the new conflicts. I see your new PRs though, thanks for the contribution! This weekend I will look into your new PRs. In the meantime, I can rebase my code.