Resolve discrepancy in RK4 between MJX and MJ.

PiperOrigin-RevId: 743185866
Change-Id: I478faca1d8733999cc9a0682e931b19ee05e3843

Author: btaba

mjx/mujoco/mjx/_src/forward.py

@@ -317,7 +317,7 @@ def euler(m: Model, d: Data) -> Data:
 @named_scope
 def rungekutta4(m: Model, d: Data) -> Data:
   """Runge-Kutta explicit order 4 integrator."""
-  d_t0 = d
+  d0 = d
   # pylint: disable=invalid-name
   A, B = _RK4_A, _RK4_B
   C = jp.tril(A).sum(axis=0)  # C(i) = sum_j A(i,j)
@@ -338,9 +338,9 @@ def f(carry, x):
         lambda k: a * k, (kqvel, d.qacc, d.act_dot)
     )
     # get intermediate RK solutions
-    kqpos = scan.flat(m, integrate_fn, 'jqv', 'q', m.jnt_type, d_t0.qpos, dqvel)
-    kact = d_t0.act + dact_dot * m.opt.timestep
-    kqvel = d_t0.qvel + dqacc * m.opt.timestep
+    kqpos = scan.flat(m, integrate_fn, 'jqv', 'q', m.jnt_type, d0.qpos, dqvel)
+    kact = d0.act + dact_dot * m.opt.timestep
+    kqvel = d0.qvel + dqacc * m.opt.timestep
     d = d.replace(qpos=kqpos, qvel=kqvel, act=kact, time=t)
     d = forward(m, d)
 
@@ -352,9 +352,10 @@ def f(carry, x):
 
   abt = jp.vstack([jp.diag(A), B[1:4], T]).T
   out, _ = jax.lax.scan(f, (qvel, qacc, act_dot, kqvel, d), abt, unroll=3)
-  qvel, qacc, act_dot, *_ = out
+  qvel, qacc, act_dot, _, d1 = out
 
-  d = _advance(m, d_t0, act_dot, qacc, qvel)
+  d = d1.replace(qpos=d0.qpos, qvel=d0.qvel, act=d0.act, time=d0.time)
+  d = _advance(m, d, act_dot, qacc, qvel)
   return d
 
 

mjx/mujoco/mjx/_src/forward_test.py

@@ -130,6 +130,7 @@ def test_rk4(self):
     _assert_attr_eq(d, dx, 'qpos')
     _assert_attr_eq(d, dx, 'act')
     _assert_attr_eq(d, dx, 'time')
+    _assert_attr_eq(d, dx, 'xpos')
 
   def test_eulerdamp(self):
     m = test_util.load_test_file('pendula.xml')