# Copyright 2023 DeepMind Technologies Limited # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """Functions to initialize, load, or save data.""" import copy from typing import List, Tuple, Union import jax from jax import numpy as jp import mujoco from mujoco.mjx._src import collision_driver from mujoco.mjx._src import constraint from mujoco.mjx._src import mesh from mujoco.mjx._src import support from mujoco.mjx._src import types import numpy as np import scipy def _strip_weak_type(tree): def f(leaf): if isinstance(leaf, jax.Array): return leaf.astype(jax.dtypes.canonicalize_dtype(leaf.dtype)) return leaf return jax.tree_util.tree_map(f, tree) def _make_option( o: mujoco.MjOption, _full_compat: bool = False # pylint: disable=invalid-name ) -> types.Option: """Returns mjx.Option given mujoco.MjOption.""" if not _full_compat: if o.integrator not in set(types.IntegratorType): raise NotImplementedError(f'{mujoco.mjtIntegrator(o.integrator)}') if o.cone not in set(types.ConeType): raise NotImplementedError(f'{mujoco.mjtCone(o.cone)}') if o.jacobian not in set(types.JacobianType): raise NotImplementedError(f'{mujoco.mjtJacobian(o.jacobian)}') if o.solver not in set(types.SolverType): raise NotImplementedError(f'{mujoco.mjtSolver(o.solver)}') for i in range(mujoco.mjtEnableBit.mjNENABLE): if o.enableflags & 2**i and 2**i not in set(types.EnableBit): raise NotImplementedError(f'{mujoco.mjtEnableBit(2**i)}') has_fluid_params = o.density > 0 or o.viscosity > 0 or o.wind.any() implicitfast = o.integrator == mujoco.mjtIntegrator.mjINT_IMPLICITFAST if not _full_compat: if implicitfast and has_fluid_params: raise NotImplementedError('implicitfast not implemented for fluid drag.') fields = {f.name: getattr(o, f.name, None) for f in types.Option.fields()} fields['integrator'] = types.IntegratorType(o.integrator) fields['cone'] = types.ConeType(o.cone) fields['jacobian'] = types.JacobianType(o.jacobian) fields['solver'] = types.SolverType(o.solver) fields['disableflags'] = types.DisableBit(o.disableflags) fields['has_fluid_params'] = has_fluid_params fields['enableflags'] = types.EnableBit(o.enableflags) return types.Option(**fields) def _make_statistic(s: mujoco.MjStatistic) -> types.Statistic: """Puts mujoco.MjStatistic onto a device, resulting in mjx.Statistic.""" return types.Statistic( meaninertia=s.meaninertia, meanmass=s.meanmass, meansize=s.meansize, extent=s.extent, center=s.center, ) def put_model( 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. 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 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 no_margin = {mujoco.mjtGeom.mjGEOM_MESH, mujoco.mjtGeom.mjGEOM_HFIELD} if no_margin.intersection({t1, t2}): if ip != -1: margin = m.pair_margin[ip] else: margin = m.geom_margin[g1] + m.geom_margin[g2] 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)]: if t == mujoco.mjtGeom.mjGEOM_MESH: mesh_geomid.add(g) # check for unsupported sensor and equality constraint combinations sensor_rne_postconstraint = ( np.any(m.sensor_type == types.SensorType.ACCELEROMETER) | np.any(m.sensor_type == types.SensorType.FORCE) | np.any(m.sensor_type == types.SensorType.TORQUE) | np.any(m.sensor_type == types.SensorType.FRAMELINACC) | np.any(m.sensor_type == types.SensorType.FRAMEANGACC) ) eq_connect_weld = np.any(m.eq_type == types.EqType.CONNECT) | np.any( m.eq_type == types.EqType.WELD ) if sensor_rne_postconstraint and eq_connect_weld: raise NotImplementedError( 'rne_postconstraint not implemented with equality constraints:' ' connect, weld.' ) for enum_field, enum_type, mj_type in ( (m.actuator_biastype, types.BiasType, mujoco.mjtBias), (m.actuator_dyntype, types.DynType, mujoco.mjtDyn), (m.actuator_gaintype, types.GainType, mujoco.mjtGain), (m.actuator_trntype, types.TrnType, mujoco.mjtTrn), (m.eq_type, types.EqType, mujoco.mjtEq), (m.sensor_type, types.SensorType, mujoco.mjtSensor), (m.wrap_type, types.WrapType, mujoco.mjtWrap), ): missing = set(enum_field) - set(enum_type) if missing and not _full_compat: raise NotImplementedError( f'{[mj_type(m) for m in missing]} not supported' ) 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} # zero out fields restricted to MuJoCo if not _full_compat: for f in types.Model.fields(): if f.metadata.get('restricted_to') == 'mujoco' and isinstance( fields[f.name], np.ndarray ): fields[f.name] = np.zeros((0,), dtype=fields[f.name].dtype) fields['dof_hasfrictionloss'] = fields['dof_frictionloss'] > 0 fields['tendon_hasfrictionloss'] = fields['tendon_frictionloss'] > 0 fields['geom_rbound_hfield'] = fields['geom_rbound'] fields['cam_mat0'] = fields['cam_mat0'].reshape((-1, 3, 3)) fields['opt'] = _make_option(m.opt, _full_compat=_full_compat) fields['stat'] = _make_statistic(m.stat) # spatial tendon wrap inside fields['wrap_inside_maxiter'] = 5 fields['wrap_inside_tolerance'] = 1.0e-4 fields['wrap_inside_z_init'] = 1.0 - 1.0e-5 fields['is_wrap_inside'] = np.zeros(0, dtype=bool) if m.nsite: # find sphere or cylinder geoms (if any exist) (wrap_id_geom,) = np.nonzero( (m.wrap_type == mujoco.mjtWrap.mjWRAP_SPHERE) | (m.wrap_type == mujoco.mjtWrap.mjWRAP_CYLINDER) ) wrap_objid_geom = m.wrap_objid[wrap_id_geom] geom_pos = m.geom_pos[wrap_objid_geom] geom_size = m.geom_size[wrap_objid_geom, 0] # find sidesites (if any exist) side_id = np.round(m.wrap_prm[wrap_id_geom]).astype(int) side = m.site_pos[side_id] # wrap inside flag fields['is_wrap_inside'] = np.array( (np.linalg.norm(side - geom_pos, axis=1) < geom_size) & (side_id >= 0) ) # Pre-compile meshes for MJX collisions. fields['mesh_convex'] = [None] * m.nmesh if not _full_compat: for i in mesh_geomid: dataid = m.geom_dataid[i] if fields['mesh_convex'][dataid] is None: fields['mesh_convex'][dataid] = mesh.convex(m, dataid) # pytype: disable=unsupported-operands fields['mesh_convex'] = tuple(fields['mesh_convex']) model = types.Model(**{k: copy.copy(v) for k, v in fields.items()}) model = jax.device_put(model, device=device) return _strip_weak_type(model) 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 float_ = jp.zeros(1, float).dtype int_ = jp.zeros(1, int).dtype contact = types.Contact( dist=np.zeros((ncon,), dtype=float_), pos=np.zeros((ncon, 3), dtype=float_), frame=np.zeros((ncon, 3, 3), dtype=float_), includemargin=np.zeros((ncon,), dtype=float_), friction=np.zeros((ncon, 5), dtype=float_), solref=np.zeros((ncon, mujoco.mjNREF), dtype=float_), solreffriction=np.zeros((ncon, mujoco.mjNREF), dtype=float_), solimp=np.zeros((ncon, mujoco.mjNIMP), dtype=float_), dim=dim, # let jax pick contact.geom int precision, for interop with # jax_enable_x64 geom1=np.full((ncon,), -1, dtype=int_), geom2=np.full((ncon,), -1, dtype=int_), geom=np.full((ncon, 2), -1, dtype=int_), efc_address=efc_address, ) if m.opt.cone == types.ConeType.ELLIPTIC and np.any(contact.dim == 1): raise NotImplementedError( 'condim=1 with ConeType.ELLIPTIC not implemented.' ) 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_), '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, np.int32), 'flexedge_J_rowadr': (m.nflexedge, np.int32), 'flexedge_J_colind': (m.nflexedge, m.nv, np.int32), 'flexedge_J': (m.nflexedge, m.nv, float_), 'flexedge_length': (m.nflexedge, float_), 'ten_wrapadr': (m.ntendon, np.int32), 'ten_wrapnum': (m.ntendon, np.int32), 'ten_J_rownnz': (m.ntendon, np.int32), 'ten_J_rowadr': (m.ntendon, np.int32), 'ten_J_colind': (m.ntendon, m.nv, np.int32), 'ten_J': (m.ntendon, m.nv, float_), 'ten_length': (m.ntendon, float_), 'wrap_obj': (m.nwrap, 2, np.int32), 'wrap_xpos': (m.nwrap, 6, float_), 'actuator_length': (m.nu, float_), 'moment_rownnz': (m.nu, np.int32), 'moment_rowadr': (m.nu, np.int32), 'moment_colind': (m.nJmom, np.int32), '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.nv, float_) if support.is_sparse(m) else (0, float_), 'bvh_aabb_dyn': (m.nbvhdynamic, 6, float_), 'bvh_active': (m.nbvh, np.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_), 'B_rownnz': (m.nbody, np.int32), 'B_rowadr': (m.nbody, np.int32), 'B_colind': (m.nB, np.int32), 'M_rownnz': (m.nv, np.int32), 'M_rowadr': (m.nv, np.int32), 'M_colind': (m.nM, np.int32), 'mapM2M': (m.nM, np.int32), 'C_rownnz': (m.nv, np.int32), 'C_rowadr': (m.nv, np.int32), 'C_colind': (m.nC, np.int32), 'mapM2C': (m.nC, np.int32), 'D_rownnz': (m.nv, np.int32), 'D_rowadr': (m.nv, np.int32), 'D_diag': (m.nv, np.int32), 'D_colind': (m.nD, np.int32), 'mapM2D': (m.nD, np.int32), 'mapD2M': (m.nM, np.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_margin': (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: np.zeros(v[:-1], dtype=v[-1]) for k, v in zero_fields.items() } d = types.Data( ne=ne, nf=nf, nl=nl, nefc=nefc, ncon=ncon, qpos=jp.array(m.qpos0, dtype=float_), contact=contact, efc_type=efc_type, eq_active=m.eq_active0, **zero_fields, ) d = jax.device_put(d, device=device) return d def _get_contact(c: mujoco._structs._MjContactList, cx: types.Contact): """Converts mjx.Contact to mujoco._structs._MjContactList.""" con_id = np.nonzero(cx.dist <= 0)[0] for field in types.Contact.fields(): value = getattr(cx, field.name)[con_id] if field.name == 'frame': value = value.reshape((-1, 9)) getattr(c, field.name)[:] = value def get_data( m: mujoco.MjModel, d: types.Data ) -> Union[mujoco.MjData, List[mujoco.MjData]]: """Gets mjx.Data from a device, resulting in mujoco.MjData or List[MjData].""" batched = len(d.qpos.shape) > 1 batch_size = d.qpos.shape[0] if batched else 1 if batched: result = [mujoco.MjData(m) for _ in range(batch_size)] else: result = mujoco.MjData(m) get_data_into(result, m, d) return result def get_data_into( result: Union[mujoco.MjData, List[mujoco.MjData]], m: mujoco.MjModel, d: types.Data, ): """Gets mjx.Data from a device into an existing mujoco.MjData or list.""" batched = isinstance(result, list) if batched and len(d.qpos.shape) < 2: raise ValueError('dst is a list, but d is not batched.') if not batched and len(d.qpos.shape) >= 2: raise ValueError('dst is a an MjData, but d is batched.') d = jax.device_get(d) batch_size = d.qpos.shape[0] if batched else 1 dof_i, dof_j = [], [] for i in range(m.nv): j = i while j > -1: dof_i.append(i) dof_j.append(j) j = m.dof_parentid[j] for i in range(batch_size): d_i = jax.tree_util.tree_map(lambda x, i=i: x[i], d) if batched else d result_i = result[i] if batched else result ncon = (d_i.contact.dist <= 0).sum() efc_active = (d_i.efc_J != 0).any(axis=1) nefc = int(efc_active.sum()) result_i.nJ = nefc * m.nv if ncon != result_i.ncon or nefc != result_i.nefc: mujoco._functions._realloc_con_efc(result_i, ncon=ncon, nefc=nefc) # pylint: disable=protected-access result_i.efc_J_rownnz[:] = np.repeat(m.nv, nefc) result_i.efc_J_rowadr[:] = np.arange(0, nefc * m.nv, m.nv) result_i.efc_J_colind[:] = np.tile(np.arange(m.nv), nefc) for field in types.Data.fields(): restricted_to = field.metadata.get('restricted_to') if restricted_to == 'mjx': continue if field.name == 'contact': _get_contact(result_i.contact, d_i.contact) # efc_address must be updated because rows were deleted above: efc_map = np.cumsum(efc_active) - 1 result_i.contact.efc_address[:] = efc_map[result_i.contact.efc_address] continue # MuJoCo actuator_moment is sparse, MJX uses a dense representation. if field.name == 'actuator_moment': moment_rownnz = np.zeros(m.nu, dtype=np.int32) moment_rowadr = np.zeros(m.nu, dtype=np.int32) moment_colind = np.zeros(m.nJmom, dtype=np.int32) actuator_moment = np.zeros(m.nJmom) if m.nu: mujoco.mju_dense2sparse( actuator_moment, d_i.actuator_moment, moment_rownnz, moment_rowadr, moment_colind, ) result_i.moment_rownnz[:] = moment_rownnz result_i.moment_rowadr[:] = moment_rowadr result_i.moment_colind[:] = moment_colind result_i.actuator_moment[:] = actuator_moment continue value = getattr(d_i, field.name) if field.name in ('nefc', 'ncon'): value = {'nefc': nefc, 'ncon': ncon}[field.name] elif field.name.endswith('xmat') or field.name == 'ximat': value = value.reshape((-1, 9)) elif field.name.startswith('efc_'): value = value[efc_active] if field.name == 'efc_J': value = value.reshape(-1) elif field.name == 'qM' and not support.is_sparse(m): value = value[dof_i, dof_j] elif field.name == 'qLD' and not support.is_sparse(m): # TODO(erikfrey): provide correct qLDs value = np.zeros(m.nM) elif field.name == 'qLDiagInv' and not support.is_sparse(m): value = np.ones(m.nv) if isinstance(value, np.ndarray) and value.shape: if restricted_to in ('mujoco', 'mjx'): continue # don't copy fields that are mujoco-only or MJX-only else: result_field = getattr(result_i, field.name) if result_field.shape != value.shape: raise ValueError( f'Input field {field.name} has shape {value.shape}, but output' f' has shape {result_field.shape}' ) result_field[:] = value else: setattr(result_i, field.name, value) def _make_contact( c: mujoco._structs._MjContactList, dim: np.ndarray, efc_address: np.ndarray, ) -> Tuple[types.Contact, np.ndarray]: """Converts mujoco.structs._MjContactList into mjx.Contact.""" fields = {f.name: getattr(c, f.name) for f in types.Contact.fields()} fields['frame'] = fields['frame'].reshape((-1, 3, 3)) # reorder contacts so that their condims match those specified in dim. # if we have fewer Contacts for a condim range, pad the range with zeros # build a map for where to find a dim-matching contact, or -1 if none contact_map = -np.ones_like(dim) for i, di in enumerate(fields['dim']): space = [j for j, dj in enumerate(dim) if di == dj and contact_map[j] == -1] if not space: # this can happen if max_geom_pairs or max_contact_points is too low raise ValueError(f'unable to place Contact[{i}], no space in condim {di}') contact_map[space[0]] = i if contact_map.size > 0: # reorganize contact according, with a zero contact at the end for -1 zero = jax.tree_util.tree_map( lambda x: np.zeros((1,) + x.shape[1:], dtype=x.dtype), fields ) zero['dist'][:] = 1e10 fields = jax.tree_util.tree_map(lambda *x: np.concatenate(x), fields, zero) fields = jax.tree_util.tree_map(lambda x: x[contact_map], fields) fields['dim'] = dim fields['efc_address'] = efc_address return types.Contact(**fields), contact_map 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) ne, nf, nl, nc = constraint.counts(efc_type) ncon, nefc = dim.size, ne + nf + nl + nc for d_val, val, name in ( (d.ncon, ncon, 'ncon'), (d.ne, ne, 'ne'), (d.nf, nf, 'nf'), (d.nl, nl, 'nl'), (d.nefc, nefc, 'nefc'), ): if d_val > val: raise ValueError(f'd.{name} too high, d.{name} = {d_val}, model = {val}') fields = { f.name: 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'): fields[fname] = fields[fname].reshape((-1, 3, 3)) # MJX does not support islanding, so only transfer the first solver_niter fields['solver_niter'] = fields['solver_niter'][0] # convert sparse representation of actuator_moment to dense matrix moment = np.zeros((m.nu, m.nv)) mujoco.mju_sparse2dense( moment, d.actuator_moment, d.moment_rownnz, d.moment_rowadr, d.moment_colind, ) fields['actuator_moment'] = moment contact, contact_map = _make_contact(d.contact, dim, efc_address) # pad efc fields: MuJoCo efc arrays are sparse for inactive constraints. # efc_J is also optionally column-sparse (typically for large nv). MJX is # neither: it contains zeros for inactive constraints, and efc_J is always # (nefc, nv). this may change in the future. if mujoco.mj_isSparse(m): efc_j = np.zeros((d.efc_J_rownnz.shape[0], m.nv)) mujoco.mju_sparse2dense( efc_j, fields['efc_J'], d.efc_J_rownnz, d.efc_J_rowadr, d.efc_J_colind, ) fields['efc_J'] = efc_j else: fields['efc_J'] = fields['efc_J'].reshape((-1 if m.nv else 0, m.nv)) # move efc rows to their correct offsets for fname in ( 'efc_J', 'efc_pos', 'efc_margin', 'efc_frictionloss', 'efc_D', 'efc_aref', 'efc_force', ): value = np.zeros((nefc, m.nv)) if fname == 'efc_J' else np.zeros(nefc) for i in range(3): value_beg = sum([ne, nf][:i]) d_beg = sum([d.ne, d.nf][:i]) size = [d.ne, d.nf, d.nl][i] value[value_beg : value_beg + size] = fields[fname][d_beg : d_beg + size] # for nc, we may reorder contacts so they match MJX order: group by dim for id_to, id_from in enumerate(contact_map): if id_from == -1: continue num_rows = dim[id_to] if num_rows > 1 and m.opt.cone == mujoco.mjtCone.mjCONE_PYRAMIDAL: num_rows = (num_rows - 1) * 2 efc_i, efc_o = d.contact.efc_address[id_from], efc_address[id_to] if efc_i == -1: continue value[efc_o : efc_o + num_rows] = fields[fname][efc_i : efc_i + num_rows] fields[fname] = value 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' and isinstance( fields[f.name], np.ndarray ): fields[f.name] = np.zeros(0, dtype=fields[f.name].dtype) # convert qM and qLD if jacobian is dense if not support.is_sparse(m): fields['qM'] = np.zeros((m.nv, m.nv)) mujoco.mj_fullM(m, fields['qM'], d.qM) # TODO(erikfrey): derive L*L' from L'*D*L instead of recomputing try: fields['qLD'], _ = scipy.linalg.cho_factor(fields['qM']) except scipy.linalg.LinAlgError: # this happens when qM is empty or unstable simulation fields['qLD'] = np.zeros((m.nv, m.nv)) fields['qLDiagInv'] = np.zeros(0) fields['contact'] = contact fields.update(ne=ne, nf=nf, nl=nl, nefc=nefc, ncon=ncon, efc_type=efc_type) # copy because device_put is async: data = types.Data(**{k: copy.copy(v) for k, v in fields.items()}) data = jax.device_put(data, device=device) return _strip_weak_type(data)