diff --git a/doc/changelog.rst b/doc/changelog.rst index 40713d29..54b706c6 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -13,20 +13,22 @@ General MJX ^^^ -4. Added ``efc_pos`` to ``mjx.Data``. +4. Added ``efc_pos`` to ``mjx.Data`` (:github:issue:`1388`). 5. Added position-dependent sensors: ``MAGNETOMETER``, ``JOINTPOS``, ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``SUBTREECOM``, ``CLOCK``. +6. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device. +7. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``. Bug fixes ^^^^^^^^^ -6. Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`, +8. Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`, contribution by :github:user:`michael-ahn`). -7. Fixed a bug wherein, for models that have both muscles and stateless actuators and used one of the implicit +9. 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 ^^^^^^^^^^^^^^^ -8. Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`). +10. Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`). Version 3.2.2 (Aug 8, 2024) diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index d10a4760..7ad1c261 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -70,15 +70,25 @@ def _make_statistic(s: mujoco.MjStatistic) -> types.Statistic: def put_model( - m: mujoco.MjModel, device=None, _check_unsupported=True + m: mujoco.MjModel, device=None, _full_compat: bool = False # pylint: disable=invalid-name ) -> types.Model: - """Puts mujoco.MjModel onto a device, resulting in mjx.Model.""" + """Puts mujoco.MjModel onto a device, resulting in mjx.Model. + + Args: + m: the model to put onto device + device: which device to use - if unspecified picks the default device + _full_compat: put all MjModel fields onto device irrespective of MJX support + This is an experimental feature. Avoid using it for now. + + Returns: + an mjx.Model placed on device + """ mesh_geomid = set() for g1, g2, ip in collision_driver.geom_pairs(m): t1, t2 = m.geom_type[[g1, g2]] # check collision function exists for type pair - if _check_unsupported and not collision_driver.has_collision_fn(t1, t2): + if not collision_driver.has_collision_fn(t1, t2) and not _full_compat: t1, t2 = mujoco.mjtGeom(t1), mujoco.mjtGeom(t2) raise NotImplementedError(f'({t1}, {t2}) collisions not implemented.') # margin/gap not supported for meshes and height fields @@ -88,7 +98,7 @@ def put_model( margin = m.pair_margin[ip] else: margin = m.geom_margin[g1] + m.geom_margin[g2] - if _check_unsupported and margin.any(): + if margin.any() and not _full_compat: t1, t2 = mujoco.mjtGeom(t1), mujoco.mjtGeom(t2) raise NotImplementedError(f'({t1}, {t2}) margin/gap not implemented.') for t, g in [(t1, g1), (t2, g2)]: @@ -106,16 +116,19 @@ def put_model( (m.wrap_type, types.WrapType, mujoco.mjtWrap), ): missing = set(enum_field) - set(enum_type) - if _check_unsupported and missing: + if missing and not _full_compat: raise NotImplementedError( f'{[mj_type(m) for m in missing]} not supported' ) - if _check_unsupported and not np.allclose(m.dof_frictionloss, 0): + if not np.allclose(m.dof_frictionloss, 0) and not _full_compat: raise NotImplementedError('dof_frictionloss is not implemented.') - mjx_only = {'mesh_convex', 'geom_rbound_hfield'} - mj_field_names = {f.name for f in types.Model.fields()} - mjx_only + mj_field_names = { + f.name + for f in types.Model.fields() + if f.metadata.get('restricted_to') != 'mjx' + } fields = {f: getattr(m, f) for f in mj_field_names} fields['geom_rbound_hfield'] = fields['geom_rbound'] fields['cam_mat0'] = fields['cam_mat0'].reshape((-1, 3, 3)) @@ -135,29 +148,156 @@ def put_model( return jax.device_put(model, device=device) -def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data: - """Allocate and initialize Data.""" +def make_data( + m: Union[types.Model, mujoco.MjModel], + device=None, + _full_compat: bool = False, # pylint: disable=invalid-name +) -> types.Data: + """Allocate and initialize Data. + + Args: + m: the model to use + device: which device to use - if unspecified picks the default device + _full_compat: create all MjData fields on device irrespective of MJX support + This is an experimental feature. Avoid using it for now. + If using this flag, also use _full_compat for put_model. + + Returns: + an initialized mjx.Data placed on device + """ dim = collision_driver.make_condim(m) efc_type = constraint.make_efc_type(m, dim) efc_address = constraint.make_efc_address(m, dim, efc_type) ne, nf, nl, nc = constraint.counts(efc_type) ncon, nefc = dim.size, ne + nf + nl + nc - contact = types.Contact( - dist=jp.zeros((ncon,), dtype=float), - pos=jp.zeros((ncon, 3), dtype=float), - frame=jp.zeros((ncon, 3, 3), dtype=float), - includemargin=jp.zeros((ncon,), dtype=float), - friction=jp.zeros((ncon, 5), dtype=float), - solref=jp.zeros((ncon, mujoco.mjNREF), dtype=float), - solreffriction=jp.zeros((ncon, mujoco.mjNREF), dtype=float), - solimp=jp.zeros((ncon, mujoco.mjNIMP), dtype=float), - dim=dim, - geom1=jp.full((ncon,), -1, dtype=jp.int32), - geom2=jp.full((ncon,), -1, dtype=jp.int32), - geom=jp.full((ncon, 2), -1, dtype=jp.int32), - efc_address=efc_address, - ) + with jax.default_device(device): + contact = types.Contact( + dist=jp.zeros((ncon,), dtype=float), + pos=jp.zeros((ncon, 3), dtype=float), + frame=jp.zeros((ncon, 3, 3), dtype=float), + includemargin=jp.zeros((ncon,), dtype=float), + friction=jp.zeros((ncon, 5), dtype=float), + solref=jp.zeros((ncon, mujoco.mjNREF), dtype=float), + solreffriction=jp.zeros((ncon, mujoco.mjNREF), dtype=float), + solimp=jp.zeros((ncon, mujoco.mjNIMP), dtype=float), + dim=dim, + geom1=jp.full((ncon,), -1, dtype=jp.int32), + geom2=jp.full((ncon,), -1, dtype=jp.int32), + geom=jp.full((ncon, 2), -1, dtype=jp.int32), + efc_address=efc_address, + ) + + zero_fields = { + 'solver_niter': (int,), + 'time': (float,), + 'qvel': (m.nv, float), + 'act': (m.na, float), + 'qacc_warmstart': (m.nv, float), + 'ctrl': (m.nu, float), + 'qfrc_applied': (m.nv, float), + 'xfrc_applied': (m.nbody, 6, float), + 'eq_active': (m.neq, jp.uint8), + 'mocap_pos': (m.nmocap, 3, float), + 'mocap_quat': (m.nmocap, 4, float), + 'qacc': (m.nv, float), + 'act_dot': (m.na, float), + 'userdata': (m.nuserdata, float), + 'sensordata': (m.nsensordata, float), + 'xpos': (m.nbody, 3, float), + 'xquat': (m.nbody, 4, float), + 'xmat': (m.nbody, 3, 3, float), + 'xipos': (m.nbody, 3, float), + 'ximat': (m.nbody, 3, 3, float), + 'xanchor': (m.njnt, 3, float), + 'xaxis': (m.njnt, 3, float), + 'geom_xpos': (m.ngeom, 3, float), + 'geom_xmat': (m.ngeom, 3, 3, float), + 'site_xpos': (m.nsite, 3, float), + 'site_xmat': (m.nsite, 3, 3, float), + 'cam_xpos': (m.ncam, 3, float), + 'cam_xmat': (m.ncam, 3, 3, float), + 'light_xpos': (m.nlight, 3, float), + 'light_xdir': (m.nlight, 3, float), + 'subtree_com': (m.nbody, 3, float), + 'cdof': (m.nv, 6, float), + 'cinert': (m.nbody, 10, float), + 'flexvert_xpos': (m.nflexvert, 3, float), + 'flexelem_aabb': (m.nflexelem, 6, float), + 'flexedge_J_rownnz': (m.nflexedge, jp.int32), + 'flexedge_J_rowadr': (m.nflexedge, jp.int32), + 'flexedge_J_colind': (m.nflexedge, m.nv, jp.int32), + 'flexedge_J': (m.nflexedge, m.nv, float), + 'flexedge_length': (m.nflexedge, float), + 'ten_wrapadr': (m.ntendon, jp.int32), + 'ten_wrapnum': (m.ntendon, jp.int32), + 'ten_J_rownnz': (m.ntendon, jp.int32), + 'ten_J_rowadr': (m.ntendon, jp.int32), + 'ten_J_colind': (m.ntendon, m.nv, jp.int32), + 'ten_J': (m.ntendon, m.nv, float), + 'ten_length': (m.ntendon, float), + 'wrap_obj': (m.nwrap, 2, jp.int32), + 'wrap_xpos': (m.nwrap, 6, float), + 'actuator_length': (m.nu, float), + 'actuator_moment': (m.nu, m.nv, float), + 'crb': (m.nbody, 10, float), + 'qM': (m.nM, float) if support.is_sparse(m) else (m.nv, m.nv, float), + 'qLD': (m.nM, float) if support.is_sparse(m) else (m.nv, m.nv, float), + 'qLDiagInv': (m.nM, float) if support.is_sparse(m) else (0, float), + 'qLDiagSqrtInv': (m.nv, float), + 'bvh_aabb_dyn': (m.nbvhdynamic, 6, float), + 'bvh_active': (m.nbvh, jp.uint8), + 'flexedge_velocity': (m.nflexedge, float), + 'ten_velocity': (m.ntendon, float), + 'actuator_velocity': (m.nu, float), + 'cvel': (m.nbody, 6, float), + 'cdof_dot': (m.nv, 6, float), + 'qfrc_bias': (m.nv, float), + 'qfrc_spring': (m.nv, float), + 'qfrc_damper': (m.nv, float), + 'qfrc_gravcomp': (m.nv, float), + 'qfrc_fluid': (m.nv, float), + 'qfrc_passive': (m.nv, float), + 'subtree_linvel': (m.nbody, 3, float), + 'subtree_angmom': (m.nbody, 3, float), + 'qH': (m.nM, float) if support.is_sparse(m) else (m.nv, m.nv, float), + 'qHDiagInv': (m.nv, float), + 'D_rownnz': (m.nv, jp.int32), + 'D_rowadr': (m.nv, jp.int32), + 'D_colind': (m.nD, jp.int32), + 'B_rownnz': (m.nbody, jp.int32), + 'B_rowadr': (m.nbody, jp.int32), + 'B_colind': (m.nB, jp.int32), + 'qDeriv': (m.nD, float), + 'qLU': (m.nD, float), + 'actuator_force': (m.nu, float), + 'qfrc_actuator': (m.nv, float), + 'qfrc_smooth': (m.nv, float), + 'qacc_smooth': (m.nv, float), + 'qfrc_constraint': (m.nv, float), + 'qfrc_inverse': (m.nv, float), + 'cacc': (m.nbody, 6, float), + 'cfrc_int': (m.nbody, 6, float), + 'cfrc_ext': (m.nbody, 6, float), + 'efc_J': (nefc, m.nv, float), + 'efc_pos': (nefc, float), + 'efc_frictionloss': (nefc, float), + 'efc_D': (nefc, float), + 'efc_aref': (nefc, float), + 'efc_force': (nefc, float), + '_qM_sparse': (m.nM, float), + '_qLD_sparse': (m.nM, float), + '_qLDiagInv_sparse': (m.nv, float), + } + + if not _full_compat: + for f in types.Data.fields(): + if f.metadata.get('restricted_to') in ('mujoco', 'mjx'): + zero_fields[f.name] = (0, zero_fields[f.name][-1]) + + zero_fields = { + k: jp.zeros(v[:-1], dtype=v[-1]) for k, v in zero_fields.items() + } d = types.Data( ne=ne, @@ -165,119 +305,10 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data: nl=nl, nefc=nefc, ncon=ncon, - solver_niter=jp.zeros((), dtype=int), - time=jp.zeros((), dtype=float), qpos=jp.array(m.qpos0), - qvel=jp.zeros((m.nv,), dtype=float), - act=jp.zeros((m.na,), dtype=float), - qacc_warmstart=jp.zeros((m.nv,), dtype=float), - ctrl=jp.zeros((m.nu,), dtype=float), - qfrc_applied=jp.zeros((m.nv,), dtype=float), - xfrc_applied=jp.zeros((m.nbody, 6), dtype=float), - eq_active=jp.zeros((m.neq,), dtype=jp.uint8), - mocap_pos=jp.zeros((m.nmocap, 3), dtype=float), - mocap_quat=jp.zeros((m.nmocap, 4), dtype=float), - qacc=jp.zeros((m.nv,), dtype=float), - act_dot=jp.zeros((m.na,), dtype=float), - userdata=jp.zeros((m.nuserdata,), dtype=float), - sensordata=jp.zeros((m.nsensordata,), dtype=float), - xpos=jp.zeros((m.nbody, 3), dtype=float), - xquat=jp.zeros((m.nbody, 4), dtype=float), - xmat=jp.zeros((m.nbody, 3, 3), dtype=float), - xipos=jp.zeros((m.nbody, 3), dtype=float), - ximat=jp.zeros((m.nbody, 3, 3), dtype=float), - xanchor=jp.zeros((m.njnt, 3), dtype=float), - xaxis=jp.zeros((m.njnt, 3), dtype=float), - geom_xpos=jp.zeros((m.ngeom, 3), dtype=float), - geom_xmat=jp.zeros((m.ngeom, 3, 3), dtype=float), - site_xpos=jp.zeros((m.nsite, 3), dtype=float), - site_xmat=jp.zeros((m.nsite, 3, 3), dtype=float), - cam_xpos=jp.zeros((m.ncam, 3), dtype=float), - cam_xmat=jp.zeros((m.ncam, 3, 3), dtype=float), - light_xpos=jp.zeros((m.nlight, 3), dtype=float), - light_xdir=jp.zeros((m.nlight, 3), dtype=float), - subtree_com=jp.zeros((m.nbody, 3), dtype=float), - cdof=jp.zeros((m.nv, 6), dtype=float), - cinert=jp.zeros((m.nbody, 10), dtype=float), - flexvert_xpos=jp.zeros((m.nflexvert, 3), dtype=float), - flexelem_aabb=jp.zeros((m.nflexelem, 6), dtype=float), - flexedge_J_rownnz=jp.zeros((m.nflexedge,), dtype=jp.int32), - flexedge_J_rowadr=jp.zeros((m.nflexedge,), dtype=jp.int32), - flexedge_J_colind=jp.zeros((m.nflexedge, m.nv), dtype=jp.int32), - flexedge_J=jp.zeros((m.nflexedge, m.nv), dtype=float), - flexedge_length=jp.zeros((m.nflexedge,), dtype=float), - ten_wrapadr=jp.zeros((m.ntendon,), dtype=jp.int32), - ten_wrapnum=jp.zeros((m.ntendon,), dtype=jp.int32), - ten_J_rownnz=jp.zeros((m.ntendon,), dtype=jp.int32), - ten_J_rowadr=jp.zeros((m.ntendon,), dtype=jp.int32), - ten_J_colind=jp.zeros((m.ntendon, m.nv), dtype=jp.int32), - ten_J=jp.zeros((m.ntendon, m.nv), dtype=float), - ten_length=jp.zeros((m.ntendon,), dtype=float), - wrap_obj=jp.zeros((m.nwrap, 2), dtype=jp.int32), - wrap_xpos=jp.zeros((m.nwrap, 6), dtype=float), - actuator_length=jp.zeros((m.nu,), dtype=float), - actuator_moment=jp.zeros((m.nu, m.nv), dtype=float), - crb=jp.zeros((m.nbody, 10), dtype=float), - qM=( - jp.zeros((m.nM,), dtype=float) - if support.is_sparse(m) - else jp.zeros((m.nv, m.nv), dtype=float) - ), - qLD=( - jp.zeros((m.nM,), dtype=float) - if support.is_sparse(m) - else jp.zeros((m.nv, m.nv), dtype=float) - ), - qLDiagInv=( - jp.zeros((m.nv,), dtype=float) if support.is_sparse(m) - else jp.zeros((0,), dtype=float) - ), - qLDiagSqrtInv=jp.zeros((m.nv,), dtype=float), - bvh_aabb_dyn=jp.zeros((m.nbvhdynamic, 6), dtype=float), - bvh_active=jp.zeros((m.nbvh,), dtype=jp.uint8), - flexedge_velocity=jp.zeros((m.nflexedge,), dtype=float), - ten_velocity=jp.zeros((m.ntendon,), dtype=float), - actuator_velocity=jp.zeros((m.nu,), dtype=float), - cvel=jp.zeros((m.nbody, 6), dtype=float), - cdof_dot=jp.zeros((m.nv, 6), dtype=float), - qfrc_bias=jp.zeros((m.nv,), dtype=float), - qfrc_spring=jp.zeros((m.nv,), dtype=float), - qfrc_damper=jp.zeros((m.nv,), dtype=float), - qfrc_gravcomp=jp.zeros((m.nv,), dtype=float), - qfrc_fluid=jp.zeros((m.nv,), dtype=float), - qfrc_passive=jp.zeros((m.nv,), dtype=float), - subtree_linvel=jp.zeros((m.nbody, 3), dtype=float), - subtree_angmom=jp.zeros((m.nbody, 3), dtype=float), - qH=jp.zeros((m.nM,), dtype=float), - qHDiagInv=jp.zeros((m.nv,), dtype=float), - D_rownnz=jp.zeros((m.nv,), dtype=jp.int32), - D_rowadr=jp.zeros((m.nv,), dtype=jp.int32), - D_colind=jp.zeros((m.nD,), dtype=jp.int32), - B_rownnz=jp.zeros((m.nbody,), dtype=jp.int32), - B_rowadr=jp.zeros((m.nbody,), dtype=jp.int32), - B_colind=jp.zeros((m.nB,), dtype=jp.int32), - qDeriv=jp.zeros((m.nD,), dtype=float), - qLU=jp.zeros((m.nD,), dtype=float), - actuator_force=jp.zeros((m.nu,), dtype=float), - qfrc_actuator=jp.zeros((m.nv,), dtype=float), - qfrc_smooth=jp.zeros((m.nv,), dtype=float), - qacc_smooth=jp.zeros((m.nv,), dtype=float), - qfrc_constraint=jp.zeros((m.nv,), dtype=float), - qfrc_inverse=jp.zeros((m.nv,), dtype=float), - cacc=jp.zeros((m.nbody, 6), dtype=float), - cfrc_int=jp.zeros((m.nbody, 6), dtype=float), - cfrc_ext=jp.zeros((m.nbody, 6), dtype=float), contact=contact, efc_type=efc_type, - efc_J=jp.zeros((nefc, m.nv), dtype=float), - efc_pos=jp.zeros((nefc,), dtype=float), - efc_frictionloss=jp.zeros((nefc,), dtype=float), - efc_D=jp.zeros((nefc,), dtype=float), - efc_aref=jp.zeros((nefc,), dtype=float), - efc_force=jp.zeros((nefc,), dtype=float), - _qM_sparse=jp.zeros((m.nM), dtype=float), - _qLD_sparse=jp.zeros((m.nM), dtype=float), - _qLDiagInv_sparse=jp.zeros((m.nv,), dtype=float), + **zero_fields ) return d @@ -348,7 +379,8 @@ def get_data_into( result_i.efc_J_colind[:] = np.tile(np.arange(m.nv), nefc) for field in types.Data.fields(): - if field.name.startswith('_') and field.name.endswith('_sparse'): + restricted_to = field.metadata.get('restricted_to') + if restricted_to == 'mjx': continue if field.name == 'contact': @@ -376,6 +408,8 @@ def get_data_into( value = np.ones(m.nv) if isinstance(value, np.ndarray) and value.shape: + if restricted_to in ('mujoco', 'mjx') and value.shape == (0,): + continue # don't copy fields that are mujoco-only or MJX-only getattr(result_i, field.name)[:] = value else: setattr(result_i, field.name, value) @@ -416,8 +450,22 @@ def _make_contact( return types.Contact(**fields), contact_map -def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data: - """Puts mujoco.MjData onto a device, resulting in mjx.Data.""" +def put_data( + m: mujoco.MjModel, d: mujoco.MjData, device=None, _full_compat: bool = False # pylint: disable=invalid-name +) -> types.Data: + """Puts mujoco.MjData onto a device, resulting in mjx.Data. + + Args: + m: the model to use + d: the data to put on device + device: which device to use - if unspecified picks the default device + _full_compat: put all MjModel fields onto device irrespective of MJX support + This is an experimental feature. Avoid using it for now. + If using this flag, also use _full_compat for put_model. + + Returns: + an mjx.Data placed on device + """ dim = collision_driver.make_condim(m) efc_type = constraint.make_efc_type(m, dim) efc_address = constraint.make_efc_address(m, dim, efc_type) @@ -434,8 +482,11 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data: if d_val > val: raise ValueError(f'd.{name} too high, d.{name} = {d_val}, model = {val}') - fields = {f.name: getattr(d, f.name) for f in types.Data.fields() - if not f.name.endswith('_sparse')} + fields = { + f.name: getattr(d, f.name) + for f in types.Data.fields() + if f.metadata.get('restricted_to') != 'mjx' + } # MJX prefers square matrices for these fields: for fname in ('xmat', 'ximat', 'geom_xmat', 'site_xmat', 'cam_xmat'): @@ -490,9 +541,6 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data: fields[fname] = value # convert qM and qLD if jacobian is dense - fields['_qM_sparse'] = fields['qM'] - fields['_qLD_sparse'] = fields['qLD'] - fields['_qLDiagInv_sparse'] = fields['qLDiagInv'] if not support.is_sparse(m): fields['qM'] = np.zeros((m.nv, m.nv)) mujoco.mj_fullM(m, fields['qM'], d.qM) @@ -504,6 +552,20 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data: fields['qLD'] = np.zeros((m.nv, m.nv)) fields['qLDiagInv'] = np.zeros(0) + if _full_compat: + # full compatibility mode, we store sparse qM regardless of jacobian setting + fields['_qM_sparse'] = fields['qM'] + fields['_qLD_sparse'] = fields['qLD'] + fields['_qLDiagInv_sparse'] = fields['qLDiagInv'] + else: + fields['_qM_sparse'] = jp.zeros(0, dtype=float) + fields['_qLD_sparse'] = jp.zeros(0, dtype=float) + fields['_qLDiagInv_sparse'] = jp.zeros(0, dtype=float) + # otherwise clear out unused arrays + for f in types.Data.fields(): + if f.metadata.get('restricted_to') == 'mujoco': + fields[f.name] = np.zeros(0, dtype=fields[f.name].dtype) + fields['contact'] = contact fields.update(ne=ne, nf=nf, nl=nl, nefc=nefc, ncon=ncon, efc_type=efc_type) diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 8b690c11..c3f2d2ff 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -14,6 +14,7 @@ # ============================================================================== """Base types used in MJX.""" +import dataclasses import enum from typing import Tuple import jax @@ -22,6 +23,13 @@ from mujoco.mjx._src.dataclasses import PyTreeNode # pylint: disable=g-importin import numpy as np +def _restricted_to(platform: str): + """Specifies whether a field exists in only MuJoCo or MJX.""" + if platform not in ('mujoco', 'mjx'): + raise ValueError(f'unknown platform: {platform}') + return dataclasses.field(metadata={'restricted_to': platform}) + + class DisableBit(enum.IntFlag): """Disable default feature bitflags. @@ -364,12 +372,12 @@ class Option(PyTreeNode): sdf_iterations: max number of iterations for gradient descent (not used) """ timestep: jax.Array - apirate: jax.Array + apirate: jax.Array = _restricted_to('mujoco') impratio: jax.Array tolerance: jax.Array ls_tolerance: jax.Array - noslip_tolerance: jax.Array - mpr_tolerance: jax.Array + noslip_tolerance: jax.Array = _restricted_to('mujoco') + mpr_tolerance: jax.Array = _restricted_to('mujoco') gravity: jax.Array wind: jax.Array magnetic: jax.Array @@ -379,20 +387,20 @@ class Option(PyTreeNode): o_solref: jax.Array o_solimp: jax.Array o_friction: jax.Array - has_fluid_params: bool + has_fluid_params: bool = _restricted_to('mjx') integrator: IntegratorType cone: ConeType jacobian: JacobianType solver: SolverType iterations: int ls_iterations: int - noslip_iterations: int - mpr_iterations: int + noslip_iterations: int = _restricted_to('mujoco') + mpr_iterations: int = _restricted_to('mujoco') disableflags: DisableBit enableflags: int disableactuator: int - sdf_initpoints: int - sdf_iterations: int + sdf_initpoints: int = _restricted_to('mujoco') + sdf_iterations: int = _restricted_to('mujoco') class Statistic(PyTreeNode): @@ -723,22 +731,22 @@ class Model(PyTreeNode): nu: int na: int nbody: int - nbvh: int - nbvhstatic: int - nbvhdynamic: int + nbvh: int = _restricted_to('mujoco') + nbvhstatic: int = _restricted_to('mujoco') + nbvhdynamic: int = _restricted_to('mujoco') njnt: int ngeom: int nsite: int ncam: int nlight: int - nflex: int - nflexvert: int - nflexedge: int - nflexelem: int - nflexelemdata: int - nflexshelldata: int - nflexevpair: int - nflextexcoord: int + nflex: int = _restricted_to('mujoco') + nflexvert: int = _restricted_to('mujoco') + nflexedge: int = _restricted_to('mujoco') + nflexelem: int = _restricted_to('mujoco') + nflexelemdata: int = _restricted_to('mujoco') + nflexshelldata: int = _restricted_to('mujoco') + nflexevpair: int = _restricted_to('mujoco') + nflextexcoord: int = _restricted_to('mujoco') nmesh: int nmeshvert: int nmeshnormal: int @@ -761,11 +769,11 @@ class Model(PyTreeNode): nM: int # pylint:disable=invalid-name nD: int # pylint:disable=invalid-name nB: int # pylint:disable=invalid-name - ntree: int + ntree: int = _restricted_to('mujoco') ngravcomp: int nuserdata: int nsensordata: int - narena: int + narena: int = _restricted_to('mujoco') opt: Option stat: Statistic qpos0: jax.Array @@ -794,11 +802,11 @@ class Model(PyTreeNode): body_margin: np.ndarray body_contype: np.ndarray body_conaffinity: np.ndarray - body_bvhadr: np.ndarray - body_bvhnum: np.ndarray - bvh_child: np.ndarray - bvh_nodeid: np.ndarray - bvh_aabb: np.ndarray + body_bvhadr: np.ndarray = _restricted_to('mujoco') + body_bvhnum: np.ndarray = _restricted_to('mujoco') + bvh_child: np.ndarray = _restricted_to('mujoco') + bvh_nodeid: np.ndarray = _restricted_to('mujoco') + bvh_aabb: np.ndarray = _restricted_to('mujoco') body_invweight0: jax.Array jnt_type: np.ndarray jnt_qposadr: np.ndarray @@ -844,7 +852,7 @@ class Model(PyTreeNode): geom_size: jax.Array geom_aabb: np.ndarray geom_rbound: jax.Array - geom_rbound_hfield: np.ndarray + geom_rbound_hfield: np.ndarray = _restricted_to('mjx') geom_pos: jax.Array geom_quat: jax.Array geom_friction: jax.Array @@ -870,54 +878,54 @@ class Model(PyTreeNode): cam_resolution: np.ndarray cam_sensorsize: np.ndarray cam_intrinsic: np.ndarray - light_mode: np.ndarray - light_bodyid: np.ndarray - light_targetbodyid: np.ndarray - light_pos: np.ndarray - light_dir: np.ndarray - light_poscom0: np.ndarray - light_pos0: np.ndarray - light_dir0: np.ndarray - flex_contype: np.ndarray - flex_conaffinity: np.ndarray - flex_condim: np.ndarray - flex_priority: np.ndarray - flex_solmix: np.ndarray - flex_solref: np.ndarray - flex_solimp: np.ndarray - flex_friction: np.ndarray - flex_margin: np.ndarray - flex_gap: np.ndarray - flex_internal: np.ndarray - flex_selfcollide: np.ndarray - flex_activelayers: np.ndarray - flex_dim: np.ndarray - flex_vertadr: np.ndarray - flex_vertnum: np.ndarray - flex_edgeadr: np.ndarray - flex_edgenum: np.ndarray - flex_elemadr: np.ndarray - flex_elemnum: np.ndarray - flex_elemdataadr: np.ndarray - flex_evpairadr: np.ndarray - flex_evpairnum: np.ndarray - flex_vertbodyid: np.ndarray - flex_edge: np.ndarray - flex_elem: np.ndarray - flex_elemlayer: np.ndarray - flex_evpair: np.ndarray - flex_vert: np.ndarray - flexedge_length0: np.ndarray - flexedge_invweight0: np.ndarray - flex_radius: np.ndarray - flex_edgestiffness: np.ndarray - flex_edgedamping: np.ndarray - flex_edgeequality: np.ndarray - flex_rigid: np.ndarray - flexedge_rigid: np.ndarray - flex_centered: np.ndarray - flex_bvhadr: np.ndarray - flex_bvhnum: np.ndarray + light_mode: np.ndarray = _restricted_to('mujoco') + light_bodyid: np.ndarray = _restricted_to('mujoco') + light_targetbodyid: np.ndarray = _restricted_to('mujoco') + light_pos: np.ndarray = _restricted_to('mujoco') + light_dir: np.ndarray = _restricted_to('mujoco') + light_poscom0: np.ndarray = _restricted_to('mujoco') + light_pos0: np.ndarray = _restricted_to('mujoco') + light_dir0: np.ndarray = _restricted_to('mujoco') + flex_contype: np.ndarray = _restricted_to('mujoco') + flex_conaffinity: np.ndarray = _restricted_to('mujoco') + flex_condim: np.ndarray = _restricted_to('mujoco') + flex_priority: np.ndarray = _restricted_to('mujoco') + flex_solmix: np.ndarray = _restricted_to('mujoco') + flex_solref: np.ndarray = _restricted_to('mujoco') + flex_solimp: np.ndarray = _restricted_to('mujoco') + flex_friction: np.ndarray = _restricted_to('mujoco') + flex_margin: np.ndarray = _restricted_to('mujoco') + flex_gap: np.ndarray = _restricted_to('mujoco') + flex_internal: np.ndarray = _restricted_to('mujoco') + flex_selfcollide: np.ndarray = _restricted_to('mujoco') + flex_activelayers: np.ndarray = _restricted_to('mujoco') + flex_dim: np.ndarray = _restricted_to('mujoco') + flex_vertadr: np.ndarray = _restricted_to('mujoco') + flex_vertnum: np.ndarray = _restricted_to('mujoco') + flex_edgeadr: np.ndarray = _restricted_to('mujoco') + flex_edgenum: np.ndarray = _restricted_to('mujoco') + flex_elemadr: np.ndarray = _restricted_to('mujoco') + flex_elemnum: np.ndarray = _restricted_to('mujoco') + flex_elemdataadr: np.ndarray = _restricted_to('mujoco') + flex_evpairadr: np.ndarray = _restricted_to('mujoco') + flex_evpairnum: np.ndarray = _restricted_to('mujoco') + flex_vertbodyid: np.ndarray = _restricted_to('mujoco') + flex_edge: np.ndarray = _restricted_to('mujoco') + flex_elem: np.ndarray = _restricted_to('mujoco') + flex_elemlayer: np.ndarray = _restricted_to('mujoco') + flex_evpair: np.ndarray = _restricted_to('mujoco') + flex_vert: np.ndarray = _restricted_to('mujoco') + flexedge_length0: np.ndarray = _restricted_to('mujoco') + flexedge_invweight0: np.ndarray = _restricted_to('mujoco') + flex_radius: np.ndarray = _restricted_to('mujoco') + flex_edgestiffness: np.ndarray = _restricted_to('mujoco') + flex_edgedamping: np.ndarray = _restricted_to('mujoco') + flex_edgeequality: np.ndarray = _restricted_to('mujoco') + flex_rigid: np.ndarray = _restricted_to('mujoco') + flexedge_rigid: np.ndarray = _restricted_to('mujoco') + flex_centered: np.ndarray = _restricted_to('mujoco') + flex_bvhadr: np.ndarray = _restricted_to('mujoco') + flex_bvhnum: np.ndarray = _restricted_to('mujoco') mesh_vertadr: np.ndarray mesh_vertnum: np.ndarray mesh_faceadr: np.ndarray @@ -929,7 +937,7 @@ class Model(PyTreeNode): mesh_graph: np.ndarray mesh_pos: np.ndarray mesh_quat: np.ndarray - mesh_convex: Tuple[ConvexMesh, ...] + mesh_convex: Tuple[ConvexMesh, ...] = _restricted_to('mjx') hfield_size: np.ndarray hfield_nrow: np.ndarray hfield_ncol: np.ndarray @@ -969,9 +977,9 @@ class Model(PyTreeNode): tendon_lengthspring: jax.Array tendon_length0: jax.Array tendon_invweight0: jax.Array - wrap_type: np.ndarray - wrap_objid: np.ndarray - wrap_prm: np.ndarray + wrap_type: np.ndarray = _restricted_to('mujoco') + wrap_objid: np.ndarray = _restricted_to('mujoco') + wrap_prm: np.ndarray = _restricted_to('mujoco') actuator_trntype: np.ndarray actuator_dyntype: np.ndarray actuator_gaintype: np.ndarray @@ -994,7 +1002,7 @@ class Model(PyTreeNode): actuator_cranklength: np.ndarray actuator_acc0: np.ndarray actuator_lengthrange: np.ndarray - actuator_plugin: np.ndarray + actuator_plugin: np.ndarray = _restricted_to('mujoco') sensor_type: np.ndarray sensor_datatype: np.ndarray sensor_needstage: np.ndarray @@ -1202,8 +1210,8 @@ class Data(PyTreeNode): xfrc_applied: jax.Array eq_active: jax.Array # mocap data: - mocap_pos: jax.Array - mocap_quat: jax.Array + mocap_pos: jax.Array = _restricted_to('mujoco') + mocap_quat: jax.Array = _restricted_to('mujoco') # dynamics: qacc: jax.Array act_dot: jax.Array @@ -1224,27 +1232,27 @@ class Data(PyTreeNode): site_xmat: jax.Array cam_xpos: jax.Array cam_xmat: jax.Array - light_xpos: jax.Array - light_xdir: jax.Array + light_xpos: jax.Array = _restricted_to('mujoco') + light_xdir: jax.Array = _restricted_to('mujoco') subtree_com: jax.Array cdof: jax.Array cinert: jax.Array - flexvert_xpos: jax.Array + flexvert_xpos: jax.Array = _restricted_to('mujoco') flexelem_aabb: jax.Array - flexedge_J_rownnz: jax.Array # pylint:disable=invalid-name - flexedge_J_rowadr: jax.Array # pylint:disable=invalid-name - flexedge_J_colind: jax.Array # pylint:disable=invalid-name - flexedge_J: jax.Array # pylint:disable=invalid-name - flexedge_length: jax.Array - ten_wrapadr: jax.Array - ten_wrapnum: jax.Array - ten_J_rownnz: jax.Array # pylint:disable=invalid-name - ten_J_rowadr: jax.Array # pylint:disable=invalid-name - ten_J_colind: jax.Array # pylint:disable=invalid-name + flexedge_J_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name + flexedge_J_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name + flexedge_J_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name + flexedge_J: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name + flexedge_length: jax.Array = _restricted_to('mujoco') + ten_wrapadr: jax.Array = _restricted_to('mujoco') + ten_wrapnum: jax.Array = _restricted_to('mujoco') + ten_J_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name + ten_J_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name + ten_J_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name ten_J: jax.Array # pylint:disable=invalid-name ten_length: jax.Array - wrap_obj: jax.Array - wrap_xpos: jax.Array + wrap_obj: jax.Array = _restricted_to('mujoco') + wrap_xpos: jax.Array = _restricted_to('mujoco') actuator_length: jax.Array actuator_moment: jax.Array crb: jax.Array @@ -1252,32 +1260,32 @@ class Data(PyTreeNode): qLD: jax.Array # pylint:disable=invalid-name qLDiagInv: jax.Array # pylint:disable=invalid-name qLDiagSqrtInv: jax.Array # pylint:disable=invalid-name - bvh_aabb_dyn: jax.Array - bvh_active: jax.Array + bvh_aabb_dyn: jax.Array = _restricted_to('mujoco') + bvh_active: jax.Array = _restricted_to('mujoco') # position, velocity dependent: - flexedge_velocity: jax.Array + flexedge_velocity: jax.Array = _restricted_to('mujoco') ten_velocity: jax.Array actuator_velocity: jax.Array cvel: jax.Array cdof_dot: jax.Array qfrc_bias: jax.Array - qfrc_spring: jax.Array - qfrc_damper: jax.Array + qfrc_spring: jax.Array = _restricted_to('mujoco') + qfrc_damper: jax.Array = _restricted_to('mujoco') qfrc_gravcomp: jax.Array qfrc_fluid: jax.Array qfrc_passive: jax.Array subtree_linvel: jax.Array subtree_angmom: jax.Array - qH: jax.Array # pylint:disable=invalid-name - qHDiagInv: jax.Array # pylint:disable=invalid-name - D_rownnz: jax.Array # pylint:disable=invalid-name - D_rowadr: jax.Array # pylint:disable=invalid-name - D_colind: jax.Array # pylint:disable=invalid-name - B_rownnz: jax.Array # pylint:disable=invalid-name - B_rowadr: jax.Array # pylint:disable=invalid-name - B_colind: jax.Array # pylint:disable=invalid-name - qDeriv: jax.Array # pylint:disable=invalid-name - qLU: jax.Array # pylint:disable=invalid-name + qH: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name + qHDiagInv: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name + D_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name + D_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name + D_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name + B_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name + B_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name + B_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name + qDeriv: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name + qLU: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name # position, velocity, control & acceleration dependent: qfrc_actuator: jax.Array actuator_force: jax.Array @@ -1302,6 +1310,6 @@ class Data(PyTreeNode): efc_force: jax.Array # sparse representation of qM, qLD, qLDiagInv, for compatibility with MuJoCo # when in dense mode - _qM_sparse: jax.Array # pylint:disable=invalid-name - _qLD_sparse: jax.Array # pylint:disable=invalid-name - _qLDiagInv_sparse: jax.Array # pylint:disable=invalid-name + _qM_sparse: jax.Array = _restricted_to('mjx') # pylint:disable=invalid-name + _qLD_sparse: jax.Array = _restricted_to('mjx') # pylint:disable=invalid-name + _qLDiagInv_sparse: jax.Array = _restricted_to('mjx') # pylint:disable=invalid-name