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

Update MJX inverse dynamics for tendon armature.

Browse Source 
PiperOrigin-RevId: 775658246
Change-Id: Ie6a71b971d4b165e0235689b7b4ee3246486e10d
This commit is contained in:
Taylor Howell
2025-06-25 06:13:59 -07:00
committed by Copybara-Service
parent da831f5d3e
commit d0faa3624e
4 changed files with 67 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/inverse.py
+6 -2
View File
@@ -93,11 +93,15 @@ def inverse(m: Model, d: Data) -> Data:
d = discrete_acc(m, d)
d = inv_constraint(m, d)
d = smooth.rne(m, d, flg_acc=True)
d = smooth.rne(m, d)
d = smooth.tendon_bias(m, d)
d = sensor.sensor_acc(m, d)
qfrc_inverse = (
d.qfrc_bias + m.dof_armature * d.qacc - d.qfrc_passive - d.qfrc_constraint
d.qfrc_bias
+ support.mul_m(m, d, d.qacc)
- d.qfrc_passive
- d.qfrc_constraint
)
if m.opt.enableflags & EnableBit.INVDISCRETE:
mjx/mujoco/mjx/_src/inverse_test.py
+36
View File
@@ -19,6 +19,7 @@ from jax import numpy as jp
import mujoco
from mujoco import mjx
from mujoco.mjx._src import support
from mujoco.mjx._src import test_util
import numpy as np
# tolerance for difference between MuJoCo and MJX calculations - mostly
@@ -109,6 +110,41 @@ class InverseTest(parameterized.TestCase):
self.assertLess(fwdinv1, 1.0e-3)
_assert_eq(dxinv.qacc, dx.qacc, 'qacc')
def test_inverse_tendon_armature(self):
m = test_util.load_test_file('tendon/armature.xml')
d = mujoco.MjData(m)
d.qvel = np.random.uniform(low=-0.01, high=0.01, size=d.qvel.shape)
d.ctrl = np.random.uniform(low=-0.01, high=0.01, size=d.ctrl.shape)
d.qfrc_applied = np.random.uniform(
low=-0.01, high=0.01, size=d.qfrc_applied.shape
)
d.xfrc_applied = np.random.uniform(
low=-0.01, high=0.01, size=d.xfrc_applied.shape
)
mujoco.mj_step(m, d, 10)
mx = mjx.put_model(m)
dx = mjx.put_data(m, d)
dx = mjx.forward(mx, dx)
dxinv = mjx.inverse(mx, dx)
fwdinv0 = jp.linalg.norm(
dxinv.qfrc_constraint - dx.qfrc_constraint, ord=np.inf
)
fwdinv1 = jp.linalg.norm(
dxinv.qfrc_inverse
- (
dx.qfrc_applied + dx.qfrc_actuator + support.xfrc_accumulate(mx, dx)
),
ord=np.inf,
)
self.assertLess(fwdinv0, 1.0e-3)
self.assertLess(fwdinv1, 1.0e-3)
_assert_eq(dxinv.qacc, dx.qacc, 'qacc')
if __name__ == '__main__':
absltest.main()
mjx/mujoco/mjx/_src/smooth_test.py
+1 -26
View File
@@ -321,32 +321,7 @@ class TendonTest(parameterized.TestCase):
@parameterized.parameters(JacobianType.DENSE, JacobianType.SPARSE)
def test_tendon_armature(self, jacobian):
"""Tests MJX tendon armature matches MuJoCo."""
m = mujoco.MjModel.from_xml_string("""
<mujoco>
<worldbody>
<site name="site0" pos="1 0 1"/>
<body>
<joint type="slide" axis="0 0 1"/>
<joint type="hinge" axis="0 1 0"/>
<geom type="box" size="0.1 0.1 0.1" mass="1" pos="1 0 0"/>
<site name="site1"/>
</body>
</worldbody>
<tendon>
<spatial armature="123">
<site site="site0"/>
<site site="site1"/>
</spatial>
<spatial armature="456">
<site site="site0"/>
<site site="site1"/>
</spatial>
</tendon>
<keyframe>
<key qpos="1.2345 1.2345" qvel="1.2345 1.2345"/>
</keyframe>
</mujoco>
""")
m = test_util.load_test_file('tendon/armature.xml')
m.opt.jacobian = jacobian
d = mujoco.MjData(m)
mujoco.mj_resetDataKeyframe(m, d, 0)
mjx/mujoco/mjx/test_data/tendon/armature.xml
+24
View File
@@ -0,0 +1,24 @@
<mujoco>
<worldbody>
<site name="site0" pos="1 0 1"/>
<body>
<joint type="slide" axis="0 0 1"/>
<joint type="hinge" axis="0 1 0"/>
<geom type="box" size="0.1 0.1 0.1" mass="1" pos="1 0 0"/>
<site name="site1"/>
</body>
</worldbody>
<tendon>
<spatial armature="123">
<site site="site0"/>
<site site="site1"/>
</spatial>
<spatial armature="456">
<site site="site0"/>
<site site="site1"/>
</spatial>
</tendon>
<keyframe>
<key qpos="1.2345 1.2345" qvel="1.2345 1.2345"/>
</keyframe>
</mujoco>