diff --git a/doc/changelog.rst b/doc/changelog.rst index 677cdbcf..bf1cbe34 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -27,6 +27,9 @@ MJX ^^^ - Added ``site_xpos`` and ``site_xmat`` to MJX. +- Added ``put_data``, ``put_model``, ``get_data`` to replace ``device_put`` and ``device_get_into``, which will be + deprecated. These new functions correctly translate fields that are the result of intermediate calculations such as + ``efc_J``. Bug fixes ^^^^^^^^^ diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py index 3ec6312a..1c8d0661 100644 --- a/mjx/mujoco/mjx/__init__.py +++ b/mjx/mujoco/mjx/__init__.py @@ -21,7 +21,10 @@ from mujoco.mjx._src.device import device_get_into from mujoco.mjx._src.device import device_put from mujoco.mjx._src.forward import forward from mujoco.mjx._src.forward import step +from mujoco.mjx._src.io import get_data from mujoco.mjx._src.io import make_data +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.smooth import com_pos from mujoco.mjx._src.smooth import com_vel diff --git a/mjx/mujoco/mjx/_src/collision_driver.py b/mjx/mujoco/mjx/_src/collision_driver.py index dde3fe92..c54278b2 100644 --- a/mjx/mujoco/mjx/_src/collision_driver.py +++ b/mjx/mujoco/mjx/_src/collision_driver.py @@ -40,7 +40,6 @@ from mujoco.mjx._src.types import DisableBit from mujoco.mjx._src.types import GeomType from mujoco.mjx._src.types import Model # pylint: enable=g-importing-member -import numpy as np # pair-wise collision functions @@ -81,6 +80,12 @@ def _add_candidate( if t1 > t2: t1, t2, g1, g2 = t2, t1, g2, g1 + # MuJoCo does not collide planes with other planes or hfields + if t1 == GeomType.PLANE and t2 == GeomType.PLANE: + return + if t1 == GeomType.PLANE and t2 == GeomType.HFIELD: + return + def mesh_key(i): convex_data = [[None] * m.ngeom] * 3 if isinstance(m, Model): @@ -284,13 +289,11 @@ def _collide_geoms( solimp=params.solimp, geom1=geom1, geom2=geom2, - dim=np.array([]), - efc_address=np.array([]), ) return con -def _max_contact_points(m: Model) -> int: +def _max_contact_points(m: Union[Model, mujoco.MjModel]) -> int: """Returns the maximum number of contact points when set as a numeric.""" for i in range(m.nnumeric): name = m.names[m.name_numericadr[i] :].decode('utf-8').split('\x00', 1)[0] @@ -334,7 +337,7 @@ def collision_candidates(m: Union[Model, mujoco.MjModel]) -> CandidateSet: return candidate_set -def ncon(m: Model) -> int: +def ncon(m: Union[Model, mujoco.MjModel]) -> int: """Returns the number of contacts computed in MJX given a model.""" if m.opt.disableflags & DisableBit.CONTACT: return 0 @@ -354,9 +357,8 @@ def ncon(m: Model) -> int: def collision(m: Model, d: Data) -> Data: """Collides geometries.""" - ncon_ = ncon(m) - if ncon_ == 0: - return d.replace(contact=Contact.zero(), ncon=0) + if ncon(m) == 0: + return d.replace(contact=Contact.zero()) candidate_set = collision_candidates(m) @@ -376,13 +378,4 @@ def collision(m: Model, d: Data) -> Data: _, idx = jax.lax.top_k(-contact.dist, k=max_contact_points) contact = jax.tree_map(lambda x, idx=idx: jp.take(x, idx, axis=0), contact) - if ncon_ != contact.dist.shape[0]: - raise RuntimeError('Number of contacts does not match ncon.') - - # TODO(robotics-simulation): move this logic to device_put - ns = d.ne + d.nf + d.nl - contact = contact.replace(efc_address=np.arange(ns, ns + ncon_ * 4, 4)) - # TODO(robotics-simulation): add support for other friction dimensions - contact = contact.replace(dim=3 * np.ones(ncon_, dtype=np.int32)) - - return d.replace(contact=contact, ncon=ncon_) + return d.replace(contact=contact) diff --git a/mjx/mujoco/mjx/_src/collision_driver_test.py b/mjx/mujoco/mjx/_src/collision_driver_test.py index b735e014..afb7bd94 100644 --- a/mjx/mujoco/mjx/_src/collision_driver_test.py +++ b/mjx/mujoco/mjx/_src/collision_driver_test.py @@ -253,7 +253,6 @@ class CapsuleCollisionTest(parameterized.TestCase): self.assertGreater(c.dist[1], 0) # extract the contact point with penetration c = jax.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact) - c = c.replace(dim=c.dim[np.array([0])]) for field in dataclasses.fields(Contact): _assert_attr_eq(c, d.contact, field.name, 'capsule_convex_edge', 1e-4) @@ -281,7 +280,6 @@ class ConvexTest(absltest.TestCase): np.testing.assert_array_less(-dx.contact.dist[2:], 0) # extract the contact points with penetration c = jax.tree_map(lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx.contact) - c = c.replace(dim=c.dim[np.array([0, 1])]) for field in dataclasses.fields(Contact): _assert_attr_eq(c, d.contact, field.name, 'box_plane', 1e-2) @@ -339,7 +337,6 @@ class ConvexTest(absltest.TestCase): np.testing.assert_array_less(-dx.contact.dist[1:], 0) # extract the contact point with penetration c = jax.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact) - c = c.replace(dim=c.dim[np.array([0])]) for field in dataclasses.fields(Contact): _assert_attr_eq(c, d.contact, field.name, 'box_box_edge', 1e-2) @@ -517,8 +514,8 @@ class TopKContactTest(absltest.TestCase): dx_all = collision_jit_fn(mx_all, dx) dx_top_k = collision_jit_fn(mx_top_k, dx) - self.assertEqual(dx_all.ncon, 3) - self.assertEqual(dx_top_k.ncon, 2) + self.assertEqual(dx_all.contact.dist.shape, (3,)) + self.assertEqual(dx_top_k.contact.dist.shape, (2,)) if __name__ == '__main__': diff --git a/mjx/mujoco/mjx/_src/constraint.py b/mjx/mujoco/mjx/_src/constraint.py index 313b43ff..50efcad3 100644 --- a/mjx/mujoco/mjx/_src/constraint.py +++ b/mjx/mujoco/mjx/_src/constraint.py @@ -14,11 +14,12 @@ # ============================================================================== """Core non-smooth constraint functions.""" -from typing import Optional, Tuple +from typing import Optional, Tuple, Union import jax from jax import numpy as jp import mujoco +from mujoco.mjx._src import collision_driver from mujoco.mjx._src import math from mujoco.mjx._src import support # pylint: disable=g-importing-member @@ -276,7 +277,7 @@ def _instantiate_limit_slide_hinge(m: Model, d: Data) -> Optional[_Efc]: def _instantiate_contact(m: Model, d: Data) -> Optional[_Efc]: """Calculates constraint rows for contacts.""" - if (m.opt.disableflags & DisableBit.CONTACT) or d.ncon == 0: + if collision_driver.ncon(m) == 0: return None @jax.vmap @@ -313,7 +314,9 @@ def _instantiate_contact(m: Model, d: Data) -> Optional[_Efc]: return _Efc(j, pos, pos, invweight, solref, solimp, frictionloss) -def count_constraints(m: Model, d: Data) -> Tuple[int, int, int, int]: +def count_constraints( + m: Union[Model, mujoco.MjModel] +) -> Tuple[int, int, int, int]: """Returns equality, friction, limit, and contact constraint counts.""" if m.opt.disableflags & DisableBit.CONSTRAINT: return 0, 0, 0, 0 @@ -333,10 +336,7 @@ def count_constraints(m: Model, d: Data) -> Tuple[int, int, int, int]: else: nl = int(m.jnt_limited.sum()) - if m.opt.disableflags & DisableBit.CONTACT: - nc = 0 - else: - nc = d.ncon * 4 + nc = collision_driver.ncon(m) * 4 return ne, nf, nl, nc @@ -344,10 +344,6 @@ def count_constraints(m: Model, d: Data) -> Tuple[int, int, int, int]: def make_constraint(m: Model, d: Data) -> Data: """Creates constraint jacobians and other supporting data.""" - ns = sum(count_constraints(m, d)[:-1]) - # TODO(robotics-simulation): make device_put set nefc/efc_address instead - d = d.tree_replace({'contact.efc_address': np.arange(ns, ns + d.ncon * 4, 4)}) - if m.opt.disableflags & DisableBit.CONSTRAINT: efcs = () else: @@ -364,7 +360,7 @@ def make_constraint(m: Model, d: Data) -> Data: if not efcs: z = jp.empty(0) d = d.replace(efc_J=jp.empty((0, m.nv))) - d = d.replace(efc_D=z, efc_aref=z, efc_frictionloss=z, nefc=0) + d = d.replace(efc_D=z, efc_aref=z, efc_frictionloss=z) return d efc = jax.tree_map(lambda *x: jp.concatenate(x), *efcs) @@ -378,6 +374,6 @@ def make_constraint(m: Model, d: Data) -> Data: aref, r = fn(efc) d = d.replace(efc_J=efc.J, efc_D=1 / r, efc_aref=aref) - d = d.replace(efc_frictionloss=efc.frictionloss, nefc=r.shape[0]) + d = d.replace(efc_frictionloss=efc.frictionloss) return d diff --git a/mjx/mujoco/mjx/_src/constraint_test.py b/mjx/mujoco/mjx/_src/constraint_test.py index 8ce1bfa6..0510e1df 100644 --- a/mjx/mujoco/mjx/_src/constraint_test.py +++ b/mjx/mujoco/mjx/_src/constraint_test.py @@ -136,13 +136,8 @@ class ConstraintTest(parameterized.TestCase): m = test_util.load_test_file('ant.xml') d = mujoco.MjData(m) - m.opt.disableflags = m.opt.disableflags & ~DisableBit.CONSTRAINT - mx, dx = mjx.device_put(m), mjx.device_put(d) - dx = constraint.make_constraint(mx, dx) - self.assertGreater(dx.efc_J.shape[0], 1) - m.opt.disableflags = m.opt.disableflags | DisableBit.CONSTRAINT - mx = mjx.device_put(m) + mx, dx = mjx.device_put(m), mjx.device_put(d) dx = constraint.make_constraint(mx, dx) self.assertEqual(dx.efc_J.shape[0], 0) diff --git a/mjx/mujoco/mjx/_src/dataclasses.py b/mjx/mujoco/mjx/_src/dataclasses.py index 0936bec6..b939ab9e 100644 --- a/mjx/mujoco/mjx/_src/dataclasses.py +++ b/mjx/mujoco/mjx/_src/dataclasses.py @@ -18,7 +18,7 @@ import copy import dataclasses import typing -from typing import Dict, Optional, Sequence, TypeVar +from typing import Any, Dict, Optional, Sequence, TypeVar import jax import numpy as np @@ -113,6 +113,10 @@ class PyTreeNode: # stub for pytype raise NotImplementedError + @classmethod + def fields(cls) -> tuple[dataclasses.Field[Any], ...]: + return dataclasses.fields(cls) + def tree_replace( self, params: Dict[str, Optional[jax.typing.ArrayLike]] ) -> 'PyTreeNode': diff --git a/mjx/mujoco/mjx/_src/device.py b/mjx/mujoco/mjx/_src/device.py index ffd687b8..327b89de 100644 --- a/mjx/mujoco/mjx/_src/device.py +++ b/mjx/mujoco/mjx/_src/device.py @@ -25,6 +25,7 @@ import mujoco from mujoco.mjx._src import collision_driver from mujoco.mjx._src import mesh from mujoco.mjx._src import types +import numpy as np _MJ_TYPE_ATTR = { mujoco.mjtBias: (mujoco.MjModel.actuator_biastype,), @@ -68,7 +69,6 @@ _TRANSFORMS = { (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 if x is not None and x.shape[0] else jp.zeros((0, 3, 3)) @@ -265,9 +265,16 @@ def device_get_into(result, value): else: if isinstance(result, mujoco.MjData): + ncon = value.contact.dist.shape[0] + nefc = value.efc_J.shape[0] mujoco._functions._realloc_con_efc( # pylint: disable=protected-access - result, ncon=value.ncon, nefc=value.nefc + result, ncon=ncon, nefc=nefc ) + result.ncon = ncon + result.nefc = nefc + efc_start = nefc - ncon * 4 + result.contact.efc_address[:] = np.arange(efc_start, nefc, 4) + result.contact.dim[:] = 3 for f in dataclasses.fields(value): # type: ignore if (type(value), f.name) in _DERIVED: diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 55381f41..060ffa50 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -14,94 +14,337 @@ # ============================================================================== """Functions to initialize, load, or save data.""" +import copy +from typing import List, Union + +import jax from jax import numpy as jp +import mujoco from mujoco.mjx._src import collision_driver from mujoco.mjx._src import constraint -# pylint: disable=g-importing-member -from mujoco.mjx._src.types import Contact -from mujoco.mjx._src.types import Data -from mujoco.mjx._src.types import Model -# pylint: enable=g-importing-member +from mujoco.mjx._src import mesh +from mujoco.mjx._src import types import numpy as np -def make_data(m: Model) -> Data: +def _put_option(o: mujoco.MjOption, device=None) -> types.Option: + """Puts mujoco.MjOption onto a device, resulting in mjx.Option.""" + if o.integrator not in set(types.IntegratorType): + raise NotImplementedError(f'{mujoco.mjtIntegrator(o.integrator)}') + + if o.cone not in set(types.ConeType): + raise NotImplementedError(f'{mujoco.mjtCone(o.cone)}') + + if o.solver not in set(types.SolverType): + raise NotImplementedError(f'{mujoco.mjtSolver(o.solver)}') + + for i in range(mujoco.mjtEnableBit.mjNENABLE): + if o.enableflags & 2**i: + raise NotImplementedError(f'{mujoco.mjtEnableBit(2 ** i)}') + + static_fields = { + f.name: copy.copy(getattr(o, f.name)) + for f in types.Option.fields() + if f.type in (int, bytes, np.ndarray) + } + static_fields['integrator'] = types.IntegratorType(o.integrator) + static_fields['cone'] = types.ConeType(o.cone) + static_fields['solver'] = types.SolverType(o.solver) + static_fields['disableflags'] = types.DisableBit(o.disableflags) + + device_fields = { + f.name: copy.copy(getattr(o, f.name)) + for f in types.Option.fields() + if f.type is jax.Array + } + device_fields = jax.device_put(device_fields, device=device) + + has_fluid_params = o.density > 0 or o.viscosity > 0 or o.wind.any() + + return types.Option( + has_fluid_params=has_fluid_params, + **static_fields, + **device_fields, + ) + + +def _put_statistic(s: mujoco.MjStatistic, device=None) -> types.Statistic: + """Puts mujoco.MjStatistic onto a device, resulting in mjx.Statistic.""" + return types.Statistic( + meaninertia=jax.device_put(s.meaninertia, device=device) + ) + + +def put_model(m: mujoco.MjModel, device=None) -> types.Model: + """Puts mujoco.MjModel onto a device, resulting in mjx.Model.""" + if m.ntendon: + raise NotImplementedError('tendons are not supported') + + if (m.geom_condim != 3).any() or (m.pair_dim != 3).any(): + raise NotImplementedError('only condim=3 is supported') + + # check collision geom types + for g1, g2, *_ in collision_driver.collision_candidates(m): + if collision_driver.get_collision_fn((g1, g2)) is None: + g1, g2 = mujoco.mjtGeom(g1), mujoco.mjtGeom(g2) + raise NotImplementedError(f'({g1}, {g2}) has no collision function') + + for enum_field, enum_type, mj_type in ( + (m.actuator_biastype, types.BiasType, mujoco.mjtBias), + (m.actuator_dyntype, types.DynType, mujoco.mjtDyn), + (m.actuator_gaintype, types.GainType, mujoco.mjtGain), + (m.actuator_trntype, types.TrnType, mujoco.mjtTrn), + (m.eq_type, types.EqType, mujoco.mjtEq), + ): + missing = set(enum_field) - set(enum_type) + if missing: + raise NotImplementedError( + f'{[mj_type(m) for m in missing]} not supported' + ) + + opt = _put_option(m.opt, device=device) + stat = _put_statistic(m.stat, device=device) + + static_fields = { + f.name: getattr(m, f.name) + for f in types.Model.fields() + if f.type in (int, bytes, np.ndarray) + } + + device_fields = { + f.name: copy.copy(getattr(m, f.name)) # copy because device_put is async + for f in types.Model.fields() + if f.type is jax.Array + } + device_fields.update(mesh.get(m)) + device_fields = jax.device_put(device_fields, device=device) + + return types.Model( + opt=opt, + stat=stat, + **static_fields, + **device_fields, + ) + + +def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data: """Allocate and initialize Data.""" + ncon = collision_driver.ncon(m) + ne, nf, nl, nc = constraint.count_constraints(m) + nefc = ne + nf + nl + nc + + zero_nv = jp.zeros(m.nv, dtype=jp.float32) + zero_nv_6 = jp.zeros((m.nv, 6), 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) + zero_nbody_3_3 = jp.zeros((m.nbody, 3, 3), dtype=jp.float32) + zero_nefc = jp.zeros(nefc, dtype=jp.float32) + zero_na = jp.zeros(m.na, dtype=jp.float32) + zero_nu = jp.zeros(m.nu, dtype=jp.float32) + zero_njnt_3 = jp.zeros((m.njnt, 3), dtype=jp.float32) + zero_nm = jp.zeros(m.nM, dtype=jp.float32) + # create first d to get num contacts and nc - d = Data( + d = types.Data( solver_niter=jp.array(0, dtype=jp.int32), - ne=0, - nf=0, - nl=0, - nefc=0, - ncon=0, - time=jp.zeros((), dtype=jp.float32), + time=jp.array(0.0), qpos=m.qpos0, - qvel=jp.zeros(m.nv, dtype=jp.float32), - act=jp.zeros(m.na, dtype=jp.float32), - qacc_warmstart=jp.zeros(m.nv, dtype=jp.float32), - ctrl=jp.zeros(m.nu, dtype=jp.float32), - qfrc_applied=jp.zeros(m.nv, dtype=jp.float32), - xfrc_applied=jp.zeros((m.nbody, 6), dtype=jp.float32), + qvel=zero_nv, + act=zero_na, + qacc_warmstart=zero_nv, + ctrl=zero_nu, + qfrc_applied=zero_nv, + xfrc_applied=zero_nbody_6, eq_active=jp.zeros(m.neq, dtype=jp.int32), - qacc=jp.zeros(m.nv, dtype=jp.float32), - act_dot=jp.zeros(m.na, dtype=jp.float32), - xpos=jp.zeros((m.nbody, 3), dtype=jp.float32), + qacc=zero_nv, + act_dot=zero_na, + xpos=zero_nbody_3, xquat=jp.zeros((m.nbody, 4), dtype=jp.float32), - xmat=jp.zeros((m.nbody, 3, 3), dtype=jp.float32), - xipos=jp.zeros((m.nbody, 3), dtype=jp.float32), - ximat=jp.zeros((m.nbody, 3, 3), dtype=jp.float32), - xanchor=jp.zeros((m.njnt, 3), dtype=jp.float32), - xaxis=jp.zeros((m.njnt, 3), dtype=jp.float32), + xmat=zero_nbody_3_3, + xipos=zero_nbody_3, + ximat=zero_nbody_3_3, + xanchor=zero_njnt_3, + xaxis=zero_njnt_3, 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), - actuator_length=jp.zeros(m.nu, dtype=jp.float32), + subtree_com=zero_nbody_3, + cdof=zero_nv_6, + cinert=zero_nbody_10, + actuator_length=zero_nu, actuator_moment=jp.zeros((m.nu, m.nv), dtype=jp.float32), - crb=jp.zeros((m.nbody, 10), dtype=jp.float32), - qM=jp.zeros(m.nM, dtype=jp.float32), - qLD=jp.zeros(m.nM, dtype=jp.float32), - qLDiagInv=jp.zeros(m.nv, dtype=jp.float32), - qLDiagSqrtInv=jp.zeros(m.nv, dtype=jp.float32), - contact=Contact.zero(), - efc_J=jp.zeros((), dtype=jp.float32), - efc_frictionloss=jp.zeros((), dtype=jp.float32), - efc_D=jp.zeros((), dtype=jp.float32), - actuator_velocity=jp.zeros(m.nu, dtype=jp.float32), - cvel=jp.zeros((m.nbody, 6), dtype=jp.float32), - cdof_dot=jp.zeros((m.nv, 6), dtype=jp.float32), - qfrc_bias=jp.zeros(m.nv, dtype=jp.float32), - qfrc_passive=jp.zeros(m.nv, dtype=jp.float32), - efc_aref=jp.zeros((), dtype=jp.float32), - actuator_force=jp.zeros(m.nu, dtype=jp.float32), - qfrc_actuator=jp.zeros(m.nv, dtype=jp.float32), - qfrc_smooth=jp.zeros(m.nv, dtype=jp.float32), - qacc_smooth=jp.zeros(m.nv, dtype=jp.float32), - qfrc_constraint=jp.zeros(m.nv, dtype=jp.float32), - qfrc_inverse=jp.zeros(m.nv, dtype=jp.float32), - efc_force=jp.zeros((), dtype=jp.float32), - ) - - # get contact data with correct shapes - ncon = collision_driver.ncon(m) - d = d.replace(contact=Contact.zero((ncon,)), ncon=ncon) - d = d.tree_replace({'contact.dim': 3 * np.ones(ncon)}) - - ne, nf, nl, nc = constraint.count_constraints(m, d) - d = d.replace(ne=ne, nf=nf, nl=nl, nefc=ne + nf + nl + nc) - ns = ne + nf + nl - d = d.tree_replace({'contact.efc_address': np.arange(ns, ns + ncon * 4, 4)}) - d = d.replace( - efc_J=jp.zeros((d.nefc, m.nv), dtype=jp.float32), - efc_frictionloss=jp.zeros(d.nefc, dtype=jp.float32), - efc_D=jp.zeros(d.nefc, dtype=jp.float32), - efc_aref=jp.zeros(d.nefc, dtype=jp.float32), - efc_force=jp.zeros(d.nefc, dtype=jp.float32), + crb=zero_nbody_10, + qM=zero_nm, + qLD=zero_nm, + qLDiagInv=zero_nv, + qLDiagSqrtInv=zero_nv, + contact=types.Contact.zero(ncon), + efc_J=jp.zeros((nefc, m.nv), dtype=jp.float32), + efc_frictionloss=zero_nefc, + efc_D=zero_nefc, + actuator_velocity=zero_nu, + cvel=zero_nbody_6, + cdof_dot=zero_nv_6, + qfrc_bias=zero_nv, + qfrc_passive=zero_nv, + efc_aref=zero_nefc, + actuator_force=zero_nu, + qfrc_actuator=zero_nv, + qfrc_smooth=zero_nv, + qacc_smooth=zero_nv, + qfrc_constraint=zero_nv, + qfrc_inverse=zero_nv, + efc_force=zero_nefc, ) return d + + +def _get_contact( + c: mujoco._structs._MjContactList, + cx: types.Contact, + efc_start: int, +): + """Converts mjx.Contact to mujoco._structs._MjContactList.""" + con_id = np.nonzero(cx.dist <= 0)[0] + for field in types.Contact.fields(): + value = getattr(cx, field.name)[con_id] + if field.name == 'frame': + value = value.reshape((-1, 9)) + getattr(c, field.name)[:] = value + + ncon = con_id.shape[0] + c.efc_address[:] = np.arange(efc_start, efc_start + ncon * 4, 4)[con_id] + + +def get_data( + m: mujoco.MjModel, d: types.Data +) -> Union[mujoco.MjData, List[mujoco.MjData]]: + """Gets mjx.Data from a device, resulting in mujoco.MjData or List[MjData].""" + dx = jax.device_get(d) + batched = len(d.qpos.shape) > 1 + batch_size = d.qpos.shape[0] if batched else 1 + ne, nf, nl, nc = constraint.count_constraints(m) + efc_type = np.array([ + mujoco.mjtConstraint.mjCNSTR_EQUALITY, + mujoco.mjtConstraint.mjCNSTR_FRICTION_DOF, + mujoco.mjtConstraint.mjCNSTR_LIMIT_JOINT, + mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL, + ]).repeat([ne, nf, nl, nc]) + + ds = [] + for i in range(batch_size): + dx_i = jax.tree_map(lambda x, i=i: x[i], dx) if batched else d + ncon = (dx_i.contact.dist <= 0).sum() + efc_active = (dx_i.efc_J != 0).any(axis=1) + efc_con = efc_type == mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL + nefc, nc = efc_active.sum(), (efc_active & efc_con).sum() + d_i = mujoco.MjData(m) + d_i.nnzJ = nefc * m.nv + mujoco._functions._realloc_con_efc(d_i, ncon=ncon, nefc=nefc) # pylint: disable=protected-access + d_i.efc_J_rownnz[:] = np.repeat(m.nv, nefc) + d_i.efc_J_rowadr[:] = np.arange(0, nefc * m.nv, m.nv) + d_i.efc_J_colind[:] = np.tile(np.arange(m.nv), nefc) + + for field in types.Data.fields(): + if field.name == 'contact': + _get_contact(d_i.contact, dx_i.contact, nefc - nc) + continue + + value = getattr(dx_i, field.name) + + if field.name in ('xmat', 'ximat', 'geom_xmat'): + value = value.reshape((-1, 9)) + + if field.name in ('efc_frictionloss', 'efc_D', 'efc_aref', 'efc_force'): + value = value[efc_active] + + if field.name == 'efc_J': + value = value[efc_active].reshape(-1) + + if value.shape: + getattr(d_i, field.name)[:] = value + else: + setattr(d_i, field.name, value) + + d_i.efc_type[:] = efc_type[efc_active] + ds.append(d_i) + + return ds if batched else ds[0] + + +def _put_contact( + c: mujoco._structs._MjContactList, ncon: int, device=None +) -> types.Contact: + """Puts mujoco.structs._MjContactList onto a device, resulting in mjx.Contact.""" + fields = { + f.name: copy.copy(getattr(c, f.name)) for f in types.Contact.fields() + } + fields['frame'] = fields['frame'].reshape((-1, 3, 3)) + pad_size = ncon - c.dist.shape[0] + pad_fn = lambda x: np.concatenate( + (x, np.zeros((pad_size,) + x.shape[1:], dtype=x.dtype)) + ) + fields = jax.tree_map(pad_fn, fields) + fields['dist'][-pad_size:] = np.inf + fields = jax.device_put(fields, device=device) + + return types.Contact(**fields) + + +def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data: + """Puts mujoco.MjData onto a device, resulting in mjx.Data.""" + ncon = collision_driver.ncon(m) + ne, nf, nl, nc = constraint.count_constraints(m) + nefc = ne + nf + nl + nc + + for d_val, val, name in ( + (d.ncon, ncon, 'ncon'), + (d.ne, ne, 'ne'), + (d.nf, nf, 'nf'), + (d.nl, nl, 'nl'), + (d.nefc, nefc, 'nefc'), + ): + if d_val > val: + raise ValueError(f'd.{name} too high, d.{name} = {d_val}, model = {val}') + + fields = { + f.name: copy.copy(getattr(d, f.name)) # copy because device_put is async + for f in types.Data.fields() + if f.type is jax.Array + } + + for fname in ('xmat', 'ximat', 'geom_xmat'): + fields[fname] = fields[fname].reshape((-1, 3, 3)) + + # pad efc fields: MuJoCo efc arrays are sparse for inactive constraints. + # efc_J is also optionally column-sparse (typically for large nv). MJX is + # neither: it contains zeros for inactive constraints, and efc_J is always + # (nefc, nv). this may change in the future. + if mujoco.mj_isSparse(m): + nr = d.efc_J_rownnz.shape[0] + efc_j = np.zeros((nr, m.nv)) + for i in range(nr): + rowadr = d.efc_J_rowadr[i] + for j in range(d.efc_J_rownnz[i]): + efc_j[i, d.efc_J_colind[rowadr + j]] = fields['efc_J'][rowadr + j] + fields['efc_J'] = efc_j + else: + fields['efc_J'] = fields['efc_J'].reshape((-1, m.nv)) + + for fname in ('efc_J', 'efc_frictionloss', 'efc_D', 'efc_aref', 'efc_force'): + value = np.zeros((nefc, m.nv)) if fname == 'efc_J' else np.zeros(nefc) + for i in range(4): + value_beg = sum([ne, nf, nl][:i]) + d_beg = sum([d.ne, d.nf, d.nl][:i]) + size = [d.ne, d.nf, d.nl, d.nefc - d.nl - d.nf - d.ne][i] + value[value_beg:value_beg+size] = fields[fname][d_beg:d_beg+size] + fields[fname] = value + + fields = jax.device_put(fields, device=device) + fields['contact'] = _put_contact(d.contact, ncon, device=device) + + return types.Data(**fields) diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index 269509a8..4c7b0065 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -17,26 +17,390 @@ from absl.testing import absltest from absl.testing import parameterized import jax +from jax import numpy as jp +import mujoco from mujoco import mjx -from mujoco.mjx._src import test_util +import numpy as np + + +_MULTIPLE_CONVEX_OBJECTS = """ + + +""" + +_MULTIPLE_CONSTRAINTS = """ + + + + + + + + + + + + + + + + + +""" class IoTest(parameterized.TestCase): - @parameterized.parameters(test_util.TEST_FILES) - def test_make_data(self, fname): - """Test that data created by make_data matches data returned by step.""" + def test_put_model(self): + m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS) + mx = mjx.put_model(m) + self.assertEqual(mx.nq, m.nq) + self.assertEqual(mx.nv, m.nv) + self.assertEqual(mx.nu, m.nu) + self.assertEqual(mx.na, m.na) + self.assertEqual(mx.nbody, m.nbody) + self.assertEqual(mx.njnt, m.njnt) + self.assertEqual(mx.ngeom, m.ngeom) + self.assertEqual(mx.nmesh, m.nmesh) + self.assertEqual(mx.npair, m.npair) + self.assertEqual(mx.nexclude, m.nexclude) + self.assertEqual(mx.neq, m.neq) + self.assertEqual(mx.nnumeric, m.nnumeric) + self.assertEqual(mx.nM, m.nM) + self.assertAlmostEqual(mx.opt.timestep, m.opt.timestep) - m = test_util.load_test_file(fname) - mx = mjx.device_put(m) - dx = mjx.make_data(mx) - dx_step = mjx.step(mx, dx) + np.testing.assert_allclose(mx.body_parentid, m.body_parentid) + np.testing.assert_allclose(mx.geom_type, m.geom_type) + np.testing.assert_allclose(mx.geom_bodyid, m.geom_bodyid) + np.testing.assert_almost_equal(mx.geom_solref, m.geom_solref) + np.testing.assert_almost_equal(mx.geom_pos, m.geom_pos) + self.assertLen(mx.geom_convex_face, 6) + self.assertLen(mx.geom_convex_vert, 6) + self.assertLen(mx.geom_convex_edge, 6) + self.assertLen(mx.geom_convex_facenormal, 6) - _, dx_treedef = jax.tree_util.tree_flatten(dx) - _, dx_step_treedef = jax.tree_util.tree_flatten(dx_step) + np.testing.assert_allclose(mx.jnt_type, m.jnt_type) + np.testing.assert_allclose(mx.jnt_dofadr, m.jnt_dofadr) + np.testing.assert_allclose(mx.jnt_bodyid, m.jnt_bodyid) + np.testing.assert_allclose(mx.jnt_limited, m.jnt_limited) + np.testing.assert_almost_equal(mx.jnt_axis, m.jnt_axis) - self.assertEqual(dx_treedef, dx_step_treedef) + np.testing.assert_allclose(mx.actuator_trntype, m.actuator_trntype) + np.testing.assert_allclose(mx.actuator_dyntype, m.actuator_dyntype) + np.testing.assert_allclose(mx.actuator_gaintype, m.actuator_gaintype) + np.testing.assert_allclose(mx.actuator_biastype, m.actuator_biastype) + np.testing.assert_allclose(mx.actuator_trnid, m.actuator_trnid) + def test_fluid_params(self): + """Test that has_fluid_params is set when fluid params are present.""" + m = mjx.put_model( + mujoco.MjModel.from_xml_string( + '' + ) + ) + self.assertTrue(m.opt.has_fluid_params) + + def test_put_model_implicit_not_implemented(self): + """Test that MJX guards against models with unimplemented features.""" + + with self.assertRaises(NotImplementedError): + mjx.put_model( + mujoco.MjModel.from_xml_string( + '' + ) + ) + + def test_put_model_cone_not_implemented(self): + with self.assertRaises(NotImplementedError): + mjx.put_model( + mujoco.MjModel.from_xml_string( + '' + ) + ) + + def test_put_model_pgs_not_implemented(self): + with self.assertRaises(NotImplementedError): + mjx.put_model( + mujoco.MjModel.from_xml_string( + '' + ) + ) + + def test_put_model_site_actuator_not_implemented(self): + with self.assertRaises(NotImplementedError): + mjx.put_model(mujoco.MjModel.from_xml_string(""" + + + + + + + + + + + """)) + + def test_put_model_tendon_not_implemented(self): + with self.assertRaises(NotImplementedError): + mjx.put_model(mujoco.MjModel.from_xml_string(""" + + + + + + + + + + + + + """)) + + def test_put_model_condim_not_implemented(self): + with self.assertRaises(NotImplementedError): + mjx.put_model(mujoco.MjModel.from_xml_string(""" + + + + + + + + + + + + """)) + + def test_put_model_cylinder_not_implemented(self): + with self.assertRaises(NotImplementedError): + mjx.put_model(mujoco.MjModel.from_xml_string(""" + + + + + + + + + + + + """)) + + def test_make_data(self): + """Test that make_data returns the correct shapes.""" + + m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS) + d = mjx.make_data(m) + + nq = 22 + nbody = 5 + ncon = 46 + nv = 19 + nefc = 185 + nm = 64 + + self.assertEqual(d.qpos.shape, (nq,)) + self.assertEqual(d.qvel.shape, (nv,)) + self.assertEqual(d.act.shape, (0,)) + self.assertEqual(d.qacc_warmstart.shape, (nv,)) + self.assertEqual(d.ctrl.shape, (1,)) + self.assertEqual(d.qfrc_applied.shape, (nv,)) + self.assertEqual(d.xfrc_applied.shape, (nbody, 6)) + self.assertEqual(d.eq_active.shape, (0,)) + self.assertEqual(d.qacc.shape, (nv,)) + self.assertEqual(d.act_dot.shape, (0,)) + self.assertEqual(d.xpos.shape, (nbody, 3)) + self.assertEqual(d.xquat.shape, (nbody, 4)) + self.assertEqual(d.xmat.shape, (nbody, 3, 3)) + self.assertEqual(d.xipos.shape, (nbody, 3)) + self.assertEqual(d.ximat.shape, (nbody, 3, 3)) + self.assertEqual(d.xanchor.shape, (4, 3)) + self.assertEqual(d.xaxis.shape, (4, 3)) + self.assertEqual(d.geom_xpos.shape, (6, 3)) + self.assertEqual(d.geom_xmat.shape, (6, 3, 3)) + self.assertEqual(d.subtree_com.shape, (nbody, 3)) + self.assertEqual(d.cdof.shape, (nv, 6)) + self.assertEqual(d.cinert.shape, (nbody, 10)) + self.assertEqual(d.crb.shape, (nbody, 10)) + self.assertEqual(d.actuator_length.shape, (1,)) + self.assertEqual(d.actuator_moment.shape, (1, nv)) + self.assertEqual(d.qM.shape, (nm,)) + self.assertEqual(d.qLD.shape, (nm,)) + self.assertEqual(d.qLDiagInv.shape, (nv,)) + self.assertEqual(d.qLDiagSqrtInv.shape, (nv,)) + self.assertEqual(d.contact.dist.shape, (ncon,)) + self.assertEqual(d.contact.pos.shape, (ncon, 3)) + self.assertEqual(d.contact.frame.shape, (ncon, 3, 3)) + self.assertEqual(d.contact.solref.shape, (ncon, 2)) + self.assertEqual(d.contact.solimp.shape, (ncon, 5)) + self.assertEqual(d.contact.geom1.shape, (ncon,)) + self.assertEqual(d.contact.geom2.shape, (ncon,)) + self.assertEqual(d.efc_J.shape, (nefc, nv)) + self.assertEqual(d.efc_frictionloss.shape, (nefc,)) + self.assertEqual(d.efc_D.shape, (nefc,)) + self.assertEqual(d.actuator_velocity.shape, (1,)) + self.assertEqual(d.cvel.shape, (nbody, 6)) + self.assertEqual(d.cdof_dot.shape, (nv, 6)) + self.assertEqual(d.qfrc_bias.shape, (nv,)) + self.assertEqual(d.qfrc_passive.shape, (nv,)) + self.assertEqual(d.efc_aref.shape, (nefc,)) + self.assertEqual(d.actuator_force.shape, (1,)) + self.assertEqual(d.qfrc_actuator.shape, (nv,)) + self.assertEqual(d.qfrc_smooth.shape, (nv,)) + self.assertEqual(d.qacc_smooth.shape, (nv,)) + self.assertEqual(d.qfrc_constraint.shape, (nv,)) + self.assertEqual(d.qfrc_inverse.shape, (nv,)) + self.assertEqual(d.efc_force.shape, (nefc,)) + + def test_put_data(self): + """Test that put_data puts the correct data for dense and sparse.""" + + m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) + d = mujoco.MjData(m) + mujoco.mj_step(m, d, 2) + dx = mjx.put_data(m, d) + + # check a few fields + np.testing.assert_allclose(dx.qpos, d.qpos) + np.testing.assert_allclose(dx.xpos, d.xpos) + np.testing.assert_allclose(dx.cvel, d.cvel) + np.testing.assert_allclose(dx.cdof_dot, d.cdof_dot) + np.testing.assert_allclose(dx.qM, d.qM) + + # 4 contacts, 2 for each capsule against the plane + self.assertEqual(dx.contact.dist.shape, (4,)) + self.assertEqual(d.ncon, 1) # however only 1 contact in this step + np.testing.assert_allclose(dx.contact.dist[0], d.contact.dist[0]) + self.assertTrue(np.isinf(dx.contact.dist[1:]).all()) + self.assertEqual(dx.contact.frame.shape, (4, 3, 3)) + np.testing.assert_allclose( + dx.contact.frame[0].reshape(9), d.contact.frame[0] + ) + np.testing.assert_allclose(dx.contact.frame[1:], 0) + + # xmat, ximat, geom_xmat are all shape transformed + self.assertEqual(dx.xmat.shape, (3, 3, 3)) + self.assertEqual(dx.ximat.shape, (3, 3, 3)) + self.assertEqual(dx.geom_xmat.shape, (3, 3, 3)) + np.testing.assert_allclose(dx.xmat.reshape((3, 9)), d.xmat) + np.testing.assert_allclose(dx.ximat.reshape((3, 9)), d.ximat) + np.testing.assert_allclose(dx.geom_xmat.reshape((3, 9)), d.geom_xmat) + + # efc_ are also shape transformed and padded + self.assertEqual(dx.efc_J.shape, (21, 8)) # nefc, nv + d_efc_j = d.efc_J.reshape((-1, 8)) + np.testing.assert_allclose(dx.efc_J[:3], d_efc_j[:3]) # connect eq + np.testing.assert_allclose(dx.efc_J[3], d_efc_j[3]) # one active limit + np.testing.assert_allclose(dx.efc_J[4], 0) # one inactive limit + np.testing.assert_allclose(dx.efc_J[5:9], d_efc_j[4:8]) # contact + np.testing.assert_allclose(dx.efc_J[9:], 0) # no contact + + # check another efc_ too + self.assertEqual(dx.efc_aref.shape, (21,)) # nefc + np.testing.assert_allclose(dx.efc_aref[:3], d.efc_aref[:3]) + np.testing.assert_allclose(dx.efc_aref[3], d.efc_aref[3]) + np.testing.assert_allclose(dx.efc_aref[4], 0) + np.testing.assert_allclose(dx.efc_aref[5:9], d.efc_aref[4:8]) + np.testing.assert_allclose(dx.efc_aref[9:], 0) + + # check sparse transform is correct + m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE + d = mujoco.MjData(m) + mujoco.mj_step(m, d, 2) + dx_from_sparse = mjx.put_data(m, d) + np.testing.assert_allclose(dx_from_sparse.efc_J, dx.efc_J, atol=1e-8) + + def test_get_data(self): + """Test that get_data makes correct MjData.""" + + m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) + d = mujoco.MjData(m) + mujoco.mj_step(m, d, 2) + dx = mjx.put_data(m, d) + d_2: mujoco.MjData = mjx.get_data(m, dx) + + # check a few fields + np.testing.assert_allclose(d_2.qpos, d.qpos) + np.testing.assert_allclose(d_2.xpos, d.xpos) + np.testing.assert_allclose(d_2.cvel, d.cvel) + np.testing.assert_allclose(d_2.cdof_dot, d.cdof_dot) + np.testing.assert_allclose(d_2.qM, d.qM) + + # only 1 contact active + self.assertEqual(d_2.contact.dist.shape, (1,)) + self.assertEqual(d_2.ncon, 1) + np.testing.assert_allclose(d_2.contact.dist, d.contact.dist) + self.assertEqual(d_2.contact.frame.shape, (1, 9)) + np.testing.assert_allclose(d_2.contact.frame, d.contact.frame) + + # xmat, ximat, geom_xmat are all shape transformed + self.assertEqual(d_2.xmat.shape, (3, 9)) + self.assertEqual(d_2.ximat.shape, (3, 9)) + self.assertEqual(d_2.geom_xmat.shape, (3, 9)) + np.testing.assert_allclose(d_2.xmat, d.xmat) + np.testing.assert_allclose(d_2.ximat, d.ximat) + np.testing.assert_allclose(d_2.geom_xmat, d.geom_xmat) + + # efc_* are also shape transformed and filtered + self.assertEqual(d_2.efc_J.shape, (64,)) # nefc * nv + np.testing.assert_allclose(d_2.efc_J, d.efc_J) + self.assertEqual(d_2.efc_aref.shape, (8,)) # nefc + np.testing.assert_allclose(d_2.efc_aref, d.efc_aref) + + # efc_address is created on demand + np.testing.assert_allclose(d_2.contact.efc_address, d.contact.efc_address) + + def test_get_data_batched(self): + """Test that get_data makes correct List[MjData] for batched Data.""" + + m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) + d = mujoco.MjData(m) + mujoco.mj_step(m, d, 2) + dx = mjx.put_data(m, d) + # second data in batch has contact dist > 0, disables contact + dx_b = jax.tree_map(lambda x: jp.stack((x, x + 0.05)), dx) + ds = mjx.get_data(m, dx_b) + self.assertLen(ds, 2) + np.testing.assert_allclose(ds[0].qpos, d.qpos) + np.testing.assert_allclose(ds[1].qpos, d.qpos + 0.05, atol=1e-8) + self.assertEqual(ds[0].ncon, 1) + self.assertEqual(ds[1].ncon, 0) if __name__ == '__main__': absltest.main() diff --git a/mjx/mujoco/mjx/_src/scan.py b/mjx/mujoco/mjx/_src/scan.py index de904e14..ae0eeb30 100644 --- a/mjx/mujoco/mjx/_src/scan.py +++ b/mjx/mujoco/mjx/_src/scan.py @@ -221,7 +221,7 @@ def flat( 'u': i, 'a': m.actuator_actadr[i], 'j': ( - m.actuator_trnid[i] + m.actuator_trnid[i, 0] if m.actuator_trntype[i] == TrnType.JOINT else np.array(-1) ), diff --git a/mjx/mujoco/mjx/_src/scan_test.py b/mjx/mujoco/mjx/_src/scan_test.py index 9b845de1..fe4448ca 100644 --- a/mjx/mujoco/mjx/_src/scan_test.py +++ b/mjx/mujoco/mjx/_src/scan_test.py @@ -210,15 +210,16 @@ class ScanTest(absltest.TestCase): m, fn, 'ujqva', 'ujqva', *args, group_by='u' ) + actuator_trnid = m.actuator_trnid[:, 0] np.testing.assert_array_equal(gear, m.actuator_gear) - np.testing.assert_array_equal(jnt_typ, m.jnt_type[m.actuator_trnid]) + np.testing.assert_array_equal(jnt_typ, m.jnt_type[actuator_trnid]) np.testing.assert_array_equal(act, jp.array([1.4, 1.1])) expected_vadr = np.concatenate( - [np.nonzero(m.dof_jntid == trnid)[0] for trnid in m.actuator_trnid] + [np.nonzero(m.dof_jntid == trnid)[0] for trnid in actuator_trnid] ) np.testing.assert_array_equal(vadr, expected_vadr) expected_qadr = np.concatenate( - [np.nonzero(scan._q_jointid(m) == i)[0] for i in m.actuator_trnid] + [np.nonzero(scan._q_jointid(m) == i)[0] for i in actuator_trnid] ) np.testing.assert_array_equal(qadr, expected_qadr) diff --git a/mjx/mujoco/mjx/_src/solver.py b/mjx/mujoco/mjx/_src/solver.py index c3f5b7c1..f8ef9bed 100644 --- a/mjx/mujoco/mjx/_src/solver.py +++ b/mjx/mujoco/mjx/_src/solver.py @@ -19,6 +19,7 @@ from typing import Optional import jax from jax import numpy as jp import mujoco +from mujoco.mjx._src import constraint from mujoco.mjx._src import math from mujoco.mjx._src import smooth # pylint: disable=g-importing-member @@ -69,12 +70,12 @@ class _Context(PyTreeNode): # TODO(robotics-team): determine nv at which sparse mul is faster M = smooth.dense_m(m, d) if m.nv < 100 else None # pylint: disable=invalid-name ma = smooth.mul_m(m, d, d.qacc) if M is None else M @ d.qacc - nv_0 = jp.zeros((m.nv,)) + nv_0 = jp.zeros(m.nv) ctx = _Context( qacc=d.qacc, qfrc_constraint=d.qfrc_constraint, Jaref=jaref, - efc_force=jp.zeros(d.nefc), + efc_force=d.efc_force, M=M, Ma=ma, grad=nv_0, @@ -111,7 +112,7 @@ class _LSPoint(PyTreeNode): @classmethod def create( cls, - d: Data, + m: Model, ctx: _Context, alpha: jax.Array, jv: jax.Array, @@ -122,7 +123,8 @@ class _LSPoint(PyTreeNode): # roughly corresponds to CGEval in mujoco/src/engine/engine_solver.c # TODO(robotics-team): change this to support friction constraints - active = ((ctx.Jaref + alpha * jv) < 0).at[:d.ne + d.nf].set(True) + ne, nf, *_ = constraint.count_constraints(m) + active = ((ctx.Jaref + alpha * jv) < 0).at[:ne + nf].set(True) quad = jax.vmap(jp.multiply)(quad, active) # only active quad_total = quad_gauss + jp.sum(quad, axis=0) @@ -177,12 +179,11 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context: Returns: context with new constraint force and costs """ - del m - # TODO(robotics-team): add friction constraints # only count active constraints - active = (ctx.Jaref < 0).at[:d.ne + d.nf].set(True) + ne, nf, *_ = constraint.count_constraints(m) + active = (ctx.Jaref < 0).at[:ne + nf].set(True) efc_force = d.efc_D * -ctx.Jaref * active qfrc_constraint = d.efc_J.T @ efc_force @@ -221,7 +222,8 @@ def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context: if m.opt.solver == SolverType.CG: mgrad = smooth.solve_m(m, d, grad) elif m.opt.solver == SolverType.NEWTON: - active = (ctx.Jaref < 0).at[:d.ne + d.nf].set(True) + ne, nf, *_ = constraint.count_constraints(m) + active = (ctx.Jaref < 0).at[:ne + nf].set(True) h = (d.efc_J.T * d.efc_D * active) @ d.efc_J h = smooth.dense_m(m, d) + h h_ = jax.scipy.linalg.cho_factor(h) @@ -265,7 +267,7 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context: quad = jp.stack((0.5 * ctx.Jaref * ctx.Jaref, jv * ctx.Jaref, 0.5 * jv * jv)) quad = (quad * d.efc_D).T - point_fn = lambda alpha: _LSPoint.create(d, ctx, alpha, jv, quad, quad_gauss) + point_fn = lambda a: _LSPoint.create(m, ctx, a, jv, quad, quad_gauss) def cond(ctx: _LSContext) -> jax.Array: done = ctx.ls_iter >= m.opt.ls_iterations diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 963109e6..19b37630 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -21,6 +21,7 @@ import jax import jax.numpy as jp import mujoco # pylint: disable=g-importing-member +from mujoco.mjx._src import dataclasses from mujoco.mjx._src.dataclasses import PyTreeNode # pylint: enable=g-importing-member import numpy as np @@ -213,7 +214,6 @@ class Option(PyTreeNode): integrator: integration mode cone: type of friction cone solver: solver algorithm - integrator: integration mode iterations: number of main solver iterations ls_iterations: maximum number of CG/Newton linesearch iterations disableflags: bit flags for disabling standard features @@ -329,8 +329,8 @@ class Model(PyTreeNode): geom_margin: include in solver if dist 'Contact': + def zero(cls, ncon: int = 0) -> 'Contact': """Returns a contact filled with zeros.""" return Contact( - dist=jp.zeros(shape), - pos=jp.zeros(shape + (3,)), - frame=jp.zeros(shape + (3, 3)), - includemargin=jp.zeros(shape), - friction=jp.zeros(shape + (5,)), - solref=jp.zeros(shape + (mujoco.mjNREF,)), - solreffriction=jp.zeros(shape + (mujoco.mjNREF,)), - solimp=jp.zeros(shape + (mujoco.mjNIMP,)), - dim=np.zeros(shape, dtype=np.int32), - geom1=jp.zeros(shape, dtype=jp.int32), - geom2=jp.zeros(shape, dtype=jp.int32), - efc_address=np.zeros(shape, dtype=np.int32), + dist=jp.zeros(ncon), + pos=jp.zeros((ncon, 3,)), + frame=jp.zeros((ncon, 3, 3)), + includemargin=jp.zeros(ncon), + friction=jp.zeros((ncon, 5)), + solref=jp.zeros((ncon, mujoco.mjNREF)), + solreffriction=jp.zeros((ncon, mujoco.mjNREF)), + solimp=jp.zeros((ncon, mujoco.mjNIMP,)), + geom1=jp.zeros(ncon, dtype=jp.int32), + geom2=jp.zeros(ncon, dtype=jp.int32), ) @@ -552,11 +546,6 @@ class Data(PyTreeNode): Attributes: solver_niter: number of solver iterations, per island (mjNISLAND,) - ne: number of equality constraints - nf: number of friction constraints - nl: number of limit constraints - nefc: number of constraints - ncon: nubmer of contacts time: simulation time qpos: position (nq,) qvel: velocity (nv,) @@ -610,12 +599,6 @@ class Data(PyTreeNode): """ # solver statistics: solver_niter: jax.Array - # sizes (variable in MJ, constant in MJX) - ne: int - nf: int - nl: int - nefc: int - ncon: int # global properties: time: jax.Array # state: diff --git a/mjx/mujoco/mjx/integration_test/collision_driver_test.py b/mjx/mujoco/mjx/integration_test/collision_driver_test.py index d49ef10a..1e28b053 100644 --- a/mjx/mujoco/mjx/integration_test/collision_driver_test.py +++ b/mjx/mujoco/mjx/integration_test/collision_driver_test.py @@ -82,7 +82,6 @@ class CollisionDriverIntegrationTest(parameterized.TestCase): mjx_contact = jax.tree_map( lambda x: x.take(np.array(idx), axis=0), dx.contact ) - mjx_contact = mjx_contact.replace(dim=mjx_contact.dim[idx]) for field in dataclasses.fields(Contact): _assert_attr_eq(mjx_contact, d.contact, field.name, seed, 1e-7)