From 35c90844ae8fbebf287bc738ce5db766a112c553 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Mon, 27 Nov 2023 10:27:20 -0800 Subject: [PATCH] Add sites to MJX. PiperOrigin-RevId: 585691582 Change-Id: I6cc594daf436231f91534f137ff259e7b39811d9 --- doc/changelog.rst | 5 +++++ mjx/mujoco/mjx/_src/device.py | 2 ++ mjx/mujoco/mjx/_src/io.py | 2 ++ mjx/mujoco/mjx/_src/smooth.py | 17 ++++++++++----- mjx/mujoco/mjx/_src/smooth_test.py | 35 ++++++++++++++++++++++++++++++ mjx/mujoco/mjx/_src/types.py | 12 ++++++++++ 6 files changed, 68 insertions(+), 5 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 64ff3ec4..57802a76 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -23,6 +23,11 @@ Plugins - Added the `pid `__ 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. diff --git a/mjx/mujoco/mjx/_src/device.py b/mjx/mujoco/mjx/_src/device.py index c39860ef..ffd687b8 100644 --- a/mjx/mujoco/mjx/_src/device.py +++ b/mjx/mujoco/mjx/_src/device.py @@ -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)) diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 67e50a7e..55381f41 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -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), diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py index 2b421f2f..624291d8 100644 --- a/mjx/mujoco/mjx/_src/smooth.py +++ b/mjx/mujoco/mjx/_src/smooth.py @@ -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 diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py index 9e518672..bd9348bb 100644 --- a/mjx/mujoco/mjx/_src/smooth_test.py +++ b/mjx/mujoco/mjx/_src/smooth_test.py @@ -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(""" + + + + + + + + + + + + + + + + """) + 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() diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 5875ad7b..963109e6 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -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