Differentiable stepping: reverse-mode autodiff through smooth dynamics, the constraint solver, and contacts

[johnnynunez]

This PR adds end-to-end reverse-mode differentiation through MuJoCo Warp stepping, enabling gradient-based policy optimization and system identification directly on MJWarp physics.

It lands in three pieces. First, reverse-mode autodiff for smooth dynamics, including the free-joint cdof_dot fix and restored backward paths required by the control-gradient chain. Second, implicit differentiation through the converged Newton fixed point rather than unrolling solver iterations. The fixed-point backward includes the direct smooth-force cotangent, packed mass-matrix dependence, active constraint reference/impedance, full constraint-Jacobian geometry, and velocity dissipation. Third, smooth contact differentiation through contact geometry and constraint assembly, with equivalent dense and packed-sparse paths and isolated per-substep AD state.

The current branch cross-validates dL/dctrl, dL/dqpos, and dL/dqvel against finite differences. The affected suite passes with 87 grad tests, and the full repository suite passes (1288 passed, 18 skipped). Hopper one-step qpos/qvel errors are 2.84e-5 / 8.95e-6 in dense mode and 3.10e-5 / 3.00e-5 in sparse mode. A 32-substep PyTorch↔Warp Hopper rollout reaches 8.68e-5 relative action-gradient error with cosine 0.9999999979. Active slide limits reach 1.29e-6 in both layouts. Pre-commit, Ruff, kernel-analyzer, and git diff --check pass.

The verified contact scope now covers both cones: elliptic contact is implemented with the coupled per-contact cone block and FD-tested at condim 1/3/4/6, including solreffriction and impratio. Ball-joint limits, tendon limits, and the equality families (joint, connect, weld, tendon, flex) carry their reference, impedance, and Jacobian state-dependence through the fixed-point backward with FD tests, and friction rows get the dissipation adjoint in the quadratic near-sticking regime. Plane-box, sphere-box, and capsule-box contacts are replayed differentiably; pairs without a differentiable replay (mesh, hfield, SDF, convex-convex) are reported by unsupported_geom_pairs() and warned about from enable_grad. Exact active-set switching points are nonsmooth, so gradients are defined and tested within fixed active regimes, and stepping under an active tape raises during CUDA graph capture rather than silently returning single-step gradients.

Authored with Mark Yang (@mar-yan24); developed and reviewed incrementally in mar-yan24#7, #5, #10, and #11 before being combined here. Credit to @mar-yan24.

[johnnynunez]

Test and benchmark results on this branch as pushed (RTX PRO 6000, CUDA 13.3).

Full test suite: uv run pytest -n 4 -> 1093 passed, 23 skipped. This includes the 29 gradient tests in grad_test.py, which cross-validate reverse-mode gradients (dL/dctrl, dL/dqpos, dL/dqvel) against finite differences through multi-step rollouts with contacts, agreeing to ~1e-4 at fp32 solver tolerances.

Forward-step throughput, mjwarp-testspeed benchmarks/humanoid/humanoid.xml --nstep 500 (steps/second, this branch vs current main):

nworld	main	this PR	delta
1	4230	3851	-9.0%
256	621692	599381	-3.6%
1024	2231204	2160472	-3.2%
4096	6178023	5903726	-4.4%

The forward-path overhead comes from re-enabling backward passes on kernels that previously had enable_backward=False and from the differentiable contact path. If this is a concern we can gate the differentiable path behind an option so the default forward-only step keeps current performance; happy to make that change if reviewers prefer it.

[johnnynunez]

Update: the forward-path overhead reported above is now eliminated. The default forward-only step matches main; AD costs are only paid when gradients are requested.

We profiled the regression with Nsight Systems (--cuda-graph-trace=node, since the kernels are hidden inside captured CUDA graphs otherwise) and found three sources, now addressed in the latest commits:

1. Backward-enabled kernel compilation. Modules on the gradient path compiled with enable_backward unconditionally, which generates slower forward code. A new ad_flags module gates this behind mjw.enable_ad(), called automatically by make_diff_data / enable_grad (or preset via MJWARP_ENABLE_AD), so the default build is forward-only.

