From c7a82c32dc33adf1a819ba4dfc7dddaafa1eae8e Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Wed, 20 Aug 2025 11:29:13 -0700 Subject: [PATCH] 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 | 25 +++++++++++++++++-- mjx/mujoco/mjx/_src/solver_test.py | 8 ------ .../third_party/mujoco_warp/_src/test_util.py | 1 + mjx/mujoco/mjx/warp/forward_test.py | 8 ------ simulate/main.cc | 5 +++- simulate/simulate.cc | 4 +-- simulate/simulate.h | 2 +- src/engine/engine_forward.c | 7 +++--- 8 files changed, 34 insertions(+), 26 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 9afb1087..26541fff 100644 --- a/doc/changelog.rst +++ b/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` by + clicking the "Pause update" toggle (off by default). + + Furthermore, this change has a numerical impact on the output of the :ref:`RK4 ` 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 - `__ + `__. MJX ^^^ @@ -21,7 +42,7 @@ MJX Version 3.3.5 (August 8, 2025) ------------------------------------ +------------------------------ General ^^^^^^^ diff --git a/mjx/mujoco/mjx/_src/solver_test.py b/mjx/mujoco/mjx/_src/solver_test.py index 0b1b3b5a..928d6a9b 100644 --- a/mjx/mujoco/mjx/_src/solver_test.py +++ b/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)) diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_util.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_util.py index be951a31..8eca4b42 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_util.py +++ b/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 diff --git a/mjx/mujoco/mjx/warp/forward_test.py b/mjx/mujoco/mjx/warp/forward_test.py index ff890658..6f50f22d 100644 --- a/mjx/mujoco/mjx/warp/forward_test.py +++ b/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 ) diff --git a/simulate/main.cc b/simulate/main.cc index afe8b720..d52e4c9d 100644 --- a/simulate/main.cc +++ b/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 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; } } diff --git a/simulate/simulate.cc b/simulate/simulate.cc index 17ac9e82..30795dca 100644 --- a/simulate/simulate.cc +++ b/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> range; diff --git a/simulate/simulate.h b/simulate/simulate.h index 91ba5fcc..c1fd960e 100644 --- a/simulate/simulate.h +++ b/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; diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index 9e19cb32..2d4044c8 100644 --- a/src/engine/engine_forward.c +++ b/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