chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Resolve discrepancy in RK4 between MJX and MJ.

Browse Source 
PiperOrigin-RevId: 743185866
Change-Id: I478faca1d8733999cc9a0682e931b19ee05e3843
This commit is contained in:
Baruch Tabanpour
2025-04-02 10:15:12 -07:00
committed by Copybara-Service
parent 40393f460a
commit 54f132f3f0
2 changed files with 8 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/forward.py
+7 -6
View File
@@ -317,7 +317,7 @@ def euler(m: Model, d: Data) -> Data:
@named_scope @named_scope
def rungekutta4(m: Model, d: Data) -> Data: def rungekutta4(m: Model, d: Data) -> Data:
"""Runge-Kutta explicit order 4 integrator.""" """Runge-Kutta explicit order 4 integrator."""
d_t0 = d d0 = d
# pylint: disable=invalid-name # pylint: disable=invalid-name
A, B = _RK4_A, _RK4_B A, B = _RK4_A, _RK4_B
C = jp.tril(A).sum(axis=0) # C(i) = sum_j A(i,j) C = jp.tril(A).sum(axis=0) # C(i) = sum_j A(i,j)
@@ -338,9 +338,9 @@ def rungekutta4(m: Model, d: Data) -> Data:
lambda k: a * k, (kqvel, d.qacc, d.act_dot) lambda k: a * k, (kqvel, d.qacc, d.act_dot)
) )
# get intermediate RK solutions # get intermediate RK solutions
kqpos = scan.flat(m, integrate_fn, 'jqv', 'q', m.jnt_type, d_t0.qpos, dqvel) kqpos = scan.flat(m, integrate_fn, 'jqv', 'q', m.jnt_type, d0.qpos, dqvel)
kact = d_t0.act + dact_dot * m.opt.timestep kact = d0.act + dact_dot * m.opt.timestep
kqvel = d_t0.qvel + dqacc * m.opt.timestep kqvel = d0.qvel + dqacc * m.opt.timestep
d = d.replace(qpos=kqpos, qvel=kqvel, act=kact, time=t) d = d.replace(qpos=kqpos, qvel=kqvel, act=kact, time=t)
d = forward(m, d) d = forward(m, d)
@@ -352,9 +352,10 @@ def rungekutta4(m: Model, d: Data) -> Data:
abt = jp.vstack([jp.diag(A), B[1:4], T]).T 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) 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 return d
mjx/mujoco/mjx/_src/forward_test.py
+1
View File
@@ -130,6 +130,7 @@ class ForwardTest(absltest.TestCase):
_assert_attr_eq(d, dx, 'qpos') _assert_attr_eq(d, dx, 'qpos')
_assert_attr_eq(d, dx, 'act') _assert_attr_eq(d, dx, 'act')
_assert_attr_eq(d, dx, 'time') _assert_attr_eq(d, dx, 'time')
_assert_attr_eq(d, dx, 'xpos')
def test_eulerdamp(self): def test_eulerdamp(self):
m = test_util.load_test_file('pendula.xml') m = test_util.load_test_file('pendula.xml')