2. AD-only workarounds running unconditionally. _rne_cfrc_backward had been rewritten to allocate per-tree-level scratch arrays each step with an O(nbody) gather kernel (needed to avoid Warp's tape output-gradient zeroing, but ~12us/step of graph gaps), and several wp.copy memcpys had become _nograd_copy kernel launches (~2.5us/step). Both now branch on requires_grad: forward-only data takes the exact upstream path, gradient-tracked data keeps the AD-safe variants.

3. Solver retained state (solver_h, solver_hfactor, solver_Jaref, kept past the solve for the implicit-diff backward pass) was allocated for every Data object. It is now zero-sized by default and allocated lazily on the first gradient-tracked solve, so forward-only stepping carries no extra memory.

Re-measured with interleaved main/branch runs on the same setup as above (humanoid, 500 steps, steps/second):

nworld	main	this PR (before)	this PR (now)
1	4100-4360	3851 (-9.0%)	4302 (parity)
1024	~2240000	2160472 (-3.2%)	~2222000 (-0.7%)
4096	~6178000	5903726 (-4.4%)	~6175000 (parity)

Nsight confirms GPU busy time per step is main 115.6ms vs branch 116.0ms over 300 graph-captured steps (+0.3%). The remaining ~0.7% at 1024 worlds comes from kernels that were restructured for AD correctness and run in both modes (e.g. _kinematics_branch unrolls joint processing because continue inside dynamic loops breaks Warp's adjoint replay, ~0.5us/step). Splitting those into forward-only and AD variants would close it entirely at the cost of duplicated kernels; we left them shared, but happy to split them if reviewers prefer strict parity.

Full test suite still passes after the changes (1093 passed, 23 skipped), including all 29 gradient cross-validation tests.

[johnnynunez]

viz @thowell

[thowell]

@johnnynunez @mar-yan24

thanks for contributing this feature to mujoco warp!

given the complexity of this implementation and the potential maintenance of this feature it would be helpful to better understand some of its key performance characteristics. how does this auto-differentiation approach compare to a finite difference approach in terms of throughput, jit time, and memory for our benchmark scenes? thanks!

[johnnynunez]

Thanks @thowell, that's the right thing to probe given the maintenance surface. I measured the three things you asked about (throughput, JIT time, memory) and, because your framing is really about whether the feature earns its keep, two more that the first three don't answer on their own: what differentiability costs the forward-only path that most users stay on, and whether the gradients are actually better than finite differences. All numbers are on an RTX PRO 6000 Blackwell, torch 2.12.0+cu130, this branch at 9491f45 versus its merge-base 083e5ef, single precision. Finite differences are batched into worlds (world 2k perturbs parameter k by +eps, world 2k+1 by -eps) so FD gets full GPU parallelism rather than a Python loop. Repro scripts are at the end.

First the maintenance question directly: what does merging this cost someone who never calls backward? I compared a forward step() on main against the same step() on this branch with gradient tracking off, graph-captured and replayed at nworld=8192.

scene     arm                 ns/step/world   device mem (MiB)
humanoid  main (083e5ef)              105.5            398.8
humanoid  branch, grad OFF           107.4            398.0
humanoid  branch, grad ON            112.5            526.0
franka    main (083e5ef)               47.5            270.0
franka    branch, grad OFF            48.0            270.0
franka    branch, grad ON             54.8            334.0

Grad-off is within 1-2% of main and adds zero device memory, which matches the intent of the gating commits: the forward-only path compiles the original kernels and allocates nothing extra. Turning grad on costs ~5% (humanoid) to ~14% (franka) on the step and a modest amount of memory for the gradient arrays (+128 and +64 MiB here). So the feature is opt-in at runtime, not a tax on everyone.

Now throughput, JIT, and memory for computing one gradient, cold compile in a fresh process each time:

scene     mode  nparam  FD worlds  JIT (s)  grads/s   mem (MiB)
franka    ad         8          1     37.0    184.7        117
franka    fd         8         16     11.7    466.3        110
humanoid  ad        21          1     30.3    208.5        115
humanoid  fd        21         42     13.0    523.1        110
cloth     fd      2706       5412     11.2      9.9       5988
cloth     ad      2706          -        -   (faults)        -

Two honest takeaways. AD compiles ~2-3x slower because it builds the backward kernels, and for a single trajectory on small models batched FD is actually faster per gradient, because AD runs at nworld=1 and leaves most of a 96 GB GPU idle while FD fills 2nparam worlds. AD's structural advantage is memory: FD's footprint grows linearly in the number of parameters (cloth needs 5412 worlds and ~6 GB for one gradient of a 2706-dim input) while AD's is set by tape depth, independent of nparam. The same trend holds across a BPTT horizon sweep on humanoid (gradient of a T-step rollout w.r.t. the whole control trajectory): batched FD keeps a ~2.4x throughput edge at every horizon because the card swallows the worlds, until 2nparam stops fitting.

The reason to carry AD anyway is accuracy, which raw throughput hides. I took a float64 reference gradient from MuJoCo's CPU C engine (mj_step, central differences at a well-conditioned eps, Richardson self-check below 1e-8 so the reference is converged) and compared warp AD against warp FD swept over eps, on the smooth Euler path where both are valid:

humanoid (nu=21), relerr vs float64 reference
  AD (no eps):        3.9e-5
  FD eps=3e-2:        8.4e-5   <- best FD
  FD eps=1e-3:        3.1e-3
  FD eps=1e-4:        2.0e-2
  FD eps=1e-6:        1.6      (float32 cancellation)

franka (nu=8), relerr vs float64 reference
  AD (no eps):        7.7e-7
  FD eps=1e-3:        1.6e-5   <- best FD
  FD eps=1e-1:        4.3e-1   (truncation)
  FD eps=1e-7:        1.9e-1   (cancellation)

FD has no single good eps: its error is a U-curve, truncation O(eps^2) on the high side and float32 cancellation blowing up on the low side, and the usable window is narrow and problem-dependent (3e-2 for humanoid, 1e-3 for franka, four to five orders of magnitude worse just outside it). AD has no eps to tune and lands at or below the best attainable FD point, about 20x more accurate than the best-tuned FD on franka. That is the maintenance justification: exact gradients with no per-problem tuning, and a cost that doesn't scale with parameter count, which is what makes long-horizon high-dimensional BPTT feasible at all.

Two scope boundaries I want to flag honestly rather than hide. franka ships with integrator=implicitfast, and there AD vs FD is only cosine 0.96 (eps-independent across 1e-2..1e-6); forcing the same model to Euler restores relerr 2e-4, so the current backward differentiates the explicit update but not the implicit velocity term (qderiv). And cloth uses the sparse mass-matrix path whose LD-solve kernels are compiled enable_backward=False, so taping backward through it currently faults; FD is the only option there today. Both are reasonable follow-ups, but worth documenting since manipulation models default to implicitfast.

Net: correct and exact on the supported path (Newton, dense, explicit/Euler), opt-in with ~1-2% forward overhead and zero memory when off, a 2-3x JIT premium and memory-not-throughput win when on, and the two gaps above tracked. Scripts below so you can rerun on your own benchmark hardware; happy to fold them into a benchmark target if useful.

scripts/bench_overhead.py (forward step() overhead: main vs branch grad-off vs grad-on)

"""Overhead probe: does merging differentiability cost anything when you DON'T
use it, and what does turning it on cost?

Measures forward step() throughput + device memory for a real batched rollout,
in three arms (run this same script from each worktree / data mode):

  --mode plain     mjw.make_data(...)             forward-only Data
  --mode diff_off  mjw.make_data(...)             on the diff branch, grad OFF
  --mode diff_on   mjw.make_diff_data(...)        grad tracking ON (no backward)

Run arm A from the main worktree (083e5ef) with --mode plain, arms B/C from the
differentiability worktree. plain (main) vs diff_off (branch) = cost this PR adds
for users who never call backward (the maintenance question). diff_off vs diff_on
= cost of enabling grad on the forward pass.

Steady-state: graph-capture the step and replay, like testspeed does.

  PYTHONPATH=/home/johnny/Projects/dev/warp uv run python scripts/bench_overhead.py \
      --scene humanoid --nworld 8192 --mode plain
"""
import argparse
import json
import time

import mujoco
import numpy as np
import warp as wp

import mujoco_warp as mjw
from etils import epath

SCENES = {
  "humanoid": "benchmarks/humanoid/humanoid.xml",
  "franka": "benchmarks/franka_emika_panda/scene.xml",
}


def main():
  ap = argparse.ArgumentParser()
  ap.add_argument("--scene", default="humanoid", choices=list(SCENES))
  ap.add_argument("--nworld", type=int, default=8192)
  ap.add_argument("--mode", required=True, choices=["plain", "diff_off", "diff_on"])
  ap.add_argument("--nstep", type=int, default=500)
  args = ap.parse_args()

  wp.config.quiet = True
  wp.init()
  wp.clear_kernel_cache()
  dev = wp.get_device()
  cuda = dev.is_cuda
  mjm = mujoco.MjSpec.from_file(epath.Path(SCENES[args.scene]).as_posix()).compile()

  if cuda:
    wp.synchronize_device()
  mem0 = wp.get_device(dev).free_memory if cuda else 0

  m = mjw.put_model(mjm)
  if args.mode == "diff_on":
    d = mjw.make_diff_data(mjm, nworld=args.nworld)
  else:
    d = mjw.make_data(mjm, nworld=args.nworld)

  # warmup + compile
  mjw.step(m, d)
  if cuda:
    wp.synchronize_device()
  mem_after = wp.get_device(dev).free_memory if cuda else 0
  mem_used = (mem0 - mem_after) / 1024**2 if cuda else 0.0

  # steady-state timing via graph capture (falls back to plain loop on CPU)
  reps = args.nstep
  if cuda:
    with wp.ScopedCapture() as cap:
      mjw.step(m, d)
    wp.synchronize_device()
    t0 = time.perf_counter()
    for _ in range(reps):
      wp.capture_launch(cap.graph)
    wp.synchronize_device()
    dt = time.perf_counter() - t0
  else:
    t0 = time.perf_counter()
    for _ in range(reps):
      mjw.step(m, d)
    dt = time.perf_counter() - t0

  steps = reps * args.nworld
  out = dict(scene=args.scene, mode=args.mode, nworld=args.nworld,
             device=("cuda" if cuda else "cpu"),
             ns_per_step_per_world=1e9 * dt / steps,
             steps_per_s=steps / dt,
             mem_mib=mem_used)
  print("RESULT_JSON " + json.dumps(out))


if __name__ == "__main__":
  main()

scripts/bench_ad_vs_fd.py (single-gradient AD vs FD: throughput, JIT, memory)

"""AD vs finite-difference benchmark for the mjwarp differentiability PR (#1423).

Answers thowell's review question: throughput, JIT/compile time, and device
memory for reverse-mode AD versus central finite differences, on benchmark
scenes.

Gradient target: dL/d(input) through a SINGLE differentiable step(), where
  input = ctrl   if the scene has actuators (nu>0)
        = qvel   otherwise (e.g. cloth: nparam = nv)
  L = sum_i wpos * qpos_i^2  +  wvel * qvel_i^2   (after one step)

Single step is deliberate: it is the regime where mjwarp's AD is validated
FD-exact (see grad_test.py::_assert_step_ctrl_grad, assert_allclose at
_FD_TOL=1e-3). Multi-step AD is approximate by a documented Warp tape
limitation (grad_test.py:999, "shared-array accumulation across steps"), so a
multi-step AD-vs-FD *accuracy* table would be misleading; we keep accuracy on
the validated single-step unit and let the cost metrics speak to scaling.

The AD-vs-FD tradeoff is exposed by nparam (= nu or nv) across scenes:
FD needs 2*nparam forward evals (batched into worlds here, the strongest form);
AD needs one forward + one backward, independent of nparam.

One MODE per process so warp's per-process kernel cache yields a valid COLD jit.
Emits one RESULT_JSON line. The driver checks ||g_ad - g_fd|| from saved vecs.

Run (from mujoco_warp/, warp dev tree on PYTHONPATH):
  PYTHONPATH=/home/johnny/Projects/dev/warp uv run python scripts/bench_ad_vs_fd.py \
      --scene humanoid --mode ad
"""

import argparse
import json
import time

import mujoco
import numpy as np
import warp as wp

import mujoco_warp as mjw
from etils import epath

SCENES = {
  "franka": "benchmarks/franka_emika_panda/scene.xml",
  "humanoid": "benchmarks/humanoid/humanoid.xml",
  "cloth": "benchmarks/cloth/scene.xml",
}


@wp.kernel
def _accum_loss_q(qpos: wp.array2d(dtype=float), wpos: float, loss: wp.array(dtype=float)):
  w, i = wp.tid()
  wp.atomic_add(loss, 0, wpos * qpos[w, i] * qpos[w, i])


@wp.kernel
def _accum_loss_v(qvel: wp.array2d(dtype=float), wvel: float, loss: wp.array(dtype=float)):
  w, i = wp.tid()
  wp.atomic_add(loss, 0, wvel * qvel[w, i] * qvel[w, i])


@wp.kernel
def _pw_loss_q(qpos: wp.array2d(dtype=float), wpos: float, lossw: wp.array(dtype=float)):
  w, i = wp.tid()
  wp.atomic_add(lossw, w, wpos * qpos[w, i] * qpos[w, i])


@wp.kernel
def _pw_loss_v(qvel: wp.array2d(dtype=float), wvel: float, lossw: wp.array(dtype=float)):
  w, i = wp.tid()
  wp.atomic_add(lossw, w, wvel * qvel[w, i] * qvel[w, i])


def _load(scene):
  spec = mujoco.MjSpec.from_file(epath.Path(SCENES[scene]).as_posix())
  return spec.compile()


def _free_mem(dev):
  return wp.get_device(dev).free_memory


def main():
  ap = argparse.ArgumentParser()
  ap.add_argument("--scene", required=True, choices=list(SCENES))
  ap.add_argument("--mode", required=True, choices=["fwd", "ad", "fd"])
  ap.add_argument("--reps", type=int, default=20)
  ap.add_argument("--eps", type=float, default=1e-3)
  ap.add_argument("--wpos", type=float, default=1.0)
  ap.add_argument("--wvel", type=float, default=0.1)
  ap.add_argument("--save_grad", default="")
  args = ap.parse_args()

  wp.config.quiet = True
  wp.init()
  wp.clear_kernel_cache()
  dev = wp.get_device()
  is_cuda = dev.is_cuda

  mjm = _load(args.scene)
  nq, nv, nu = mjm.nq, mjm.nv, mjm.nu
  use_ctrl = nu > 0
  nparam = nu if use_ctrl else nv

  rng = np.random.default_rng(0)
  if use_ctrl:
    inp_base = (0.01 * rng.standard_normal(nu)).astype(np.float32)
  else:
    inp_base = (0.01 * rng.standard_normal(nv)).astype(np.float32)

  if is_cuda:
    wp.synchronize_device()
  mem0 = _free_mem(dev) if is_cuda else 0
  m = mjw.put_model(mjm)

  def base_state(nworld):
    qp = np.tile(mjm.qpos0.astype(np.float32), (nworld, 1))
    return qp

  def reset(d, nworld, qvel_np=None):
    d.qpos.assign(base_state(nworld))
    if qvel_np is None:
      d.qvel.zero_()
    else:
      d.qvel.assign(qvel_np)
    d.qacc.zero_()
    if hasattr(d, "qacc_warmstart"):
      d.qacc_warmstart.zero_()
    d.act.zero_()
    d.time.zero_()

  def add_loss(loss, d, nworld):
    wp.launch(_accum_loss_q, dim=(nworld, nq), inputs=[d.qpos, args.wpos], outputs=[loss])
    wp.launch(_accum_loss_v, dim=(nworld, nv), inputs=[d.qvel, args.wvel], outputs=[loss])

  # ----------------------------------------------------------------- AD --
  if args.mode == "ad":
    d = mjw.make_diff_data(mjm, nworld=1)
    if use_ctrl:
      inp = wp.zeros((1, nu), dtype=float, requires_grad=True)
      inp.assign(inp_base.reshape(1, nu))
    else:
      inp = wp.zeros((1, nv), dtype=float, requires_grad=True)
      inp.assign(inp_base.reshape(1, nv))

    def grad_once():
      reset(d, 1)
      if use_ctrl:
        wp.copy(d.ctrl, inp)
      else:
        wp.copy(d.qvel, inp)
      loss = wp.zeros(1, dtype=float, requires_grad=True)
      tape = wp.Tape()
      with tape:
        if use_ctrl:
          wp.copy(d.ctrl, inp)
        mjw.step(m, d)
        add_loss(loss, d, 1)
      tape.backward(loss=loss)
      g = inp.grad.numpy().copy().reshape(-1)
      L = float(loss.numpy()[0])
      tape.zero()
      return L, g

    t0 = time.perf_counter()
    L, g = grad_once()
    if is_cuda:
      wp.synchronize_device()
    jit = time.perf_counter() - t0
    mem = mem0 - _free_mem(dev) if is_cuda else 0
    ts = []
    for _ in range(args.reps):
      if is_cuda:
        wp.synchronize_device()
      t = time.perf_counter()
      grad_once()
      if is_cuda:
        wp.synchronize_device()
      ts.append(time.perf_counter() - t)
    tg = float(np.median(ts))
    if args.save_grad:
      np.save(args.save_grad, g)
    out = dict(scene=args.scene, mode="ad", device=("cuda" if is_cuda else "cpu"),
               input=("ctrl" if use_ctrl else "qvel"), nu=nu, nv=nv, nq=nq,
               nparam=nparam, nworld=1, jit_s=jit, grad_s=tg, grads_per_s=1.0 / tg,
               mem_mib=mem / 1024**2, loss=L)

  # ----------------------------------------------------------------- FD --
  elif args.mode == "fd":
    nworld = 2 * nparam
    d = mjw.make_data(mjm, nworld=nworld)
    pert = np.tile(inp_base.reshape(1, -1), (nworld, 1))
    for k in range(nparam):
      pert[2 * k, k] += args.eps
      pert[2 * k + 1, k] -= args.eps
    pert_wp = wp.array(pert.astype(np.float32), dtype=float)

    def grad_once():
      if use_ctrl:
        reset(d, nworld)
        wp.copy(d.ctrl, pert_wp)
      else:
        reset(d, nworld, qvel_np=pert.astype(np.float32))
      mjw.step(m, d)
      lossw = wp.zeros(nworld, dtype=float)
      wp.launch(_pw_loss_q, dim=(nworld, nq), inputs=[d.qpos, args.wpos], outputs=[lossw])
      wp.launch(_pw_loss_v, dim=(nworld, nv), inputs=[d.qvel, args.wvel], outputs=[lossw])
      Lw = lossw.numpy().astype(np.float64)
      g = np.empty(nparam, dtype=np.float64)
      for k in range(nparam):
        g[k] = (Lw[2 * k] - Lw[2 * k + 1]) / (2.0 * args.eps)
      return g

    t0 = time.perf_counter()
    g = grad_once()
    if is_cuda:
      wp.synchronize_device()
    jit = time.perf_counter() - t0
    mem = mem0 - _free_mem(dev) if is_cuda else 0
    ts = []
    for _ in range(args.reps):
      if is_cuda:
        wp.synchronize_device()
      t = time.perf_counter()
      grad_once()
      if is_cuda:
        wp.synchronize_device()
      ts.append(time.perf_counter() - t)
    tg = float(np.median(ts))
    if args.save_grad:
      np.save(args.save_grad, g.astype(np.float32))
    out = dict(scene=args.scene, mode="fd", device=("cuda" if is_cuda else "cpu"),
               input=("ctrl" if use_ctrl else "qvel"), nu=nu, nv=nv, nq=nq,
               nparam=nparam, nworld=nworld, jit_s=jit, grad_s=tg, grads_per_s=1.0 / tg,
               mem_mib=mem / 1024**2)

  # ---------------------------------------------------------------- FWD --
  else:
    d = mjw.make_data(mjm, nworld=1)
    inp = wp.array(inp_base.reshape(1, -1), dtype=float)

    def fwd_once():
      if use_ctrl:
        reset(d, 1)
        wp.copy(d.ctrl, inp)
      else:
        reset(d, 1, qvel_np=inp_base.reshape(1, -1))
      mjw.step(m, d)
      return float(d.qpos.numpy().sum())

    t0 = time.perf_counter()
    fwd_once()
    if is_cuda:
      wp.synchronize_device()
    jit = time.perf_counter() - t0
    mem = mem0 - _free_mem(dev) if is_cuda else 0
    ts = []
    for _ in range(args.reps):
      if is_cuda:
        wp.synchronize_device()
      t = time.perf_counter()
      fwd_once()
      if is_cuda:
        wp.synchronize_device()
      ts.append(time.perf_counter() - t)
    tf = float(np.median(ts))
    out = dict(scene=args.scene, mode="fwd", device=("cuda" if is_cuda else "cpu"),
               input=("ctrl" if use_ctrl else "qvel"), nu=nu, nv=nv, nq=nq,
               nparam=nparam, nworld=1, jit_s=jit, fwd_s=tf, mem_mib=mem / 1024**2)

  print("RESULT_JSON " + json.dumps(out))


if __name__ == "__main__":
  main()

scripts/bench_accuracy.py (accuracy: AD (no eps) vs FD (eps sweep) vs float64 MuJoCo reference)

"""Accuracy: AD (eps-free) vs finite differences (eps-tuned), against a
float64 ground-truth gradient computed with MuJoCo's CPU C engine.

The maintenance question's flip side: if AD only matches FD, why carry it? The
answer is that FD has no single good eps -- its error is a U-curve (truncation
O(eps^2) for central differences, swamped by float32 cancellation as eps->0),
so you must tune eps per problem and still eat a noise floor. AD has no eps and
hits a fixed low error. We show both curves against the same reference.

Ground truth: dL/dctrl through one mujoco.mj_step (float64 C engine), central
differences at float64 with a well-conditioned eps. This is independent of warp.

Compared:
  warp AD  (float32, this PR)      -- one number, no eps
  warp FD  (float32, batched)      -- swept over eps in [1e-1 .. 1e-7]

Scene: humanoid (Newton, Euler; the FD-exact smooth path). Loss = sum qpos^2.

  PYTHONPATH=/home/johnny/Projects/dev/warp uv run python scripts/bench_accuracy.py
"""
import argparse
import json

import mujoco
import numpy as np
import warp as wp

import mujoco_warp as mjw
from etils import epath

SCENES = {
  "humanoid": "benchmarks/humanoid/humanoid.xml",
  "franka": "benchmarks/franka_emika_panda/scene.xml",
}


def mj_loss_f64(mjm, mjd0, ctrl):
  """One mj_step in float64, return sum(qpos^2)+0.1*sum(qvel^2)."""
  d = mujoco.MjData(mjm)
  d.qpos[:] = mjd0.qpos
  d.qvel[:] = mjd0.qvel
  d.act[:] = mjd0.act
  d.ctrl[:] = ctrl
  mujoco.mj_step(mjm, d)
  return float(np.sum(d.qpos**2) + 0.1 * np.sum(d.qvel**2))


@wp.kernel
def _lq(qpos: wp.array2d(dtype=float), wv: float, loss: wp.array(dtype=float)):
  w, i = wp.tid()
  wp.atomic_add(loss, 0, wv * qpos[w, i] * qpos[w, i])


@wp.kernel
def _lv(qvel: wp.array2d(dtype=float), wv: float, loss: wp.array(dtype=float)):
  w, i = wp.tid()
  wp.atomic_add(loss, 0, wv * qvel[w, i] * qvel[w, i])


@wp.kernel
def _pwq(qpos: wp.array2d(dtype=float), wv: float, lw: wp.array(dtype=float)):
  w, i = wp.tid()
  wp.atomic_add(lw, w, wv * qpos[w, i] * qpos[w, i])


@wp.kernel
def _pwv(qvel: wp.array2d(dtype=float), wv: float, lw: wp.array(dtype=float)):
  w, i = wp.tid()
  wp.atomic_add(lw, w, wv * qvel[w, i] * qvel[w, i])


def main():
  ap = argparse.ArgumentParser()
  ap.add_argument("--scene", default="humanoid", choices=list(SCENES))
  args = ap.parse_args()
  wp.config.quiet = True
  wp.init()
  wp.clear_kernel_cache()
  mjm = mujoco.MjSpec.from_file(epath.Path(SCENES[args.scene]).as_posix()).compile()
  mjm.opt.integrator = 0  # Euler, smooth FD-exact path
  nu, nq, nv = mjm.nu, mjm.nq, mjm.nv

  rng = np.random.default_rng(0)
  ctrl0 = (0.01 * rng.standard_normal(nu)).astype(np.float64)

  mjd0 = mujoco.MjData(mjm)
  mujoco.mj_forward(mjm, mjd0)

  # ---- float64 reference gradient via well-conditioned central FD (C engine) ----
  ref_eps = 1e-6
  g_ref = np.zeros(nu)
  for k in range(nu):
    cp = ctrl0.copy(); cp[k] += ref_eps
    cm = ctrl0.copy(); cm[k] -= ref_eps
    g_ref[k] = (mj_loss_f64(mjm, mjd0, cp) - mj_loss_f64(mjm, mjd0, cm)) / (2 * ref_eps)
  # Richardson check at 2*eps to confirm the reference is converged
  g_ref2 = np.zeros(nu)
  for k in range(nu):
    cp = ctrl0.copy(); cp[k] += 2 * ref_eps
    cm = ctrl0.copy(); cm[k] -= 2 * ref_eps
    g_ref2[k] = (mj_loss_f64(mjm, mjd0, cp) - mj_loss_f64(mjm, mjd0, cm)) / (4 * ref_eps)
  ref_selfcheck = float(np.linalg.norm(g_ref - g_ref2) / (np.linalg.norm(g_ref) + 1e-30))

  m = mjw.put_model(mjm)
  ctrl32 = ctrl0.astype(np.float32)

  def relerr(g):
    return float(np.linalg.norm(g - g_ref) / (np.linalg.norm(g_ref) + 1e-30))

  # ---- warp AD (float32), no eps ----
  d = mjw.make_diff_data(mjm, nworld=1)
  inp = wp.zeros((1, nu), dtype=float, requires_grad=True)
  inp.assign(ctrl32.reshape(1, nu))
  d.qpos.assign(mjd0.qpos.astype(np.float32).reshape(1, -1))
  d.qvel.zero_(); d.qacc.zero_(); d.act.zero_()
  loss = wp.zeros(1, dtype=float, requires_grad=True)
  tape = wp.Tape()
  with tape:
    wp.copy(d.ctrl, inp)
    mjw.step(m, d)
    wp.launch(_lq, dim=(1, nq), inputs=[d.qpos, 1.0], outputs=[loss])
    wp.launch(_lv, dim=(1, nv), inputs=[d.qvel, 0.1], outputs=[loss])
  tape.backward(loss=loss)
  g_ad = inp.grad.numpy().reshape(-1).astype(np.float64)
  ad_relerr = relerr(g_ad)

  # ---- warp FD (float32) swept over eps ----
  fd_curve = {}
  for eps in [1e-1, 3e-2, 1e-2, 3e-3, 1e-3, 3e-4, 1e-4, 3e-5, 1e-5, 1e-6, 1e-7]:
    nw = 2 * nu
    d2 = mjw.make_data(mjm, nworld=nw)
    pert = np.tile(ctrl32.reshape(1, -1), (nw, 1))
    for k in range(nu):
      pert[2 * k, k] += eps; pert[2 * k + 1, k] -= eps
    pwm = wp.array(pert.astype(np.float32), dtype=float)
    d2.qpos.assign(np.tile(mjd0.qpos.astype(np.float32), (nw, 1)))
    d2.qvel.zero_(); d2.qacc.zero_(); d2.act.zero_()
    wp.copy(d2.ctrl, pwm)
    mjw.step(m, d2)
    lw = wp.zeros(nw, dtype=float)
    wp.launch(_pwq, dim=(nw, nq), inputs=[d2.qpos, 1.0], outputs=[lw])
    wp.launch(_pwv, dim=(nw, nv), inputs=[d2.qvel, 0.1], outputs=[lw])
    L = lw.numpy().astype(np.float64)
    g_fd = np.array([(L[2 * k] - L[2 * k + 1]) / (2 * eps) for k in range(nu)])
    fd_curve[f"{eps:.0e}"] = relerr(g_fd)

  best_eps = min(fd_curve, key=fd_curve.get)
  out = dict(scene=args.scene, nu=nu, ref_eps=ref_eps, ref_selfcheck=ref_selfcheck,
             ad_relerr=ad_relerr, fd_curve=fd_curve,
             fd_best_eps=best_eps, fd_best_relerr=fd_curve[best_eps])
  print("RESULT_JSON " + json.dumps(out))
  print("\nReference self-check (||g(eps)-g(2eps)||/||g||):", f"{ref_selfcheck:.2e}", "(small = converged)")
  print(f"AD relerr (no eps): {ad_relerr:.3e}")
  print("FD relerr vs eps:")
  for e, r in fd_curve.items():
    mark = "  <- best" if e == best_eps else ""
    print(f"  eps={e}: {r:.3e}{mark}")


if __name__ == "__main__":
  main()

scripts/bench_horizon_scaling.py (BPTT regime: AD vs FD cost over horizon T)

"""Horizon scaling: AD vs FD COST (throughput + memory) for a T-step rollout
gradient on humanoid (the FD-exact Euler scene). nparam = T * nu.

This is the BPTT regime thowell's question targets: gradient of a rollout loss
w.r.t. the whole control trajectory. AD = one fwd + one bwd over T steps,
independent of nparam. FD = 2*nparam = 2*T*nu forward rollouts (batched into
worlds). We report median grad time and device memory per mode per T.

We measure COST only (not accuracy) here: multi-step AD accuracy is approximate
by a documented Warp tape limitation (grad_test.py:999); single-step accuracy is
established separately. Cost is exactly what thowell asked about.

One (mode, T) per process for a clean cold jit is overkill for the cost trend,
so we loop T in-process AFTER a warmup compile (the rollout kernels are reused;
only the python loop length changes), and report steady-state timings.

Run:
  PYTHONPATH=/home/johnny/Projects/dev/warp uv run python scripts/bench_horizon_scaling.py
"""
import json
import time

import mujoco
import numpy as np
import warp as wp

import mujoco_warp as mjw
from etils import epath

SCENE = "benchmarks/humanoid/humanoid.xml"
HORIZONS = [1, 2, 4, 8, 16, 32]
REPS = 10


@wp.kernel
def _apply(seq: wp.array3d(dtype=float), t: int, out: wp.array2d(dtype=float)):
  w, u = wp.tid()
  out[w, u] = seq[w, t, u]


@wp.kernel
def _lq(qpos: wp.array2d(dtype=float), wv: float, loss: wp.array(dtype=float)):
  w, i = wp.tid()
  wp.atomic_add(loss, 0, wv * qpos[w, i] * qpos[w, i])


@wp.kernel
def _pwq(qpos: wp.array2d(dtype=float), wv: float, lw: wp.array(dtype=float)):
  w, i = wp.tid()
  wp.atomic_add(lw, w, wv * qpos[w, i] * qpos[w, i])


def main():
  wp.config.quiet = True
  wp.init()
  wp.clear_kernel_cache()
  dev = wp.get_device()
  cuda = dev.is_cuda
  mjm = mujoco.MjSpec.from_file(epath.Path(SCENE).as_posix()).compile()
  m = mjw.put_model(mjm)
  nu, nq = mjm.nu, mjm.nq
  rng = np.random.default_rng(0)
  qpos0 = mjm.qpos0.astype(np.float32)

  def memnow():
    if cuda:
      wp.synchronize_device()
      return wp.get_device(dev).free_memory
    return 0

  results = []
  for T in HORIZONS:
    nparam = T * nu
    cb = (0.01 * rng.standard_normal((T, nu))).astype(np.float32)

    # ---- AD: nworld=1, tape over T steps ----
    free0 = memnow()
    d = mjw.make_diff_data(mjm, nworld=1)
    seq = wp.zeros((1, T, nu), dtype=float, requires_grad=True)
    seq.assign(cb.reshape(1, T, nu))

    def ad_once():
      d.qpos.assign(qpos0.reshape(1, -1)); d.qvel.zero_(); d.qacc.zero_(); d.act.zero_(); d.time.zero_()
      loss = wp.zeros(1, dtype=float, requires_grad=True)
      tape = wp.Tape()
      with tape:
        for t in range(T):
          wp.launch(_apply, dim=(1, nu), inputs=[seq, t], outputs=[d.ctrl])
          mjw.step(m, d)
          qp = wp.clone(d.qpos)
          wp.launch(_lq, dim=(1, nq), inputs=[qp, 1.0], outputs=[loss])
      tape.backward(loss=loss)
      g = seq.grad.numpy().copy()
      tape.zero()
      return g

    ad_once()  # warmup/compile
    ad_mem = (free0 - memnow()) / 1024**2 if cuda else 0.0
    ts = []
    for _ in range(REPS):
      t0 = time.perf_counter(); ad_once()
      if cuda: wp.synchronize_device()
      ts.append(time.perf_counter() - t0)
    ad_s = float(np.median(ts))
    del d, seq

    # ---- FD: nworld = 2*nparam, single batched rollout ----
    free0 = memnow()
    nworld = 2 * nparam
    d = mjw.make_data(mjm, nworld=nworld)
    pert = np.tile(cb.reshape(1, T * nu), (nworld, 1))
    for k in range(nparam):
      pert[2 * k, k] += 1e-3; pert[2 * k + 1, k] -= 1e-3
    seq = wp.array(pert.reshape(nworld, T, nu).astype(np.float32), dtype=float)

    def fd_once():
      d.qpos.assign(np.tile(qpos0, (nworld, 1))); d.qvel.zero_(); d.qacc.zero_(); d.act.zero_(); d.time.zero_()
      lw = wp.zeros(nworld, dtype=float)
      for t in range(T):
        wp.launch(_apply, dim=(nworld, nu), inputs=[seq, t], outputs=[d.ctrl])
        mjw.step(m, d)
        wp.launch(_pwq, dim=(nworld, nq), inputs=[d.qpos, 1.0], outputs=[lw])
      L = lw.numpy()
      return (L[0::2] - L[1::2]) / 2e-3

    fd_once()  # warmup
    fd_mem = (free0 - memnow()) / 1024**2 if cuda else 0.0
    ts = []
    for _ in range(REPS):
      t0 = time.perf_counter(); fd_once()
      if cuda: wp.synchronize_device()
      ts.append(time.perf_counter() - t0)
    fd_s = float(np.median(ts))
    del d, seq

    row = dict(T=T, nparam=nparam, fd_nworld=nworld,
               ad_s=ad_s, fd_s=fd_s, ad_grads_per_s=1.0 / ad_s, fd_grads_per_s=1.0 / fd_s,
               ad_mem_mib=ad_mem, fd_mem_mib=fd_mem, speedup_ad_over_fd=fd_s / ad_s)
    results.append(row)
    print("ROW " + json.dumps(row), flush=True)

  print("RESULTS_JSON " + json.dumps(results))


if __name__ == "__main__":
  main()

[johnnynunez]

Pushed a follow-up commit (614c011) fixing a crash in the backward Hessian solve for any model with more than 32 DOFs. The blocked-Cholesky path in _solve_hessian_system reshapes the right-hand side to (nworld, nv_pad, 1), but the incoming adjoint is only nv wide, so for nv > 32 the reshape failed with "Reshaped array must have the same total size as the original" (for example nv=81, nv_pad=96 on a sparse humanoid). The smaller per-world tile path was unaffected, which is why this stayed hidden. The fix copies the adjoint into an nv_pad-wide zero buffer before the reshape; the trailing rows stay zero, matching the padding DOFs that _padding_h_adjoint already handles, so the solve is unchanged for the real DOFs.

I also added a regression test covering the blocked-Cholesky branch for both dense and sparse jacobians at nv=40. There was no AD coverage above nv=32 before, which is why this was never caught; the test fails with the reshape error without the fix. Verified on an RTX PRO 6000 (sm_120) with the full grad suite green on GPU (31 passed) and CPU (16 passed, 13 skipped).

[johnnynunez]

Pushed a second follow-up (6da38f0) fixing a CUDA 719 illegal access in the adjoint blocked Cholesky solve. The adjoint passed m.nv as the runtime matrix size, but the blocked kernels load tiles at offsets that step by the tile size of 16. When nv is not a multiple of 16 the backward-substitution loop starts at nv - block_size and lands on unaligned offsets; with aligned tile loads and bounds checking off this reads out of bounds and faults. Any model with nv not a multiple of 16 and above the 32-DOF tile cutoff hit it, for example nv=33 or nv=81.

The forward solve already pads the Hessian to nv_pad (via _padding_h) and iterates at nv_pad, so the fix is to do the same on the adjoint side: pass nv_pad rather than nv as the runtime size in both the stored-factor solve and the full factorize-and-solve path. The padded system is SPD with an identity block on the padding diagonal and a zero right-hand side there, so the leading nv rows of the solution are unchanged.

Added a unit test that runs the blocked factorize+solve kernel on a known SPD matrix and checks against numpy for nv in {33, 40, 48, 64, 81}, covering aligned and unaligned sizes; the unaligned cases fault without the fix. Full grad suite green on GPU (36 passed) and CPU (16 passed, 18 skipped) on an RTX PRO 6000.

[johnnynunez]

Pushed a third follow-up (255d522) fixing illegal memory accesses when differentiating through flex (cloth and soft body) dynamics. The flex force and kinematics kernels (_flex_vertices and _flex_edges in smooth.py, _flex_elasticity and _flex_bending in passive.py, and _flex_normals in bvh.py) scan all flexes to find the one owning the current vertex, edge, or element and store its index in a local. That local was read after the loop without being initialized, so a thread whose element is not owned by any flex indexed the flex arrays with a stale id. The forward pass mostly got away with reading a harmless slot, but the autodiff-generated backward kernels indexed the adjoint arrays with the stale or negative id and faulted. A single cloth step backward reproduced it immediately.

The fix initializes the id to -1 and returns early when no owning flex is found, matching the pattern the flex collision kernels already use. Forward behavior is unchanged for valid threads. _flex_normals is in a module that currently disables backward, but it is fixed too for consistency and so it stays correct if backward is enabled there later, which matters for differentiable contact and tactile work.

Added a regression test that differentiates through a cloth grid step and checks for finite, nonzero gradients; it faults without the fix. Verified clean under compute-sanitizer with zero errors, and the flex forward tests still match MuJoCo on both GPU and CPU.

[johnnynunez]

Fixed a fourth issue found while chasing a NaN in the tactile SDF model (test_data/collision_sdf/tactile.xml).

The model produced NaN already in the forward pass, but only when started from make_data rather than from a put_data primed by mj_forward, which is why it looked grad-specific at first.
Bisecting the data fields showed the divergence was driven entirely by xaxis being zero for one joint.
The ball body in that model has five joints (three slide plus two hinge), and the unrolled joint loop introduced in the smooth-dynamics autodiff commits only processed up to four joints per body.
The fifth and sixth joints were never visited, so they kept their zero-initialized xaxis, and a zero joint axis propagates into the dynamics as NaN.

A non-free body can hold up to six joints (three slide plus three hinge), so I extended the unroll to six, which is the maximum.
Kinematics now matches MuJoCo for any joint count, and the forward NaN is gone.
The regression test covers five- and six-joint bodies and fails on the previous four-joint unroll.
Verified on GPU and CPU; the existing kinematics grad test still passes, and an AD-vs-FD check on a six-joint body agrees to 1.7e-4.

[johnnynunez]

Fixed a fifth issue found while sweeping subsystems that were not yet covered by the gradient tests (tendons, equality, passive springs, gravcomp).

A differentiable step on any model containing a spatial tendon faulted in the backward pass with an illegal memory access (CUDA 700) launching the adjoint of _spatial_site_tendon.
The forward was fine; the fault was in the autodiff-generated backward.
The cause was in the shared _accumulate_jac_chain helper, which walked the body chain toward the root and matched each dof against the tendon's sparse column indices using a pointer that decremented across both the outer body loop and the inner dof loop, and read the matched sparse index after the nested while loop that defined it.
That cross-scope pointer pattern is something Warp's reverse-mode autodiff cannot turn into a correct adjoint.
I replaced it with a self-contained forward scan of the sparse row per dof, which yields the same result because each dof appears at most once in a tendon row.
Since _spatial_geom_tendon shares the same helper, geom-wrap (pulley) tendons are now backward-safe too, which I confirmed separately.

Verified on GPU and CPU.
The forward ten_J and ten_length still match MuJoCo, the existing 51 tendon tests pass, and an AD-vs-FD check on the spatial-tendon Jacobian agrees.
The regression test runs a full differentiable step over a two-link chain spanned by a spatial tendon and faults on the previous code.

For the record, the other subsystems I swept in the same pass (fixed tendons, equality connect, equality joint coupling, passive joint springs, and gravcomp) already produced finite gradients through a differentiable step, so this was the only gap among them.

[johnnynunez]

While checking that the gradients are usable for control/RL on contact-rich tasks (not just matching finite differences on single steps), I found and fixed a control-gradient bug, and isolated a second one through contact that I want to flag.

Fixed here (1-line change plus tests): the diagonal-DOF inertia solve _factor_solve_simple was compiled with backward enabled, while the dense and sparse Cholesky solve kernels deliberately set enable_backward=False because the M-inverse gradient is supplied by the custom adjoint in _record_fwd_accel_adjoint. For models whose DOFs take the simple (diagonal) path, this double-counted the M-inverse VJP and doubled dL/dctrl. It's checkable in closed form: a 1-DOF unit-mass slider with a unit-gear motor, no gravity, no contact, one semi-implicit Euler step has d(qpos)/d(ctrl) = dt^2, and autodiff returned exactly 2*dt^2. Marking _factor_solve_simple enable_backward=False to match the other solve kernels restores the analytic value. Added GradIntegratorAnalyticTest (pins the single-step gradient to dt^2, fails at 2x without the fix).

Still open (added as an xfail test, GradContactMultiStepTest): gradients through a persistent contact over a multi-step rollout disagree with finite differences. I traced it to the retained solver Hessian used by the implicit-diff backward pass holding only the mass matrix M, not M + J^T D J — i.e. d.solver_h reads back as M even with an active contact (verified on a 1- and 2-DOF slider resting on the floor, nacon>=1). So the backward solves M x = b instead of H x = b and returns the free-body gradient, ignoring the contact stiffness. This is what breaks first-order policy optimization (SHAC-style) on hopper/cheetah while cartpole (no contact) trains fine. I scoped the fix for the H-retention out of this change since it touches the solver's retained-state plumbing; happy to take it as a follow-up or coordinate with whoever is closest to the implicit-diff path.

All of grad/smooth/solver/forward tests pass with the fix (398 passed, 2 xfailed for the documented contact issue). Measured on an RTX PRO 6000 (sm_120).

[johnnynunez]

Update (2026-06-29): commit c1ef142 completes the current fixed-point contact adjoint work. The current branch removes all demo, renderer, report, and generated-media artifacts from contrib/; they remain outside the upstream contribution.

The earlier strong-normal / pyramidal-friction ~2× gap is resolved. The backward now uses the converged Newton fixed point: it solves the retained post-solve Hessian, routes the direct smooth-force cotangent, differentiates the packed mass matrix through M → cdof/crb → cinert → qpos, and includes active aref, impedance, full-J geometry, and velocity-dissipation terms. Dense and packed-sparse Jacobians use the same colind-indexed VJP. Every active contact row, including effectively 1-DOF rows, consumes the native aref/vel cotangent after the explicit VJP so the position chain is not counted twice. Per-substep state and overwritten solver/contact intermediates are isolated for temporal composition.

Fresh verification on RTX PRO 6000 Blackwell / CUDA 13.3:

• affected suite: 217/217 passed (49 grad + 64 smooth + 104 forward);
• Hopper one-step local qpos/qvel VJP: dense 2.84e-5 / 8.95e-6, sparse 3.10e-5 / 3.00e-5 relative error;
• 32-substep PyTorch↔Warp Hopper action gradient: 8.68e-5 relative error, cosine 0.9999999979;
• active slide limit, dense and sparse: 1.29e-6 relative error;
• replicated sparse backward at 64/256 worlds: max inter-world gradient deviation 0.0;
• pre-commit, Ruff, kernel-analyzer, and git diff --check: pass.

Mutation checks are causal: removing the direct qfrc_smooth VJP makes the action gradient exactly zero; removing the mass term moves Hopper qpos error to ~4.96e-2; removing sparse full-J geometry moves it to ~9.59e-2; removing sparse limit expansion moves the active-limit error to ~0.997.

The validation boundary is explicit: these results cover MuJoCo's default pyramidal friction cone. Elliptic contact requires a coupled per-contact K block rather than diag(D) and is not claimed by this commit. Ball-joint limits, tendon limits, equality constraints, and independent friction constraints also remain outside the independently verified set. Exact active-set switches are nonsmooth; derivatives are validated within fixed active regimes.

[johnnynunez]

Contact and fixed-active constraint update pushed in cd82756.

This revision completes the converged Newton fixed-point VJP for active constraints. In addition to the existing pyramidal-contact path, elliptic CONE contacts now use the exact coupled per-contact block from the forward solver rather than a rowwise diag(D) approximation. The dense and packed-sparse paths cover condim=3 and condim=6, anisotropic friction, per-row solreffriction, and direct contact-frame Jacobian geometry. The differentiable replay also routes impedance/reference/Jacobian terms for equality constraints, tendon limits, and ball-joint limits, with per-substep tendon state isolated for multistep tapes.

Verification on an RTX PRO 6000 Blackwell:

• uv run pytest -q -n 8 --tb=short: 1054 passed, 7 skipped
• focused fixed-active regressions for equality, tendon, ball limits, Euler damping, and elliptic contacts: 17 passed
• Ruff, kernel-analyzer, and git diff --check: pass
• every elliptic finite-difference perturbation asserts the same contact count, constraint type, and CONE state before comparison
• mutation check: replacing tangential solreffriction with the normal solref makes both new condim=6 dense/sparse tests fail; restoring the per-row reference makes all four condim=3/6 cases pass

This supersedes the PR body's earlier statement that elliptic contact was not covered. The scope is intentionally local to a fixed contact topology and fixed active regime; exact contact creation/removal and active-set switching points remain nonsmooth. Contact-geometry derivatives retain the smooth primitive-pair coverage already described in the PR.

[johnnynunez]

Merged current main and pushed eeac6f2, a round of gradient-correctness fixes, each validated against central finite differences (grad_test grows from 61 to 87 tests):

• integrators: RK4 records one correctly-scoped solver adjoint per stage with isolated stage intermediates; implicit records an LDL-based adjoint instead of returning exactly zero dL/dctrl; implicitfast gets the (M - dt dF/dv)^-1 M correction; eulerdamp differentiates polynomial damping instead of capturing it as a record-time constant
• stepping: step1/step2 performs the same per-substep gradient isolation as step(); actuator_moment, actuator_length, and act_dot join the isolation set, fixing progressive corruption and sign flips for spatial-tendon and muscle actuators over rollouts
• solver backward: equality connect, weld, tendon, and flex rows carry their aref, Jacobian, and impedance state-dependence (connect and weld gradients were sign-flipped before); friction dof and tendon rows get the dissipation adjoint in the quadratic near-sticking regime via joint and tendon solref/solimp; CG solves assemble and factorize the Hessian at convergence so the same implicit backward applies, and backward fails with a clear host error when retained buffers were not populated instead of a CUDA illegal memory access
• previously dead gradient paths now flow: xfrc_applied, fluid forces (box and ellipsoid), flex passive elasticity and bending, smooth position- and velocity-stage sensors (with the differentiable/blocked sensor split documented in enable_grad), mocap inputs, and delayed actuators
• contacts: the smooth replay skips flex contacts instead of indexing geom_bodyid[-1] out of bounds (the grad-mode forward silently diverged from the plain forward on flex scenes); plane-box, sphere-box, and capsule-box are replayed differentiably; the capsule-capsule parallel branch selects its second contact by index rather than value dedup

The PR description's contact-scope paragraph is updated to match. Full suite passes (1288 passed, 18 skipped).