From a68141eeffaac14f69e8464ab1ca477bc90e0cc5 Mon Sep 17 00:00:00 2001 From: Erik Frey Date: Mon, 19 Aug 2024 11:11:34 -0700 Subject: [PATCH] Implicitfast integration for MJX. PiperOrigin-RevId: 664889120 Change-Id: Id3bd3916fe821ab2af79c52e53e14837ae829b2a --- doc/changelog.rst | 8 +++-- doc/mjx.rst | 4 +-- mjx/mujoco/mjx/__init__.py | 3 +- mjx/mujoco/mjx/_src/forward.py | 42 ++++++++++++++++++++++++ mjx/mujoco/mjx/_src/forward_test.py | 6 ++++ mjx/mujoco/mjx/_src/io.py | 7 +++- mjx/mujoco/mjx/_src/io_test.py | 8 +++++ mjx/mujoco/mjx/_src/types.py | 4 ++- mjx/mujoco/mjx/test_data/constraints.xml | 2 +- 9 files changed, 75 insertions(+), 9 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 87cbebea..9b47e771 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -23,17 +23,19 @@ MJX 7. Added velocity-dependent sensors: ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``. 8. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device. 9. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``. +10. Added support for :ref:`implicitfast integration` for all cases except + :doc:`fluid drag `. Bug fixes ^^^^^^^^^ -10. Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`, +11. Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`, contribution by :github:user:`michael-ahn`). -11. Fixed a bug wherein, for models that have both muscles and stateless actuators and used one of the implicit +12. Fixed a bug wherein, for models that have both muscles and stateless actuators and used one of the implicit integrators, wrong derivatives would be computed. Python bindings ^^^^^^^^^^^^^^^ -12. Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`). +13. Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`). Version 3.2.2 (Aug 8, 2024) diff --git a/doc/mjx.rst b/doc/mjx.rst index 9aafa764..90a43e87 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -202,7 +202,7 @@ The following features are **fully supported** in MJX: * - :ref:`Equality ` - ``CONNECT``, ``WELD``, ``JOINT``, ``TENDON`` * - :ref:`Integrator ` - - ``EULER``, ``RK4`` + - ``EULER``, ``RK4``, ``IMPLICITFAST`` (``IMPLICITFAST`` not supported with :doc:`fluid drag `) * - :ref:`Cone ` - ``PYRAMIDAL``, ``ELLIPTIC`` * - :ref:`Condim ` @@ -229,7 +229,7 @@ The following features are **in development** and coming soon: * - :ref:`Constraint ` - :ref:`Frictionloss `, ``FRICTION_DOF`` * - :ref:`Integrator ` - - ``IMPLICIT``, ``IMPLICITFAST`` + - ``IMPLICIT`` * - Dynamics - :ref:`Inverse ` * - :ref:`Actuator Dynamics ` diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py index 2b442d77..a4a8a5da 100644 --- a/mjx/mujoco/mjx/__init__.py +++ b/mjx/mujoco/mjx/__init__.py @@ -23,6 +23,7 @@ from mujoco.mjx._src.forward import fwd_acceleration from mujoco.mjx._src.forward import fwd_actuation from mujoco.mjx._src.forward import fwd_position from mujoco.mjx._src.forward import fwd_velocity +from mujoco.mjx._src.forward import implicit from mujoco.mjx._src.forward import rungekutta4 from mujoco.mjx._src.forward import step from mujoco.mjx._src.io import get_data @@ -32,9 +33,9 @@ from mujoco.mjx._src.io import put_data from mujoco.mjx._src.io import put_model from mujoco.mjx._src.passive import passive from mujoco.mjx._src.ray import ray +from mujoco.mjx._src.sensor import sensor_acc from mujoco.mjx._src.sensor import sensor_pos from mujoco.mjx._src.sensor import sensor_vel -from mujoco.mjx._src.sensor import sensor_acc from mujoco.mjx._src.smooth import camlight from mujoco.mjx._src.smooth import com_pos from mujoco.mjx._src.smooth import com_vel diff --git a/mjx/mujoco/mjx/_src/forward.py b/mjx/mujoco/mjx/_src/forward.py index a90ef854..456e748c 100644 --- a/mjx/mujoco/mjx/_src/forward.py +++ b/mjx/mujoco/mjx/_src/forward.py @@ -348,6 +348,46 @@ def rungekutta4(m: Model, d: Data) -> Data: return d +@named_scope +def implicit(m: Model, d: Data) -> Data: + """Integrates fully implicit in velocity.""" + + qderiv = None + + # qDeriv += d qfrc_actuator / d qvel + if not m.opt.disableflags & DisableBit.ACTUATION: + affine_bias = m.actuator_biastype == BiasType.AFFINE + bias_vel = m.actuator_biasprm[:, 2] * affine_bias + affine_gain = m.actuator_gaintype == GainType.AFFINE + gain_vel = m.actuator_gainprm[:, 2] * affine_gain + ctrl = d.ctrl.at[m.actuator_dyntype != DynType.NONE].set(d.act) + vel = bias_vel + gain_vel * ctrl + qderiv = d.actuator_moment.T @ jp.diag(vel) @ d.actuator_moment + + # qDeriv += d qfrc_passive / d qvel + if not m.opt.disableflags & DisableBit.PASSIVE: + if qderiv is None: + qderiv = -jp.diag(m.dof_damping) + else: + qderiv -= jp.diag(m.dof_damping) + if m.ntendon: + qderiv -= d.ten_J.T @ jp.diag(m.tendon_damping) @ d.ten_J + # TODO(robotics-simulation): fluid drag model + if m.opt.has_fluid_params: + raise NotImplementedError('fluid drag not supported for implicitfast') + + qacc = d.qacc + if qderiv is not None: + # TODO(robotics-simulation): use smooth.factor_m / solve_m here: + qm = support.full_m(m, d) if support.is_sparse(m) else d.qM + qm -= m.opt.timestep * qderiv + qh, _ = jax.scipy.linalg.cho_factor(qm) + qfrc = d.qfrc_smooth + d.qfrc_constraint + qacc = jax.scipy.linalg.cho_solve((qh, False), qfrc) + + return _advance(m, d, d.act_dot, qacc) + + @named_scope def forward(m: Model, d: Data) -> Data: """Forward dynamics.""" @@ -377,6 +417,8 @@ def step(m: Model, d: Data) -> Data: d = euler(m, d) elif m.opt.integrator == IntegratorType.RK4: d = rungekutta4(m, d) + elif m.opt.integrator == IntegratorType.IMPLICITFAST: + d = implicit(m, d) else: raise NotImplementedError(f'integrator {m.opt.integrator} not implemented.') diff --git a/mjx/mujoco/mjx/_src/forward_test.py b/mjx/mujoco/mjx/_src/forward_test.py index 0a777797..a93f04e0 100644 --- a/mjx/mujoco/mjx/_src/forward_test.py +++ b/mjx/mujoco/mjx/_src/forward_test.py @@ -68,6 +68,12 @@ class ForwardTest(absltest.TestCase): _assert_attr_eq(d, dx, 'qpos') _assert_attr_eq(d, dx, 'time') + # implicitfast + m.opt.integrator = mujoco.mjtIntegrator.mjINT_IMPLICITFAST + dx = jax.jit(mjx.implicit)(mx, mjx.put_data(m, d)) + mujoco.mj_implicit(m, d) + _assert_attr_eq(d, dx, 'qpos') + def test_step(self): m = test_util.load_test_file('constraints.xml') d = mujoco.MjData(m) diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 7ad1c261..4caf0be9 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -47,13 +47,18 @@ def _make_option(o: mujoco.MjOption) -> types.Option: if o.enableflags & 2**i: raise NotImplementedError(f'{mujoco.mjtEnableBit(2 ** i)}') + has_fluid_params = o.density > 0 or o.viscosity > 0 or o.wind.any() + implicitfast = o.integrator == mujoco.mjtIntegrator.mjINT_IMPLICITFAST + if implicitfast and has_fluid_params: + raise NotImplementedError('implicitfast not implemented for fluid drag.') + fields = {f.name: getattr(o, f.name, None) for f in types.Option.fields()} fields['integrator'] = types.IntegratorType(o.integrator) fields['cone'] = types.ConeType(o.cone) fields['jacobian'] = types.JacobianType(o.jacobian) fields['solver'] = types.SolverType(o.solver) fields['disableflags'] = types.DisableBit(o.disableflags) - fields['has_fluid_params'] = o.density > 0 or o.viscosity > 0 or o.wind.any() + fields['has_fluid_params'] = has_fluid_params return types.Option(**fields) diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index 020e8e23..97c61382 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -208,6 +208,14 @@ class ModelIOTest(parameterized.TestCase): """)) + def test_implicitfast_fluid_not_implemented(self): + with self.assertRaises(NotImplementedError): + mjx.put_model(mujoco.MjModel.from_xml_string(""" + + """)) + class DataIOTest(parameterized.TestCase): """IO tests for mjx.Data.""" diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 15e8bb18..4a14db17 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -90,10 +90,12 @@ class IntegratorType(enum.IntEnum): Members: EULER: semi-implicit Euler RK4: 4th-order Runge Kutta + IMPLICITFAST: implicit in velocity, no rne derivative """ EULER = mujoco.mjtIntegrator.mjINT_EULER RK4 = mujoco.mjtIntegrator.mjINT_RK4 - # unsupported: IMPLICIT, IMPLICITFAST + IMPLICITFAST = mujoco.mjtIntegrator.mjINT_IMPLICITFAST + # unsupported: IMPLICIT class GeomType(enum.IntEnum): diff --git a/mjx/mujoco/mjx/test_data/constraints.xml b/mjx/mujoco/mjx/test_data/constraints.xml index 8f7e5643..a1476a45 100644 --- a/mjx/mujoco/mjx/test_data/constraints.xml +++ b/mjx/mujoco/mjx/test_data/constraints.xml @@ -83,7 +83,7 @@ - +