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

doc/changelog.rst

@@ -7,8 +7,29 @@ Upcoming version (not yet released)
 
 General
 ^^^^^^^
+.. admonition:: Breaking API changes
+   :class: attention
+
+   - The update of ``mjData.qacc_warmstart`` was moved from the end of the solver call (:ref:`mj_fwdConstraint`) to the
+     end of :ref:`mj_step`, and is now updated with all other state variables. This change makes :ref:`mj_forward`
+     fully idempotent.
+
+     Before this change, calling :ref:`mj_forward` repeatedly would make the constraint solver converge,
+     since each subsequent call would start from the previously updated ``qacc_warmstart`` value.
+     Indeed, this is precisely what happened in the viewer, which calls :ref:`mj_forward` repeatedly in PAUSE mode.
+
+     **Migration:** If your code depended on this behavior, you can recover it by updating manually after each
+     :ref:`mj_forward`: ``qacc_warmstart ← qacc``. The behavior is available in :ref:`simulate<saSimulate>` by
+     clicking the "Pause update" toggle (off by default).
+
+     Furthermore, this change has a numerical impact on the output of the :ref:`RK4 <geIntegrators>` integrator. Before
+     this change, due to the ``qacc_warmstart`` update occurring after each of the four Runge-Kutta substeps, the solver
+     convergence of RK4 was faster, at the cost of unprincipled integration. This change makes the RK4 integration
+     principled and well-defined. Since this change to RK4 is effectively a bug fix, migration to the previous behavior
+     is not provided.
+
 - Added support for shells with a curved reference configuration. See this `example
-   <https://github.com/google-deepmind/mujoco/blob/main/model/flex/basket.xml>`__
+  <https://github.com/google-deepmind/mujoco/blob/main/model/flex/basket.xml>`__.
 
 MJX
 ^^^
@@ -21,7 +42,7 @@ MJX
 
 
 Version 3.3.5 (August 8, 2025)
------------------------------------
+------------------------------
 
 General
 ^^^^^^^

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
       )

simulate/main.cc

@@ -285,7 +285,7 @@ mjModel* LoadModel(const char* file, mj::Simulate& sim) {
 
 // simulate in background thread (while rendering in main thread)
 void PhysicsLoop(mj::Simulate& sim) {
-  // cpu-sim syncronization point
+  // cpu-sim synchronization point
   std::chrono::time_point<mj::Simulate::Clock> syncCPU;
   mjtNum syncSim = 0;
 
@@ -438,6 +438,9 @@ void PhysicsLoop(mj::Simulate& sim) {
         else {
           // run mj_forward, to update rendering and joint sliders
           mj_forward(m, d);
+          if (sim.pause_update) {
+            mju_copy(d->qacc_warmstart, d->qacc, m->nv);
+          }
           sim.speed_changed = true;
         }
       }

simulate/simulate.cc

@@ -1922,8 +1922,8 @@ void Simulate::Sync(bool state_only) {
     return;
   }
 
-  bool update_profiler = this->profiler && (this->pause_update || this->run);
-  bool update_sensor = this->sensor && (this->pause_update || this->run);
+  bool update_profiler = this->profiler;
+  bool update_sensor = this->sensor;
 
   for (int i = 0; i < m_->njnt; ++i) {
     std::optional<std::pair<mjtNum, mjtNum>> range;

simulate/simulate.h

@@ -193,7 +193,7 @@ class Simulate {
   int info = 0;
   int profiler = 0;
   int sensor = 0;
-  int pause_update = 1;
+  int pause_update = 0;
   int fullscreen = 0;
   int vsync = 1;
   int busywait = 0;

src/engine/engine_forward.c

@@ -709,7 +709,6 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
   // no constraints: copy unconstrained acc, clear forces, return
   if (!nefc) {
     mju_copy(d->qacc, d->qacc_smooth, nv);
-    mju_copy(d->qacc_warmstart, d->qacc_smooth, nv);
     mju_zeroInt(d->solver_niter, mjNISLAND);
     TM_END(mjTIMER_CONSTRAINT);
     return;
@@ -783,9 +782,6 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
     }
   }
 
-  // save result for next step warmstart
-  mju_copy(d->qacc_warmstart, d->qacc, nv);
-
   // run noslip solver if enabled
   if (m->opt.noslip_iterations > 0) {
     mj_solNoSlip(m, d, m->opt.noslip_iterations);
@@ -837,6 +833,9 @@ static void mj_advance(const mjModel* m, mjData* d,
       }
     }
   }
+
+  // save qacc for next step warmstart
+  mju_copy(d->qacc_warmstart, d->qacc, m->nv);
 }
 
 // Euler integrator, semi-implicit in velocity, possibly skipping factorisation