diff --git a/doc/changelog.rst b/doc/changelog.rst index cb27d8b4..40b74e01 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -26,21 +26,23 @@ MJX 5. Added :at:`site` transmission. 6. Updated MJX colab tutorial with more stable quadruped environment. 7. Added ``mjx.ray`` which mirrors :ref:`mj_ray` for planes, spheres, capsules, boxes, and meshes. +8. Added ``mjx.is_sparse`` which mirrors :ref:`mj_isSparse` and ``mjx.full_m`` which mirrors :ref:`mj_fullM`. +9. Added support for specifying sparse or dense mass matrices via :ref:`option-jacobian`. Python bindings ^^^^^^^^^^^^^^^ -8. Improved the implmentation of the :ref:`rollout` module. Note the changes below are breaking, dependent - code will require modification. +10. Improved the implmentation of the :ref:`rollout` module. Note the changes below are breaking, dependent + code will require modification. - - Uses :ref:`mjSTATE_FULLPHYSICS` as state spec, enabling divergence detection by inspecting time. - - Allows user-defined control spec for any combination of :ref:`user input` fields as controls. - - Outputs are no longer squeezed and always have dim=3. + - Uses :ref:`mjSTATE_FULLPHYSICS` as state spec, enabling divergence detection by inspecting time. + - Allows user-defined control spec for any combination of :ref:`user input` fields as controls. + - Outputs are no longer squeezed and always have dim=3. Bug fixes ^^^^^^^^^ -9. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes - :github:issue:`1270`. -10. Fixed a bug in the :ref:`muscle model` that led to non-zero values outside the lower +11. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes + :github:issue:`1270`. +12. Fixed a bug in the :ref:`muscle model` that led to non-zero values outside the lower bound of the length range. Fixes :github:issue:`1342`. diff --git a/doc/mjx.rst b/doc/mjx.rst index 251e05fc..50c47389 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -349,3 +349,11 @@ For MJX to perform well, some configuration parameters should be adjusted from t :ref:`option-flag` element Disabling ``eulerdamp`` can help performance and is often not needed for stability. + +:ref:`option-jacobian` element + Explicitly setting "dense" or "sparse" may speed up simulation depending on your device. Modern TPUs have specialized + hardware for rapidly operating over sparse matrices, whereas GPUs tend to be faster with dense matrices as long as + they fit onto the device. As such, the behavior in MJX for the default "auto" setting is sparse if ``nv`` is 60 or + greater, or if MJX detects a TPU as the default backend, otherwise "dense". For TPU, using "sparse" with the + Newton solver can speed up simulation by 2x to 3x. For GPU, choosing "dense" may impart a more modest speedup of 10% + to 20%, as long as the dense matrices can fit on the device. diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py index f4953c60..737fe2c5 100644 --- a/mjx/mujoco/mjx/__init__.py +++ b/mjx/mujoco/mjx/__init__.py @@ -39,8 +39,10 @@ from mujoco.mjx._src.smooth import com_vel from mujoco.mjx._src.smooth import crb from mujoco.mjx._src.smooth import factor_m from mujoco.mjx._src.smooth import kinematics -from mujoco.mjx._src.smooth import mul_m from mujoco.mjx._src.smooth import rne from mujoco.mjx._src.smooth import transmission from mujoco.mjx._src.solver import solve +from mujoco.mjx._src.support import is_sparse +from mujoco.mjx._src.support import full_m +from mujoco.mjx._src.support import mul_m from mujoco.mjx._src.types import * diff --git a/mjx/mujoco/mjx/_src/device_test.py b/mjx/mujoco/mjx/_src/device_test.py index 71137c5d..6126f707 100644 --- a/mjx/mujoco/mjx/_src/device_test.py +++ b/mjx/mujoco/mjx/_src/device_test.py @@ -75,6 +75,7 @@ class DeviceTest(parameterized.TestCase): def testdevice_get(self, fname): """Test getting MjData from a device.""" m = test_util.load_test_file(fname) + m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE # force sparse for testing mx = device.device_put(m) dx = mjx.make_data(mx) d = mujoco.MjData(m) @@ -85,6 +86,7 @@ class DeviceTest(parameterized.TestCase): def testdevice_get_batched(self, fname): """Test getting MjData from a device.""" m = test_util.load_test_file(fname) + m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE # force sparse for testing mx = device.device_put(m) batch_size = 32 diff --git a/mjx/mujoco/mjx/_src/forward.py b/mjx/mujoco/mjx/_src/forward.py index 1ff268e7..682cc88e 100644 --- a/mjx/mujoco/mjx/_src/forward.py +++ b/mjx/mujoco/mjx/_src/forward.py @@ -66,7 +66,7 @@ def fwd_position(m: Model, d: Data) -> Data: d = smooth.kinematics(m, d) d = smooth.com_pos(m, d) d = smooth.crb(m, d) - d = smooth.factor_m(m, d, d.qM) + d = smooth.factor_m(m, d) d = collision_driver.collision(m, d) d = constraint.make_constraint(m, d) d = smooth.transmission(m, d) @@ -288,8 +288,8 @@ def euler(m: Model, d: Data) -> Data: qacc = d.qacc if not m.opt.disableflags & DisableBit.EULERDAMP: # TODO(robotics-simulation): can this be done with a smaller perf hit - mh = d.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping) - dh = smooth.factor_m(m, d, mh) + dh = d.replace(qM=d.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping)) + dh = smooth.factor_m(m, dh) qfrc = d.qfrc_smooth + d.qfrc_constraint qacc = smooth.solve_m(m, dh, qfrc) return _advance(m, d, d.act_dot, qacc) diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 5b406185..dd91de94 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -23,8 +23,10 @@ import mujoco from mujoco.mjx._src import collision_driver from mujoco.mjx._src import constraint from mujoco.mjx._src import mesh +from mujoco.mjx._src import support from mujoco.mjx._src import types import numpy as np +import scipy def _put_option(o: mujoco.MjOption, device=None) -> types.Option: @@ -35,6 +37,9 @@ def _put_option(o: mujoco.MjOption, device=None) -> types.Option: if o.cone not in set(types.ConeType): raise NotImplementedError(f'{mujoco.mjtCone(o.cone)}') + if o.jacobian not in set(types.JacobianType): + raise NotImplementedError(f'{mujoco.mjtJacobian(o.jacobian)}') + if o.solver not in set(types.SolverType): raise NotImplementedError(f'{mujoco.mjtSolver(o.solver)}') @@ -49,6 +54,7 @@ def _put_option(o: mujoco.MjOption, device=None) -> types.Option: } static_fields['integrator'] = types.IntegratorType(o.integrator) static_fields['cone'] = types.ConeType(o.cone) + static_fields['jacobian'] = types.JacobianType(o.jacobian) static_fields['solver'] = types.SolverType(o.solver) static_fields['disableflags'] = types.DisableBit(o.disableflags) @@ -137,8 +143,10 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data: ne, nf, nl, nc = constraint.count_constraints(m) nefc = ne + nf + nl + nc + zero_0 = jp.zeros(0, dtype=jp.float32) zero_nv = jp.zeros(m.nv, dtype=jp.float32) zero_nv_6 = jp.zeros((m.nv, 6), dtype=jp.float32) + zero_nv_nv = jp.zeros((m.nv, m.nv), dtype=jp.float32) zero_nbody_3 = jp.zeros((m.nbody, 3), dtype=jp.float32) zero_nbody_6 = jp.zeros((m.nbody, 6), dtype=jp.float32) zero_nbody_10 = jp.zeros((m.nbody, 10), dtype=jp.float32) @@ -180,10 +188,9 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data: actuator_length=zero_nu, actuator_moment=jp.zeros((m.nu, m.nv), dtype=jp.float32), crb=zero_nbody_10, - qM=zero_nm, - qLD=zero_nm, - qLDiagInv=zero_nv, - qLDiagSqrtInv=zero_nv, + qM=zero_nm if support.is_sparse(m) else zero_nv_nv, + qLD=zero_nm if support.is_sparse(m) else zero_nv_nv, + qLDiagInv=zero_nv if support.is_sparse(m) else zero_0, contact=types.Contact.zero(ncon), efc_J=jp.zeros((nefc, m.nv), dtype=jp.float32), efc_frictionloss=zero_nefc, @@ -237,6 +244,14 @@ def get_data( mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL, ]).repeat([ne, nf, nl, nc]) + dof_i, dof_j = [], [] + for i in range(m.nv): + j = i + while j > -1: + dof_i.append(i) + dof_j.append(j) + j = m.dof_parentid[j] + ds = [] for i in range(batch_size): dx_i = jax.tree_map(lambda x, i=i: x[i], dx) if batched else d @@ -267,6 +282,15 @@ def get_data( if field.name == 'efc_J': value = value[efc_active].reshape(-1) + if field.name == 'qM' and not support.is_sparse(m): + value = value[dof_i, dof_j] + + if field.name == 'qLD' and not support.is_sparse(m): + value = value[dof_i, dof_j] + + if field.name == 'qLDiagInv' and not support.is_sparse(m): + value = np.ones(m.nv) + if value.shape: getattr(d_i, field.name)[:] = value else: @@ -346,6 +370,18 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data: value[value_beg:value_beg+size] = fields[fname][d_beg:d_beg+size] fields[fname] = value + # convert qM and qLD if jacobian is dense + if not support.is_sparse(m): + fields['qM'] = np.zeros((m.nv, m.nv)) + mujoco.mj_fullM(m, fields['qM'], d.qM) + # TODO(erikfrey): derive L*L' from L'*D*L instead of recomputing + try: + fields['qLD'], _ = scipy.linalg.cho_factor(fields['qM']) + except scipy.linalg.LinAlgError: + # this happens when qM is empty or unstable simulation + fields['qLD'] = np.zeros((m.nv, m.nv)) + fields['qLDiagInv'] = np.zeros(0) + fields = jax.device_put(fields, device=device) fields['contact'] = _put_contact(d.contact, ncon, device=device) diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index 8b25f359..b402c407 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -25,7 +25,7 @@ import numpy as np _MULTIPLE_CONVEX_OBJECTS = """ -