Avoid placing arrays on device for unused MJX fields. Add `device` parameter to mjx.make_data.
PiperOrigin-RevId: 663808411 Change-Id: Ic55875f87f5c36b9dfe21e5ed5891d64c7003cab
This commit is contained in:
Erik Frey
committed by
Copybara-Service
Copybara-Service
parent
ce9c0ea351
commit
390bce2352
+6
-4
@@ -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)
|
||||
|
||||
+205
-143
@@ -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)
|
||||
|
||||
|
||||
+123
-115
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user