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

Fix bug where MJX data tendon fields are incorrect after calling put_data.

Browse Source 
PiperOrigin-RevId: 677911639
Change-Id: I75f536fa348e2941d9ec2d454f475cd067e6d3af
This commit is contained in:
Erik Frey
2024-09-23 13:11:16 -07:00
committed by Copybara-Service
parent 5219893813
commit 5838f84777
3 changed files with 32 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+1
View File
@@ -16,6 +16,7 @@ General
Bug fixes
^^^^^^^^^
- Fixed a bug where ``actuator_force`` was not set in MJX (:github:issue:`2068`).
- Fixed bug where MJX data tendon fields were incorrect after calling ``mjx.put_data``.
Version 3.2.3 (Sep 16, 2024)
mjx/mujoco/mjx/_src/io_test.py
+24 -4
View File
@@ -68,8 +68,8 @@ _MULTIPLE_CONSTRAINTS = """
<freejoint/>
<geom type="capsule" size=".2 .05"/>
<body name="cap2" pos=".6 -.3 .3">
<joint axis="0 1 0" type="hinge" range="-45 45"/>
<joint axis="1 0 0" type="hinge" range="-0.001 0.001"/>
<joint name="joint1" axis="0 1 0" type="hinge" range="-45 45"/>
<joint name="joint2" axis="1 0 0" type="hinge" range="-0.001 0.001"/>
<geom type="capsule" size=".2 .05"/>
<site pos="-0.214 -0.078 0" quat="0.664 0.664 -0.242 -0.242"/>
</body>
@@ -78,6 +78,12 @@ _MULTIPLE_CONSTRAINTS = """
<equality>
<connect body1="cap2" anchor="0 0 1"/>
</equality>
<tendon>
<fixed name="tendon_1" limited="false" stiffness=".1" damping=".2">
<joint joint="joint1" coef=".1"/>
<joint joint="joint2" coef="-.2"/>
</fixed>
</tendon>
</mujoco>
"""
@@ -85,8 +91,11 @@ _MULTIPLE_CONSTRAINTS = """
class ModelIOTest(parameterized.TestCase):
"""IO tests for mjx.Model."""
def test_put_model(self):
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS)
@parameterized.parameters(
_MULTIPLE_CONVEX_OBJECTS, _MULTIPLE_CONSTRAINTS
)
def test_put_model(self, xml):
m = mujoco.MjModel.from_xml_string(xml)
mx = mjx.put_model(m)
def assert_not_weak_type(x):
if isinstance(x, jax.Array):
@@ -126,6 +135,10 @@ class ModelIOTest(parameterized.TestCase):
np.testing.assert_allclose(mx.actuator_biastype, m.actuator_biastype)
np.testing.assert_allclose(mx.actuator_trnid, m.actuator_trnid)
np.testing.assert_equal(mx.wrap_type, m.wrap_type)
np.testing.assert_equal(mx.wrap_objid, m.wrap_objid)
np.testing.assert_equal(mx.wrap_prm, m.wrap_prm)
def test_fluid_params(self):
"""Test that has_fluid_params is set when fluid params are present."""
m = mjx.put_model(
@@ -338,6 +351,13 @@ class DataIOTest(parameterized.TestCase):
np.testing.assert_allclose(dx.geom_xmat.reshape((3, 9)), d.geom_xmat)
np.testing.assert_allclose(dx.site_xmat.reshape((1, 9)), d.site_xmat)
# tendon data is correct
np.testing.assert_allclose(dx.ten_length, d.ten_length)
np.testing.assert_equal(dx.ten_wrapadr, np.zeros((1,)))
np.testing.assert_equal(dx.ten_wrapnum, np.zeros((1,)))
np.testing.assert_equal(dx.wrap_obj, np.zeros((2, 2)))
np.testing.assert_equal(dx.wrap_xpos, np.zeros((2, 6)))
# efc_ are also shape transformed and padded
self.assertEqual(dx.efc_J.shape, (45, 8)) # nefc, nv
d_efc_j = d.efc_J.reshape((-1, 8))
mjx/mujoco/mjx/_src/types.py
+7 -7
View File
@@ -1029,9 +1029,9 @@ class Model(PyTreeNode):
tendon_length0: jax.Array
tendon_invweight0: jax.Array
tendon_hasfrictionloss: np.ndarray = _restricted_to('mjx')
wrap_type: np.ndarray = _restricted_to('mujoco')
wrap_objid: np.ndarray = _restricted_to('mujoco')
wrap_prm: np.ndarray = _restricted_to('mujoco')
wrap_type: np.ndarray
wrap_objid: np.ndarray
wrap_prm: np.ndarray
actuator_trntype: np.ndarray
actuator_dyntype: np.ndarray
actuator_gaintype: np.ndarray
@@ -1297,15 +1297,15 @@ class Data(PyTreeNode):
flexedge_J_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
flexedge_J: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
flexedge_length: jax.Array = _restricted_to('mujoco')
ten_wrapadr: jax.Array = _restricted_to('mujoco')
ten_wrapnum: jax.Array = _restricted_to('mujoco')
ten_wrapadr: jax.Array
ten_wrapnum: jax.Array
ten_J_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
ten_J_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
ten_J_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
ten_J: jax.Array # pylint:disable=invalid-name
ten_length: jax.Array
wrap_obj: jax.Array = _restricted_to('mujoco')
wrap_xpos: jax.Array = _restricted_to('mujoco')
wrap_obj: jax.Array
wrap_xpos: jax.Array
actuator_length: jax.Array
actuator_moment: jax.Array
crb: jax.Array