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.

PiperOrigin-RevId: 797826112 Change-Id: Ibc51624adc3ec42c265958bf77696b33231431c1
2025-08-21 10:43:09 -07:00
parent 2f28473bc1
commit 3dec35a91e
2 changed files with 3 additions and 1 deletions
@@ -334,6 +334,9 @@ def _advance(
  # advance time
  time = d.time + m.opt.timestep

+  # save qacc for next step warmstart
+  d = d.replace(qacc_warmstart=d.qacc)
+
  return d.replace(act=act, qpos=qpos, time=time)


@@ -602,7 +602,6 @@ def solve(m: Model, d: Data) -> Data:
    ctx = jax.lax.while_loop(cond, body, ctx)

  d = d.tree_replace({
-      'qacc_warmstart': ctx.qacc,
      'qfrc_constraint': ctx.qfrc_constraint,
      'qacc': ctx.qacc,
      '_impl.efc_force': ctx.efc_force,