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

Add sites to MJX.

Browse Source 
PiperOrigin-RevId: 585691582
Change-Id: I6cc594daf436231f91534f137ff259e7b39811d9
This commit is contained in:
Baruch Tabanpour
2023-11-27 10:27:20 -08:00
committed by Copybara-Service
parent 37bb03f435
commit 35c90844ae
6 changed files with 68 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+5
View File
@@ -23,6 +23,11 @@ Plugins
- Added the `pid <https://github.com/deepmind/mujoco/blob/main/plugin/actuator/README.md>`__ actuator plugin, a
configurable PID controller that implements the Integral term, which is not available with native MuJoCo actuators.
MJX
^^^
- Added ``site_xpos`` and ``site_xmat`` to MJX.
Bug fixes
^^^^^^^^^
- Fix bug in Cartesian actuation with movable refsite, as when using body-centric Cartesian actuators on a quadruped.
mjx/mujoco/mjx/_src/device.py
+2
View File
@@ -67,6 +67,7 @@ _TRANSFORMS = {
(types.Data, 'ximat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
(types.Data, 'xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
(types.Data, 'geom_xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
(types.Data, 'site_xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
(types.Model, 'actuator_trnid'): lambda x: x[:, 0],
(types.Contact, 'frame'): (
lambda x: x.reshape(x.shape[:-1] + (3, 3)) # pylint: disable=g-long-lambda
@@ -78,6 +79,7 @@ _INVERSE_TRANSFORMS = {
(types.Data, 'ximat'): lambda x: x.reshape(x.shape[:-2] + (9,)),
(types.Data, 'xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)),
(types.Data, 'geom_xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)),
(types.Data, 'site_xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)),
(types.Contact, 'frame'): (
lambda x: x.reshape(x.shape[:-2] + (9,)) # pylint: disable=g-long-lambda
if x is not None and x.shape[0] else jp.zeros((0, 9))
mjx/mujoco/mjx/_src/io.py
+2
View File
@@ -56,6 +56,8 @@ def make_data(m: Model) -> Data:
xaxis=jp.zeros((m.njnt, 3), dtype=jp.float32),
geom_xpos=jp.zeros((m.ngeom, 3), dtype=jp.float32),
geom_xmat=jp.zeros((m.ngeom, 3, 3), dtype=jp.float32),
site_xpos=jp.zeros((m.nsite, 3), dtype=jp.float32),
site_xmat=jp.zeros((m.nsite, 3, 3), dtype=jp.float32),
subtree_com=jp.zeros((m.nbody, 3), dtype=jp.float32),
cdof=jp.zeros((m.nv, 6), dtype=jp.float32),
cinert=jp.zeros((m.nbody, 10), dtype=jp.float32),
mjx/mujoco/mjx/_src/smooth.py
+12 -5
View File
@@ -101,13 +101,20 @@ def kinematics(m: Model, d: Data) -> Data:
# TODO(erikfrey): confirm that quats are more performant for mjx than mats
xipos, ximat = local_to_global(xpos, xquat, m.body_ipos, m.body_iquat)
geom_xpos, geom_xmat = local_to_global(
xpos[m.geom_bodyid], xquat[m.geom_bodyid], m.geom_pos, m.geom_quat
)
d = d.replace(qpos=qpos, xanchor=xanchor, xaxis=xaxis, xpos=xpos)
d = d.replace(xquat=xquat, xmat=xmat, xipos=xipos, ximat=ximat)
d = d.replace(geom_xpos=geom_xpos, geom_xmat=geom_xmat)
if m.ngeom:
geom_xpos, geom_xmat = local_to_global(
xpos[m.geom_bodyid], xquat[m.geom_bodyid], m.geom_pos, m.geom_quat
)
d = d.replace(geom_xpos=geom_xpos, geom_xmat=geom_xmat)
if m.nsite:
site_xpos, site_xmat = local_to_global(
xpos[m.site_bodyid], xquat[m.site_bodyid], m.site_pos, m.site_quat
)
d = d.replace(site_xpos=site_xpos, site_xmat=site_xmat)
return d
mjx/mujoco/mjx/_src/smooth_test.py
+35
View File
@@ -162,5 +162,40 @@ class DisableGravityTest(absltest.TestCase):
)
class SiteTest(absltest.TestCase):
def test_site(self):
"""Tests that site positions and orientations match MuJoCo."""
m = mujoco.MjModel.from_xml_string("""
<mujoco>
<worldbody>
<site name="origin"/>
<body>
<joint type="free"/>
<geom pos="1 0 1" type="box" size="0.1 0.01 0.01"/>
<site name="s1" pos="1.3 0 0"/> <!-- pos only -->
<site name="s2"/> <!-- no pos, no quat -->
<site name="s3" quat="1 0 1 0"/> <!- quat only -->
<site name="s4" pos="0 1.5 0" quat="1 0 1 0"/>
<site name="s5" pos="1 0 1"/> <!-- same as ipos -->
<site/>
</body>
</worldbody>
</mujoco>
""")
d = mujoco.MjData(m)
mx = mjx.device_put(m)
dx = mjx.device_put(d)
mujoco.mj_forward(m, d)
dx = mjx.forward(mx, dx)
np.testing.assert_array_almost_equal(dx.site_xpos, d.site_xpos)
np.testing.assert_array_almost_equal(
dx.site_xmat, d.site_xmat.reshape((-1, 3, 3))
)
if __name__ == '__main__':
absltest.main()
mjx/mujoco/mjx/_src/types.py
+12
View File
@@ -260,6 +260,7 @@ class Model(PyTreeNode):
nbody: number of bodies
njnt: number of joints
ngeom: number of geoms
nsite: number of sites
nmesh: number of meshes
npair: number of predefined geom pairs
nexclude: number of excluded geom pairs
@@ -327,6 +328,9 @@ class Model(PyTreeNode):
geom_friction: friction for (slide, spin, roll) (ngeom, 3)
geom_margin: include in solver if dist<margin-gap (ngeom,)
geom_gap: include in solver if dist<margin-gap (ngeom,)
site_bodyid: id of site's body (nsite,)
site_pos: local position offset rel. to body (nsite,3)
site_quat: local orientation offset rel. to body (nsite,4)
geom_convex_face: vertex face data, MJX only (ngeom,)
geom_convex_vert: vertex data, MJX only (ngeom,)
geom_convex_edge: unique edge data, MJX only (ngeom,)
@@ -377,6 +381,7 @@ class Model(PyTreeNode):
nbody: int
njnt: int
ngeom: int
nsite: int
nmesh: int
npair: int
nexclude: int
@@ -444,6 +449,9 @@ class Model(PyTreeNode):
geom_friction: jax.Array
geom_margin: jax.Array
geom_gap: jax.Array
site_bodyid: np.ndarray
site_pos: jax.Array
site_quat: jax.Array
pair_dim: np.ndarray
pair_geom1: np.ndarray
pair_geom2: np.ndarray
@@ -569,6 +577,8 @@ class Data(PyTreeNode):
xaxis: Cartesian joint axis (njnt, 3)
geom_xpos: Cartesian geom position (ngeom, 3)
geom_xmat: Cartesian geom orientation (ngeom, 3, 3)
site_xpos: Cartesian site position (nsite, 3)
site_xmat: Cartesian site orientation (nsite, 9)
subtree_com: center of mass of each subtree (nbody, 3)
cdof: com-based motion axis of each dof (nv, 6)
cinert: com-based body inertia and mass (nbody, 10)
@@ -631,6 +641,8 @@ class Data(PyTreeNode):
xaxis: jax.Array
geom_xpos: jax.Array
geom_xmat: jax.Array
site_xpos: jax.Array
site_xmat: jax.Array
subtree_com: jax.Array
cdof: jax.Array
cinert: jax.Array