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

Don't populate _qM_sparse, _qLD_sparse, _qLDiagInv_sparse fields if they were never set via _full_compat. Fixes #2188

Browse Source 
PiperOrigin-RevId: 691538016
Change-Id: Ib0385166af2dfbbb46cf1986cb692e1fda467a0b
This commit is contained in:
Baruch Tabanpour
2024-10-30 13:42:10 -07:00
committed by Copybara-Service
parent a9737c60e8
commit 26ccaeb78e
2 changed files with 8 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/smooth.py
+5 -2
View File
@@ -291,7 +291,7 @@ def crb(m: Model, d: Data) -> Data:
crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d.cdof)
qm = support.make_m(m, crb_cdof, d.cdof, m.dof_armature)
d = d.replace(qM=qm)
if support.is_sparse(m):
if support.is_sparse(m) and d._qM_sparse.size > 0: # pylint: disable=protected-access
d = d.replace(_qM_sparse=qm)
return d
@@ -353,7 +353,10 @@ def factor_m(m: Model, d: Data) -> Data:
qld = (qld / qld[jp.array(madr_ds)]).at[m.dof_Madr].set(qld_diag)
d = d.replace(qLD=qld, qLDiagInv=1 / qld_diag)
d = d.replace(_qLD_sparse=d.qLD, _qLDiagInv_sparse=d.qLDiagInv)
if d._qLD_sparse.size > 0: # pylint: disable=protected-access
d = d.replace(_qLD_sparse=d.qLD)
if d._qLDiagInv_sparse.size > 0: # pylint: disable=protected-access
d = d.replace(_qLDiagInv_sparse=d.qLDiagInv)
return d
mjx/mujoco/mjx/_src/smooth_test.py
+3
View File
@@ -87,10 +87,13 @@ class SmoothTest(absltest.TestCase):
dx = jax.jit(mjx.crb)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx, 'crb')
_assert_attr_eq(d, dx, 'qM')
_assert_eq(dx._qM_sparse, np.zeros(0), '_qM_sparse')
# factor_m
dx = jax.jit(mjx.factor_m)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx, 'qLD')
_assert_attr_eq(d, dx, 'qLDiagInv')
_assert_eq(dx._qLD_sparse, np.zeros(0), '_qLD_sparse')
_assert_eq(dx._qLDiagInv_sparse, np.zeros(0), '_qLDiagInv_sparse')
# com_vel
dx = jax.jit(mjx.com_vel)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx, 'cvel')