Do not update qacc_warmstart at the end of the solver call; instead, update it at the same time as all other state variables. This change makes mj_forward idempotent.

See changelog herein for motivation and discussion.

PiperOrigin-RevId: 797392841
Change-Id: If50e9f09e1d07e8363169d3589c32097aeca2432

mjx/mujoco/mjx/_src/solver_test.py

@@ -78,11 +78,7 @@ class SolverTest(parameterized.TestCase):
       mjx_cost = ctx.cost - ctx.gauss
       _assert_eq(mj_cost, mjx_cost, 'cost')
 
-      # mj_forward overwrites qacc_warmstart, so let's restore it to what it was
-      # before the step so that MJX does not have a trivial solution
-      warmstart = d.qacc_warmstart.copy()
       mujoco.mj_forward(m, d)
-      d.qacc_warmstart = warmstart
       dx = jax.jit(mjx.solve)(mjx.put_model(m), mjx.put_data(m, d))
 
       # MJX finds very similar solutions with the newton solver
@@ -120,11 +116,7 @@ class SolverTest(parameterized.TestCase):
     # significant constraint forces keyframe 2
     mujoco.mj_resetDataKeyframe(m, d, 2)
 
-    # mj_forward overwrites qacc_warmstart, so let's restore it to what it was
-    # at the beginning of the step so that MJX does not have a trivial solution
-    warmstart = d.qacc_warmstart.copy()
     mujoco.mj_forward(m, d)
-    d.qacc_warmstart = warmstart
 
     dx = jax.jit(mjx.solve)(mjx.put_model(m), mjx.put_data(m, d))
 

mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_util.py

@@ -146,6 +146,7 @@ def fixture(
     mjd.mocap_quat = mocap_quat
 
   mujoco.mj_forward(mjm, mjd)
+  mjd.qacc_warmstart = mjd.qacc
   m = io.put_model(mjm)
   if ls_parallel is not None:
     m.opt.ls_parallel = ls_parallel

mjx/mujoco/mjx/warp/forward_test.py

@@ -192,14 +192,6 @@ class ForwardTest(parameterized.TestCase):
       tu.assert_attr_eq(dx, d, 'qfrc_smooth')
       tu.assert_attr_eq(dx, d, 'qacc_smooth')
 
-      # solve
-      np.testing.assert_allclose(
-          dx.qacc_warmstart,
-          d.qacc_warmstart,
-          err_msg='qacc_warmstart',
-          rtol=1e-5,
-          atol=1.0,
-      )
       np.testing.assert_allclose(
           dx.qacc, d.qacc, err_msg='qacc', rtol=1e-5, atol=1.0
       )