chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Avoid placing arrays on device for unused MJX fields. Add `device` parameter to mjx.make_data.

Browse Source 
PiperOrigin-RevId: 663808411
Change-Id: Ic55875f87f5c36b9dfe21e5ed5891d64c7003cab
This commit is contained in:
Erik Frey
2024-08-16 12:01:32 -07:00
committed by Copybara-Service
parent ce9c0ea351
commit 390bce2352
3 changed files with 334 additions and 262 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+6 -4
View File
@@ -13,20 +13,22 @@ General
MJX 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``, 5. Added position-dependent sensors: ``MAGNETOMETER``, ``JOINTPOS``, ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``,
``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``SUBTREECOM``, ``CLOCK``. ``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 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`). 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. integrators, wrong derivatives would be computed.
Python bindings 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) Version 3.2.2 (Aug 8, 2024)
mjx/mujoco/mjx/_src/io.py
+205 -143
View File
@@ -70,15 +70,25 @@ def _make_statistic(s: mujoco.MjStatistic) -> types.Statistic:
def put_model( 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: ) -> 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() mesh_geomid = set()
for g1, g2, ip in collision_driver.geom_pairs(m): for g1, g2, ip in collision_driver.geom_pairs(m):
t1, t2 = m.geom_type[[g1, g2]] t1, t2 = m.geom_type[[g1, g2]]
# check collision function exists for type pair # 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) t1, t2 = mujoco.mjtGeom(t1), mujoco.mjtGeom(t2)
raise NotImplementedError(f'({t1}, {t2}) collisions not implemented.') raise NotImplementedError(f'({t1}, {t2}) collisions not implemented.')
# margin/gap not supported for meshes and height fields # margin/gap not supported for meshes and height fields
@@ -88,7 +98,7 @@ def put_model(
margin = m.pair_margin[ip] margin = m.pair_margin[ip]
else: else:
margin = m.geom_margin[g1] + m.geom_margin[g2] 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) t1, t2 = mujoco.mjtGeom(t1), mujoco.mjtGeom(t2)
raise NotImplementedError(f'({t1}, {t2}) margin/gap not implemented.') raise NotImplementedError(f'({t1}, {t2}) margin/gap not implemented.')
for t, g in [(t1, g1), (t2, g2)]: for t, g in [(t1, g1), (t2, g2)]:
@@ -106,16 +116,19 @@ def put_model(
(m.wrap_type, types.WrapType, mujoco.mjtWrap), (m.wrap_type, types.WrapType, mujoco.mjtWrap),
): ):
missing = set(enum_field) - set(enum_type) missing = set(enum_field) - set(enum_type)
if _check_unsupported and missing: if missing and not _full_compat:
raise NotImplementedError( raise NotImplementedError(
f'{[mj_type(m) for m in missing]} not supported' 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.') raise NotImplementedError('dof_frictionloss is not implemented.')
mjx_only = {'mesh_convex', 'geom_rbound_hfield'} mj_field_names = {
mj_field_names = {f.name for f in types.Model.fields()} - mjx_only 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 = {f: getattr(m, f) for f in mj_field_names}
fields['geom_rbound_hfield'] = fields['geom_rbound'] fields['geom_rbound_hfield'] = fields['geom_rbound']
fields['cam_mat0'] = fields['cam_mat0'].reshape((-1, 3, 3)) fields['cam_mat0'] = fields['cam_mat0'].reshape((-1, 3, 3))
@@ -135,29 +148,156 @@ def put_model(
return jax.device_put(model, device=device) return jax.device_put(model, device=device)
def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data: def make_data(
"""Allocate and initialize 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) dim = collision_driver.make_condim(m)
efc_type = constraint.make_efc_type(m, dim) efc_type = constraint.make_efc_type(m, dim)
efc_address = constraint.make_efc_address(m, dim, efc_type) efc_address = constraint.make_efc_address(m, dim, efc_type)
ne, nf, nl, nc = constraint.counts(efc_type) ne, nf, nl, nc = constraint.counts(efc_type)
ncon, nefc = dim.size, ne + nf + nl + nc ncon, nefc = dim.size, ne + nf + nl + nc
contact = types.Contact( with jax.default_device(device):
dist=jp.zeros((ncon,), dtype=float), contact = types.Contact(
pos=jp.zeros((ncon, 3), dtype=float), dist=jp.zeros((ncon,), dtype=float),
frame=jp.zeros((ncon, 3, 3), dtype=float), pos=jp.zeros((ncon, 3), dtype=float),
includemargin=jp.zeros((ncon,), dtype=float), frame=jp.zeros((ncon, 3, 3), dtype=float),
friction=jp.zeros((ncon, 5), dtype=float), includemargin=jp.zeros((ncon,), dtype=float),
solref=jp.zeros((ncon, mujoco.mjNREF), dtype=float), friction=jp.zeros((ncon, 5), dtype=float),
solreffriction=jp.zeros((ncon, mujoco.mjNREF), dtype=float), solref=jp.zeros((ncon, mujoco.mjNREF), dtype=float),
solimp=jp.zeros((ncon, mujoco.mjNIMP), dtype=float), solreffriction=jp.zeros((ncon, mujoco.mjNREF), dtype=float),
dim=dim, solimp=jp.zeros((ncon, mujoco.mjNIMP), dtype=float),
geom1=jp.full((ncon,), -1, dtype=jp.int32), dim=dim,
geom2=jp.full((ncon,), -1, dtype=jp.int32), geom1=jp.full((ncon,), -1, dtype=jp.int32),
geom=jp.full((ncon, 2), -1, dtype=jp.int32), geom2=jp.full((ncon,), -1, dtype=jp.int32),
efc_address=efc_address, 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( d = types.Data(
ne=ne, ne=ne,
@@ -165,119 +305,10 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
nl=nl, nl=nl,
nefc=nefc, nefc=nefc,
ncon=ncon, ncon=ncon,
solver_niter=jp.zeros((), dtype=int),
time=jp.zeros((), dtype=float),
qpos=jp.array(m.qpos0), 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, contact=contact,
efc_type=efc_type, efc_type=efc_type,
efc_J=jp.zeros((nefc, m.nv), dtype=float), **zero_fields
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),
) )
return d return d
@@ -348,7 +379,8 @@ def get_data_into(
result_i.efc_J_colind[:] = np.tile(np.arange(m.nv), nefc) result_i.efc_J_colind[:] = np.tile(np.arange(m.nv), nefc)
for field in types.Data.fields(): 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 continue
if field.name == 'contact': if field.name == 'contact':
@@ -376,6 +408,8 @@ def get_data_into(
value = np.ones(m.nv) value = np.ones(m.nv)
if isinstance(value, np.ndarray) and value.shape: 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 getattr(result_i, field.name)[:] = value
else: else:
setattr(result_i, field.name, value) setattr(result_i, field.name, value)
@@ -416,8 +450,22 @@ def _make_contact(
return types.Contact(**fields), contact_map return types.Contact(**fields), contact_map
def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data: def put_data(
"""Puts mujoco.MjData onto a device, resulting in mjx.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) dim = collision_driver.make_condim(m)
efc_type = constraint.make_efc_type(m, dim) efc_type = constraint.make_efc_type(m, dim)
efc_address = constraint.make_efc_address(m, dim, efc_type) 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: if d_val > val:
raise ValueError(f'd.{name} too high, d.{name} = {d_val}, model = {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() fields = {
if not f.name.endswith('_sparse')} 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: # MJX prefers square matrices for these fields:
for fname in ('xmat', 'ximat', 'geom_xmat', 'site_xmat', 'cam_xmat'): 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 fields[fname] = value
# convert qM and qLD if jacobian is dense # 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): if not support.is_sparse(m):
fields['qM'] = np.zeros((m.nv, m.nv)) fields['qM'] = np.zeros((m.nv, m.nv))
mujoco.mj_fullM(m, fields['qM'], d.qM) 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['qLD'] = np.zeros((m.nv, m.nv))
fields['qLDiagInv'] = np.zeros(0) 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['contact'] = contact
fields.update(ne=ne, nf=nf, nl=nl, nefc=nefc, ncon=ncon, efc_type=efc_type) fields.update(ne=ne, nf=nf, nl=nl, nefc=nefc, ncon=ncon, efc_type=efc_type)
mjx/mujoco/mjx/_src/types.py
+123 -115
View File
@@ -14,6 +14,7 @@
# ============================================================================== # ==============================================================================
"""Base types used in MJX.""" """Base types used in MJX."""
import dataclasses
import enum import enum
from typing import Tuple from typing import Tuple
import jax import jax
@@ -22,6 +23,13 @@ from mujoco.mjx._src.dataclasses import PyTreeNode # pylint: disable=g-importin
import numpy as np 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): class DisableBit(enum.IntFlag):
"""Disable default feature bitflags. """Disable default feature bitflags.
@@ -364,12 +372,12 @@ class Option(PyTreeNode):
sdf_iterations: max number of iterations for gradient descent (not used) sdf_iterations: max number of iterations for gradient descent (not used)
""" """
timestep: jax.Array timestep: jax.Array
apirate: jax.Array apirate: jax.Array = _restricted_to('mujoco')
impratio: jax.Array impratio: jax.Array
tolerance: jax.Array tolerance: jax.Array
ls_tolerance: jax.Array ls_tolerance: jax.Array
noslip_tolerance: jax.Array noslip_tolerance: jax.Array = _restricted_to('mujoco')
mpr_tolerance: jax.Array mpr_tolerance: jax.Array = _restricted_to('mujoco')
gravity: jax.Array gravity: jax.Array
wind: jax.Array wind: jax.Array
magnetic: jax.Array magnetic: jax.Array
@@ -379,20 +387,20 @@ class Option(PyTreeNode):
o_solref: jax.Array o_solref: jax.Array
o_solimp: jax.Array o_solimp: jax.Array
o_friction: jax.Array o_friction: jax.Array
has_fluid_params: bool has_fluid_params: bool = _restricted_to('mjx')
integrator: IntegratorType integrator: IntegratorType
cone: ConeType cone: ConeType
jacobian: JacobianType jacobian: JacobianType
solver: SolverType solver: SolverType
iterations: int iterations: int
ls_iterations: int ls_iterations: int
noslip_iterations: int noslip_iterations: int = _restricted_to('mujoco')
mpr_iterations: int mpr_iterations: int = _restricted_to('mujoco')
disableflags: DisableBit disableflags: DisableBit
enableflags: int enableflags: int
disableactuator: int disableactuator: int
sdf_initpoints: int sdf_initpoints: int = _restricted_to('mujoco')
sdf_iterations: int sdf_iterations: int = _restricted_to('mujoco')
class Statistic(PyTreeNode): class Statistic(PyTreeNode):
@@ -723,22 +731,22 @@ class Model(PyTreeNode):
nu: int nu: int
na: int na: int
nbody: int nbody: int
nbvh: int nbvh: int = _restricted_to('mujoco')
nbvhstatic: int nbvhstatic: int = _restricted_to('mujoco')
nbvhdynamic: int nbvhdynamic: int = _restricted_to('mujoco')
njnt: int njnt: int
ngeom: int ngeom: int
nsite: int nsite: int
ncam: int ncam: int
nlight: int nlight: int
nflex: int nflex: int = _restricted_to('mujoco')
nflexvert: int nflexvert: int = _restricted_to('mujoco')
nflexedge: int nflexedge: int = _restricted_to('mujoco')
nflexelem: int nflexelem: int = _restricted_to('mujoco')
nflexelemdata: int nflexelemdata: int = _restricted_to('mujoco')
nflexshelldata: int nflexshelldata: int = _restricted_to('mujoco')
nflexevpair: int nflexevpair: int = _restricted_to('mujoco')
nflextexcoord: int nflextexcoord: int = _restricted_to('mujoco')
nmesh: int nmesh: int
nmeshvert: int nmeshvert: int
nmeshnormal: int nmeshnormal: int
@@ -761,11 +769,11 @@ class Model(PyTreeNode):
nM: int # pylint:disable=invalid-name nM: int # pylint:disable=invalid-name
nD: int # pylint:disable=invalid-name nD: int # pylint:disable=invalid-name
nB: int # pylint:disable=invalid-name nB: int # pylint:disable=invalid-name
ntree: int ntree: int = _restricted_to('mujoco')
ngravcomp: int ngravcomp: int
nuserdata: int nuserdata: int
nsensordata: int nsensordata: int
narena: int narena: int = _restricted_to('mujoco')
opt: Option opt: Option
stat: Statistic stat: Statistic
qpos0: jax.Array qpos0: jax.Array
@@ -794,11 +802,11 @@ class Model(PyTreeNode):
body_margin: np.ndarray body_margin: np.ndarray
body_contype: np.ndarray body_contype: np.ndarray
body_conaffinity: np.ndarray body_conaffinity: np.ndarray
body_bvhadr: np.ndarray body_bvhadr: np.ndarray = _restricted_to('mujoco')
body_bvhnum: np.ndarray body_bvhnum: np.ndarray = _restricted_to('mujoco')
bvh_child: np.ndarray bvh_child: np.ndarray = _restricted_to('mujoco')
bvh_nodeid: np.ndarray bvh_nodeid: np.ndarray = _restricted_to('mujoco')
bvh_aabb: np.ndarray bvh_aabb: np.ndarray = _restricted_to('mujoco')
body_invweight0: jax.Array body_invweight0: jax.Array
jnt_type: np.ndarray jnt_type: np.ndarray
jnt_qposadr: np.ndarray jnt_qposadr: np.ndarray
@@ -844,7 +852,7 @@ class Model(PyTreeNode):
geom_size: jax.Array geom_size: jax.Array
geom_aabb: np.ndarray geom_aabb: np.ndarray
geom_rbound: jax.Array geom_rbound: jax.Array
geom_rbound_hfield: np.ndarray geom_rbound_hfield: np.ndarray = _restricted_to('mjx')
geom_pos: jax.Array geom_pos: jax.Array
geom_quat: jax.Array geom_quat: jax.Array
geom_friction: jax.Array geom_friction: jax.Array
@@ -870,54 +878,54 @@ class Model(PyTreeNode):
cam_resolution: np.ndarray cam_resolution: np.ndarray
cam_sensorsize: np.ndarray cam_sensorsize: np.ndarray
cam_intrinsic: np.ndarray cam_intrinsic: np.ndarray
light_mode: np.ndarray light_mode: np.ndarray = _restricted_to('mujoco')
light_bodyid: np.ndarray light_bodyid: np.ndarray = _restricted_to('mujoco')
light_targetbodyid: np.ndarray light_targetbodyid: np.ndarray = _restricted_to('mujoco')
light_pos: np.ndarray light_pos: np.ndarray = _restricted_to('mujoco')
light_dir: np.ndarray light_dir: np.ndarray = _restricted_to('mujoco')
light_poscom0: np.ndarray light_poscom0: np.ndarray = _restricted_to('mujoco')
light_pos0: np.ndarray light_pos0: np.ndarray = _restricted_to('mujoco')
light_dir0: np.ndarray light_dir0: np.ndarray = _restricted_to('mujoco')
flex_contype: np.ndarray flex_contype: np.ndarray = _restricted_to('mujoco')
flex_conaffinity: np.ndarray flex_conaffinity: np.ndarray = _restricted_to('mujoco')
flex_condim: np.ndarray flex_condim: np.ndarray = _restricted_to('mujoco')
flex_priority: np.ndarray flex_priority: np.ndarray = _restricted_to('mujoco')
flex_solmix: np.ndarray flex_solmix: np.ndarray = _restricted_to('mujoco')
flex_solref: np.ndarray flex_solref: np.ndarray = _restricted_to('mujoco')
flex_solimp: np.ndarray flex_solimp: np.ndarray = _restricted_to('mujoco')
flex_friction: np.ndarray flex_friction: np.ndarray = _restricted_to('mujoco')
flex_margin: np.ndarray flex_margin: np.ndarray = _restricted_to('mujoco')
flex_gap: np.ndarray flex_gap: np.ndarray = _restricted_to('mujoco')
flex_internal: np.ndarray flex_internal: np.ndarray = _restricted_to('mujoco')
flex_selfcollide: np.ndarray flex_selfcollide: np.ndarray = _restricted_to('mujoco')
flex_activelayers: np.ndarray flex_activelayers: np.ndarray = _restricted_to('mujoco')
flex_dim: np.ndarray flex_dim: np.ndarray = _restricted_to('mujoco')
flex_vertadr: np.ndarray flex_vertadr: np.ndarray = _restricted_to('mujoco')
flex_vertnum: np.ndarray flex_vertnum: np.ndarray = _restricted_to('mujoco')
flex_edgeadr: np.ndarray flex_edgeadr: np.ndarray = _restricted_to('mujoco')
flex_edgenum: np.ndarray flex_edgenum: np.ndarray = _restricted_to('mujoco')
flex_elemadr: np.ndarray flex_elemadr: np.ndarray = _restricted_to('mujoco')
flex_elemnum: np.ndarray flex_elemnum: np.ndarray = _restricted_to('mujoco')
flex_elemdataadr: np.ndarray flex_elemdataadr: np.ndarray = _restricted_to('mujoco')
flex_evpairadr: np.ndarray flex_evpairadr: np.ndarray = _restricted_to('mujoco')
flex_evpairnum: np.ndarray flex_evpairnum: np.ndarray = _restricted_to('mujoco')
flex_vertbodyid: np.ndarray flex_vertbodyid: np.ndarray = _restricted_to('mujoco')
flex_edge: np.ndarray flex_edge: np.ndarray = _restricted_to('mujoco')
flex_elem: np.ndarray flex_elem: np.ndarray = _restricted_to('mujoco')
flex_elemlayer: np.ndarray flex_elemlayer: np.ndarray = _restricted_to('mujoco')
flex_evpair: np.ndarray flex_evpair: np.ndarray = _restricted_to('mujoco')
flex_vert: np.ndarray flex_vert: np.ndarray = _restricted_to('mujoco')
flexedge_length0: np.ndarray flexedge_length0: np.ndarray = _restricted_to('mujoco')
flexedge_invweight0: np.ndarray flexedge_invweight0: np.ndarray = _restricted_to('mujoco')
flex_radius: np.ndarray flex_radius: np.ndarray = _restricted_to('mujoco')
flex_edgestiffness: np.ndarray flex_edgestiffness: np.ndarray = _restricted_to('mujoco')
flex_edgedamping: np.ndarray flex_edgedamping: np.ndarray = _restricted_to('mujoco')
flex_edgeequality: np.ndarray flex_edgeequality: np.ndarray = _restricted_to('mujoco')
flex_rigid: np.ndarray flex_rigid: np.ndarray = _restricted_to('mujoco')
flexedge_rigid: np.ndarray flexedge_rigid: np.ndarray = _restricted_to('mujoco')
flex_centered: np.ndarray flex_centered: np.ndarray = _restricted_to('mujoco')
flex_bvhadr: np.ndarray flex_bvhadr: np.ndarray = _restricted_to('mujoco')
flex_bvhnum: np.ndarray flex_bvhnum: np.ndarray = _restricted_to('mujoco')
mesh_vertadr: np.ndarray mesh_vertadr: np.ndarray
mesh_vertnum: np.ndarray mesh_vertnum: np.ndarray
mesh_faceadr: np.ndarray mesh_faceadr: np.ndarray
@@ -929,7 +937,7 @@ class Model(PyTreeNode):
mesh_graph: np.ndarray mesh_graph: np.ndarray
mesh_pos: np.ndarray mesh_pos: np.ndarray
mesh_quat: np.ndarray mesh_quat: np.ndarray
mesh_convex: Tuple[ConvexMesh, ...] mesh_convex: Tuple[ConvexMesh, ...] = _restricted_to('mjx')
hfield_size: np.ndarray hfield_size: np.ndarray
hfield_nrow: np.ndarray hfield_nrow: np.ndarray
hfield_ncol: np.ndarray hfield_ncol: np.ndarray
@@ -969,9 +977,9 @@ class Model(PyTreeNode):
tendon_lengthspring: jax.Array tendon_lengthspring: jax.Array
tendon_length0: jax.Array tendon_length0: jax.Array
tendon_invweight0: jax.Array tendon_invweight0: jax.Array
wrap_type: np.ndarray wrap_type: np.ndarray = _restricted_to('mujoco')
wrap_objid: np.ndarray wrap_objid: np.ndarray = _restricted_to('mujoco')
wrap_prm: np.ndarray wrap_prm: np.ndarray = _restricted_to('mujoco')
actuator_trntype: np.ndarray actuator_trntype: np.ndarray
actuator_dyntype: np.ndarray actuator_dyntype: np.ndarray
actuator_gaintype: np.ndarray actuator_gaintype: np.ndarray
@@ -994,7 +1002,7 @@ class Model(PyTreeNode):
actuator_cranklength: np.ndarray actuator_cranklength: np.ndarray
actuator_acc0: np.ndarray actuator_acc0: np.ndarray
actuator_lengthrange: np.ndarray actuator_lengthrange: np.ndarray
actuator_plugin: np.ndarray actuator_plugin: np.ndarray = _restricted_to('mujoco')
sensor_type: np.ndarray sensor_type: np.ndarray
sensor_datatype: np.ndarray sensor_datatype: np.ndarray
sensor_needstage: np.ndarray sensor_needstage: np.ndarray
@@ -1202,8 +1210,8 @@ class Data(PyTreeNode):
xfrc_applied: jax.Array xfrc_applied: jax.Array
eq_active: jax.Array eq_active: jax.Array
# mocap data: # mocap data:
mocap_pos: jax.Array mocap_pos: jax.Array = _restricted_to('mujoco')
mocap_quat: jax.Array mocap_quat: jax.Array = _restricted_to('mujoco')
# dynamics: # dynamics:
qacc: jax.Array qacc: jax.Array
act_dot: jax.Array act_dot: jax.Array
@@ -1224,27 +1232,27 @@ class Data(PyTreeNode):
site_xmat: jax.Array site_xmat: jax.Array
cam_xpos: jax.Array cam_xpos: jax.Array
cam_xmat: jax.Array cam_xmat: jax.Array
light_xpos: jax.Array light_xpos: jax.Array = _restricted_to('mujoco')
light_xdir: jax.Array light_xdir: jax.Array = _restricted_to('mujoco')
subtree_com: jax.Array subtree_com: jax.Array
cdof: jax.Array cdof: jax.Array
cinert: jax.Array cinert: jax.Array
flexvert_xpos: jax.Array flexvert_xpos: jax.Array = _restricted_to('mujoco')
flexelem_aabb: jax.Array flexelem_aabb: jax.Array
flexedge_J_rownnz: jax.Array # pylint:disable=invalid-name flexedge_J_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
flexedge_J_rowadr: jax.Array # pylint:disable=invalid-name flexedge_J_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
flexedge_J_colind: jax.Array # pylint:disable=invalid-name flexedge_J_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
flexedge_J: jax.Array # pylint:disable=invalid-name flexedge_J: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
flexedge_length: jax.Array flexedge_length: jax.Array = _restricted_to('mujoco')
ten_wrapadr: jax.Array ten_wrapadr: jax.Array = _restricted_to('mujoco')
ten_wrapnum: jax.Array ten_wrapnum: jax.Array = _restricted_to('mujoco')
ten_J_rownnz: jax.Array # pylint:disable=invalid-name ten_J_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
ten_J_rowadr: jax.Array # pylint:disable=invalid-name ten_J_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
ten_J_colind: jax.Array # 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_J: jax.Array # pylint:disable=invalid-name
ten_length: jax.Array ten_length: jax.Array
wrap_obj: jax.Array wrap_obj: jax.Array = _restricted_to('mujoco')
wrap_xpos: jax.Array wrap_xpos: jax.Array = _restricted_to('mujoco')
actuator_length: jax.Array actuator_length: jax.Array
actuator_moment: jax.Array actuator_moment: jax.Array
crb: jax.Array crb: jax.Array
@@ -1252,32 +1260,32 @@ class Data(PyTreeNode):
qLD: jax.Array # pylint:disable=invalid-name qLD: jax.Array # pylint:disable=invalid-name
qLDiagInv: jax.Array # pylint:disable=invalid-name qLDiagInv: jax.Array # pylint:disable=invalid-name
qLDiagSqrtInv: jax.Array # pylint:disable=invalid-name qLDiagSqrtInv: jax.Array # pylint:disable=invalid-name
bvh_aabb_dyn: jax.Array bvh_aabb_dyn: jax.Array = _restricted_to('mujoco')
bvh_active: jax.Array bvh_active: jax.Array = _restricted_to('mujoco')
# position, velocity dependent: # position, velocity dependent:
flexedge_velocity: jax.Array flexedge_velocity: jax.Array = _restricted_to('mujoco')
ten_velocity: jax.Array ten_velocity: jax.Array
actuator_velocity: jax.Array actuator_velocity: jax.Array
cvel: jax.Array cvel: jax.Array
cdof_dot: jax.Array cdof_dot: jax.Array
qfrc_bias: jax.Array qfrc_bias: jax.Array
qfrc_spring: jax.Array qfrc_spring: jax.Array = _restricted_to('mujoco')
qfrc_damper: jax.Array qfrc_damper: jax.Array = _restricted_to('mujoco')
qfrc_gravcomp: jax.Array qfrc_gravcomp: jax.Array
qfrc_fluid: jax.Array qfrc_fluid: jax.Array
qfrc_passive: jax.Array qfrc_passive: jax.Array
subtree_linvel: jax.Array subtree_linvel: jax.Array
subtree_angmom: jax.Array subtree_angmom: jax.Array
qH: jax.Array # pylint:disable=invalid-name qH: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
qHDiagInv: jax.Array # pylint:disable=invalid-name qHDiagInv: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
D_rownnz: jax.Array # pylint:disable=invalid-name D_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
D_rowadr: jax.Array # pylint:disable=invalid-name D_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
D_colind: jax.Array # pylint:disable=invalid-name D_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
B_rownnz: jax.Array # pylint:disable=invalid-name B_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
B_rowadr: jax.Array # pylint:disable=invalid-name B_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
B_colind: jax.Array # pylint:disable=invalid-name B_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
qDeriv: jax.Array # pylint:disable=invalid-name qDeriv: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
qLU: jax.Array # pylint:disable=invalid-name qLU: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
# position, velocity, control & acceleration dependent: # position, velocity, control & acceleration dependent:
qfrc_actuator: jax.Array qfrc_actuator: jax.Array
actuator_force: jax.Array actuator_force: jax.Array
@@ -1302,6 +1310,6 @@ class Data(PyTreeNode):
efc_force: jax.Array efc_force: jax.Array
# sparse representation of qM, qLD, qLDiagInv, for compatibility with MuJoCo # sparse representation of qM, qLD, qLDiagInv, for compatibility with MuJoCo
# when in dense mode # when in dense mode
_qM_sparse: jax.Array # pylint:disable=invalid-name _qM_sparse: jax.Array = _restricted_to('mjx') # pylint:disable=invalid-name
_qLD_sparse: jax.Array # pylint:disable=invalid-name _qLD_sparse: jax.Array = _restricted_to('mjx') # pylint:disable=invalid-name
_qLDiagInv_sparse: jax.Array # pylint:disable=invalid-name _qLDiagInv_sparse: jax.Array = _restricted_to('mjx') # pylint:disable=invalid-name