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

Internal change.

Browse Source 
PiperOrigin-RevId: 704273388
Change-Id: I6de2aa5bf80353dae4f2faf408c4b4c1c2e14f15
This commit is contained in:
Google DeepMind
2024-12-09 07:07:34 -08:00
committed by Copybara-Service
parent f7f1c84a1e
commit b8344c191f
1 changed files with 28 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/io.py
+28 -22
View File
@@ -37,28 +37,32 @@ def _strip_weak_type(tree):
return jax.tree_util.tree_map(f, tree)
def _make_option(o: mujoco.MjOption) -> types.Option:
def _make_option(
o: mujoco.MjOption, _full_compat: bool = False
) -> types.Option:
"""Returns mjx.Option given mujoco.MjOption."""
if o.integrator not in set(types.IntegratorType):
raise NotImplementedError(f'{mujoco.mjtIntegrator(o.integrator)}')
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.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.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)}')
if o.solver not in set(types.SolverType):
raise NotImplementedError(f'{mujoco.mjtSolver(o.solver)}')
for i in range(mujoco.mjtEnableBit.mjNENABLE):
if o.enableflags & 2**i:
raise NotImplementedError(f'{mujoco.mjtEnableBit(2 ** i)}')
for i in range(mujoco.mjtEnableBit.mjNENABLE):
if o.enableflags & 2**i:
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 implicitfast and has_fluid_params:
raise NotImplementedError('implicitfast not implemented for fluid drag.')
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)
@@ -182,16 +186,17 @@ def put_model(
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)
fields['opt'] = _make_option(m.opt, _full_compat=_full_compat)
fields['stat'] = _make_statistic(m.stat)
# Pre-compile meshes for MJX collisions.
fields['mesh_convex'] = [None] * m.nmesh
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'])
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()})
@@ -496,9 +501,10 @@ 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,):
if restricted_to in ('mujoco', 'mjx'):
continue # don't copy fields that are mujoco-only or MJX-only
getattr(result_i, field.name)[:] = value
else:
getattr(result_i, field.name)[:] = value
else:
setattr(result_i, field.name, value)