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

Rename backend_impl to impl. Remove _full_compat from callsites.

Browse Source 
PiperOrigin-RevId: 770400560
Change-Id: I7bec79aeae07322a30b15cc04d4dd05077741a8b
This commit is contained in:
Baruch Tabanpour
2025-06-11 18:04:51 -07:00
committed by Copybara-Service
parent f978b4eea3
commit 65deedbc75
4 changed files with 217 additions and 217 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/io.py
+93 -93
View File
@@ -41,51 +41,51 @@ def _is_cuda_gpu_device(device: jax.Device) -> bool:
return device in cuda_devices return device in cuda_devices
def _resolve_backend_impl( def _resolve_impl(
device: jax.Device, device: jax.Device,
) -> types.BackendImpl: ) -> types.Impl:
"""Pick a default backend impl based on the device specified.""" """Pick a default implementation based on the device specified."""
if _is_cuda_gpu_device(device): if _is_cuda_gpu_device(device):
# TODO(btaba): Remove flag once Warp is ready to launch. # TODO(btaba): Remove flag once Warp is ready to launch.
mjx_warp_enabled = os.environ.get('MJX_WARP_ENABLED', 'f').lower() == 'true' mjx_warp_enabled = os.environ.get('MJX_WARP_ENABLED', 'f').lower() == 'true'
if mjx_warp_enabled: if mjx_warp_enabled:
logging.debug('Picking default backend implementation: Warp.') logging.debug('Picking default implementation: Warp.')
return types.BackendImpl.WARP return types.Impl.WARP
logging.info('MJX Warp is disabled via MJX_WARP_ENABLED=false.') logging.info('MJX Warp is disabled via MJX_WARP_ENABLED=false.')
if device.platform in ('gpu', 'tpu'): if device.platform in ('gpu', 'tpu'):
logging.debug('Picking default backend implementation: JAX.') logging.debug('Picking default implementation: JAX.')
return types.BackendImpl.JAX return types.Impl.JAX
if device.platform == 'cpu': if device.platform == 'cpu':
mjx_c_default = ( mjx_c_default = (
os.environ.get('MJX_C_DEFAULT_ENABLED', 'f').lower() == 'true' os.environ.get('MJX_C_DEFAULT_ENABLED', 'f').lower() == 'true'
) )
if mjx_c_default: if mjx_c_default:
logging.debug('Picking default backend implementation: C.') logging.debug('Picking default implementation: C.')
return types.BackendImpl.C return types.Impl.C
return types.BackendImpl.JAX return types.Impl.JAX
raise ValueError(f'Unsupported device: {device}') raise ValueError(f'Unsupported device: {device}')
def _resolve_device( def _resolve_device(
backend_impl: types.BackendImpl, impl: types.Impl,
) -> jax.Device: ) -> jax.Device:
"""Resolves a device based on the backend implementation.""" """Resolves a device based on the implementation."""
backend_impl = types.BackendImpl(backend_impl) impl = types.Impl(impl)
if backend_impl == types.BackendImpl.JAX: if impl == types.Impl.JAX:
device_0 = jax.devices()[0] device_0 = jax.devices()[0]
logging.debug('Picking default device: %s.', device_0) logging.debug('Picking default device: %s.', device_0)
return device_0 return device_0
if backend_impl == types.BackendImpl.C: if impl == types.Impl.C:
cpu_0 = jax.devices('cpu')[0] cpu_0 = jax.devices('cpu')[0]
logging.debug('Picking default device: %s', cpu_0) logging.debug('Picking default device: %s', cpu_0)
return cpu_0 return cpu_0
if backend_impl == types.BackendImpl.WARP: if impl == types.Impl.WARP:
# WARP backend requires a CUDA GPU. # WARP implementation requires a CUDA GPU.
cuda_gpus = [d for d in jax.devices('cuda')] cuda_gpus = [d for d in jax.devices('cuda')]
if not cuda_gpus: if not cuda_gpus:
raise AssertionError( raise AssertionError(
@@ -96,64 +96,64 @@ def _resolve_device(
logging.debug('Picking default device: %s', cuda_gpus[0]) logging.debug('Picking default device: %s', cuda_gpus[0])
return cuda_gpus[0] return cuda_gpus[0]
raise ValueError(f'Unsupported backend implementation: {backend_impl}') raise ValueError(f'Unsupported implementation: {impl}')
def _check_backend_impl_device_compatibility( def _check_impl_device_compatibility(
backend_impl: Union[str, types.BackendImpl], impl: Union[str, types.Impl],
device: jax.Device, device: jax.Device,
) -> None: ) -> None:
"""Checks that the backend implementation is compatible with the device.""" """Checks that the implementation is compatible with the device."""
if backend_impl is None: if impl is None:
raise ValueError('No backend implementation specified.') raise ValueError('No implementation specified.')
backend_impl = types.BackendImpl(backend_impl) impl = types.Impl(impl)
if backend_impl == types.BackendImpl.WARP: if impl == types.Impl.WARP:
if not _is_cuda_gpu_device(device): if not _is_cuda_gpu_device(device):
raise AssertionError( raise AssertionError(
'Warp backend implementation requires a CUDA GPU device, got ' 'Warp implementation requires a CUDA GPU device, got '
f'{device}.' f'{device}.'
) )
mjx_warp_enabled = os.environ.get('MJX_WARP_ENABLED', 'f').lower() == 'true' mjx_warp_enabled = os.environ.get('MJX_WARP_ENABLED', 'f').lower() == 'true'
if not mjx_warp_enabled: if not mjx_warp_enabled:
raise AssertionError( raise AssertionError(
'Warp backend implementation is disabled via MJX_WARP_ENABLED=false.' 'Warp implementation is disabled via MJX_WARP_ENABLED=false.'
) )
is_cpu_device = device.platform == 'cpu' is_cpu_device = device.platform == 'cpu'
if backend_impl == types.BackendImpl.C: if impl == types.Impl.C:
if not is_cpu_device: if not is_cpu_device:
raise AssertionError( raise AssertionError(
f'C backend implementation requires a CPU device, got {device}.' f'C implementation requires a CPU device, got {device}.'
) )
# NB: JAX backend works with any device. # NB: JAX implementation works with any device.
def _resolve_backend_impl_and_device( def _resolve_impl_and_device(
backend_impl: Optional[Union[str, types.BackendImpl]], impl: Optional[Union[str, types.Impl]],
device: Optional[jax.Device] = None, device: Optional[jax.Device] = None,
) -> Tuple[types.BackendImpl, jax.Device]: ) -> Tuple[types.Impl, jax.Device]:
"""Resolves a backend implementation and device.""" """Resolves a implementation and device."""
if backend_impl: if impl:
backend_impl = types.BackendImpl(backend_impl) impl = types.Impl(impl)
has_backend_impl, has_device = backend_impl is not None, device is not None has_impl, has_device = impl is not None, device is not None
if (has_backend_impl, has_device) == (True, True): if (has_impl, has_device) == (True, True):
pass pass
elif (has_backend_impl, has_device) == (True, False): elif (has_impl, has_device) == (True, False):
device = _resolve_device(backend_impl) device = _resolve_device(impl)
elif (has_backend_impl, has_device) == (False, True): elif (has_impl, has_device) == (False, True):
backend_impl = _resolve_backend_impl(device) impl = _resolve_impl(device)
else: else:
device = jax.devices(jax.default_backend())[0] device = jax.devices(jax.default_backend())[0]
logging.info('Using JAX default device: %s.', device) logging.info('Using JAX default device: %s.', device)
backend_impl = _resolve_backend_impl(device) impl = _resolve_impl(device)
_check_backend_impl_device_compatibility(backend_impl, device) _check_impl_device_compatibility(impl, device)
return backend_impl, device # pytype: disable=bad-return-type return impl, device # pytype: disable=bad-return-type
def _strip_weak_type(tree): def _strip_weak_type(tree):
@@ -167,7 +167,7 @@ def _strip_weak_type(tree):
def _put_option( def _put_option(
o: mujoco.MjOption, o: mujoco.MjOption,
backend_impl: types.BackendImpl, impl: types.Impl,
impl_fields: Optional[dict[str, Any]] = None, impl_fields: Optional[dict[str, Any]] = None,
) -> types.Option: ) -> types.Option:
"""Returns mjx.Option given mujoco.MjOption.""" """Returns mjx.Option given mujoco.MjOption."""
@@ -195,7 +195,7 @@ def _put_option(
fields['disableflags'] = types.DisableBit(o.disableflags) fields['disableflags'] = types.DisableBit(o.disableflags)
fields['enableflags'] = types.EnableBit(o.enableflags) fields['enableflags'] = types.EnableBit(o.enableflags)
if backend_impl == types.BackendImpl.JAX: if impl == types.Impl.JAX:
has_fluid_params = o.density > 0 or o.viscosity > 0 or o.wind.any() has_fluid_params = o.density > 0 or o.viscosity > 0 or o.wind.any()
implicitfast = o.integrator == mujoco.mjtIntegrator.mjINT_IMPLICITFAST implicitfast = o.integrator == mujoco.mjtIntegrator.mjINT_IMPLICITFAST
if implicitfast and has_fluid_params: if implicitfast and has_fluid_params:
@@ -203,12 +203,12 @@ def _put_option(
fields['has_fluid_params'] = has_fluid_params fields['has_fluid_params'] = has_fluid_params
return types.OptionJAX(**fields, **(impl_fields or {})) return types.OptionJAX(**fields, **(impl_fields or {}))
if backend_impl == types.BackendImpl.C: if impl == types.Impl.C:
c_field_keys = types.OptionC.__annotations__.keys() - fields.keys() c_field_keys = types.OptionC.__annotations__.keys() - fields.keys()
c_fields = {k: getattr(o, k, None) for k in c_field_keys} c_fields = {k: getattr(o, k, None) for k in c_field_keys}
return types.OptionC(**fields, **c_fields, **(impl_fields or {})) return types.OptionC(**fields, **c_fields, **(impl_fields or {}))
raise NotImplementedError(f'Unsupported backend: {backend_impl}') raise NotImplementedError(f'Unsupported implementation: {impl}')
def _put_statistic(s: mujoco.MjStatistic) -> types.Statistic: def _put_statistic(s: mujoco.MjStatistic) -> types.Statistic:
@@ -283,7 +283,7 @@ def _put_model_jax(
mj_field_names = {f.name for f in types.Model.fields() if f.name != '_impl'} mj_field_names = {f.name for f in types.Model.fields() if f.name != '_impl'}
fields = {f: getattr(m, f) for f in mj_field_names} fields = {f: getattr(m, f) for f in mj_field_names}
fields['cam_mat0'] = fields['cam_mat0'].reshape((-1, 3, 3)) fields['cam_mat0'] = fields['cam_mat0'].reshape((-1, 3, 3))
fields['opt'] = _put_option(m.opt, types.BackendImpl.JAX) fields['opt'] = _put_option(m.opt, types.Impl.JAX)
fields['stat'] = _put_statistic(m.stat) fields['stat'] = _put_statistic(m.stat)
fields_jax = {} fields_jax = {}
@@ -340,7 +340,7 @@ def _put_model_c(
mj_field_names = {f.name for f in types.Model.fields() if f.name != '_impl'} mj_field_names = {f.name for f in types.Model.fields() if f.name != '_impl'}
fields = {f: getattr(m, f) for f in mj_field_names} fields = {f: getattr(m, f) for f in mj_field_names}
fields['cam_mat0'] = fields['cam_mat0'].reshape((-1, 3, 3)) fields['cam_mat0'] = fields['cam_mat0'].reshape((-1, 3, 3))
fields['opt'] = _put_option(m.opt, backend_impl=types.BackendImpl.C) fields['opt'] = _put_option(m.opt, impl=types.Impl.C)
fields['stat'] = _put_statistic(m.stat) fields['stat'] = _put_statistic(m.stat)
c_impl_keys = ( c_impl_keys = (
@@ -359,7 +359,7 @@ def _put_model_c(
def put_model( def put_model(
m: mujoco.MjModel, m: mujoco.MjModel,
device: Optional[jax.Device] = None, device: Optional[jax.Device] = None,
backend_impl: Optional[Union[str, types.BackendImpl]] = None, impl: Optional[Union[str, types.Impl]] = None,
_full_compat: bool = False, # pylint: disable=invalid-name _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.
@@ -367,7 +367,7 @@ def put_model(
Args: Args:
m: the model to put onto device m: the model to put onto device
device: which device to use - if unspecified picks the default device device: which device to use - if unspecified picks the default device
backend_impl: backend implementation to use impl: implementation to use
_full_compat: put all MjModel fields onto device irrespective of MJX support _full_compat: put all MjModel fields onto device irrespective of MJX support
This is an experimental feature. Avoid using it for now. This is an experimental feature. Avoid using it for now.
@@ -375,28 +375,29 @@ def put_model(
an mjx.Model placed on device an mjx.Model placed on device
Raises: Raises:
ValueError: if backend_impl is not supported ValueError: if impl is not supported
DeprecationWarning: if _full_compat is True DeprecationWarning: if _full_compat is True
""" """
if _full_compat: if _full_compat:
warnings.warn( warnings.warn(
'mjx.put_model(..., _full_compat=True) is deprecated. Use' 'mjx.put_model(..., _full_compat=True) is deprecated and will be'
' mjx.put_model(..., backend_impl=types.BackendImpl.C) instead.', ' removed in MuJoCo >=3.4. Use mjx.put_model(..., impl=types.Impl.C)'
' instead.',
DeprecationWarning, DeprecationWarning,
stacklevel=2, stacklevel=2,
) )
backend_impl = types.BackendImpl.C impl = types.Impl.C
backend_impl, device = _resolve_backend_impl_and_device(backend_impl, device) impl, device = _resolve_impl_and_device(impl, device)
if backend_impl == types.BackendImpl.JAX: if impl == types.Impl.JAX:
return _put_model_jax(m, device) return _put_model_jax(m, device)
elif backend_impl == types.BackendImpl.C: elif impl == types.Impl.C:
return _put_model_c(m, device) return _put_model_c(m, device)
elif backend_impl == types.BackendImpl.WARP: elif impl == types.Impl.WARP:
raise NotImplementedError('Warp backend not implemented yet.') raise NotImplementedError('Warp implementation not implemented yet.')
else: else:
raise ValueError(f'Unsupported backend implementation: {backend_impl}') raise ValueError(f'Unsupported implementation: {impl}')
def _make_data_public_fields(m: types.Model) -> Dict[str, Any]: def _make_data_public_fields(m: types.Model) -> Dict[str, Any]:
@@ -696,7 +697,7 @@ def _make_data_c(
def make_data( def make_data(
m: Union[types.Model, mujoco.MjModel], m: Union[types.Model, mujoco.MjModel],
device: Optional[jax.Device] = None, device: Optional[jax.Device] = None,
backend_impl: Optional[Union[str, types.BackendImpl]] = None, impl: Optional[Union[str, types.Impl]] = None,
_full_compat: bool = False, # pylint: disable=invalid-name _full_compat: bool = False, # pylint: disable=invalid-name
) -> types.Data: ) -> types.Data:
"""Allocate and initialize Data. """Allocate and initialize Data.
@@ -704,7 +705,7 @@ def make_data(
Args: Args:
m: the model to use m: the model to use
device: which device to use - if unspecified picks the default device device: which device to use - if unspecified picks the default device
backend_impl: backend implementation to use impl: implementation to use ('jax', 'warp')
_full_compat: put all fields onto device irrespective of MJX support This is _full_compat: put all fields onto device irrespective of MJX support This is
an experimental feature. Avoid using it for now. If using this flag, also an experimental feature. Avoid using it for now. If using this flag, also
use _full_compat for put_model. use _full_compat for put_model.
@@ -713,35 +714,34 @@ def make_data(
an initialized mjx.Data placed on device an initialized mjx.Data placed on device
Raises: Raises:
ValueError: if the model's backend_impl does not match the make_data ValueError: if the model's impl does not match the make_data impl
backend_impl NotImplementedError: if the impl is not implemented yet
NotImplementedError: if the backend_impl is not implemented yet
DeprecationWarning: if _full_compat is used DeprecationWarning: if _full_compat is used
""" """
if _full_compat: if _full_compat:
warnings.warn( warnings.warn(
'mjx.make_data(..., _full_compat=True) is deprecated. Use' 'mjx.make_data(..., _full_compat=True) is deprecated. Use'
' mjx.make_data(..., backend_impl=types.BackendImpl.C) instead.', ' mjx.make_data(..., impl=types.Impl.C) instead.',
DeprecationWarning, DeprecationWarning,
stacklevel=2, stacklevel=2,
) )
backend_impl = types.BackendImpl.C impl = types.Impl.C
backend_impl, device = _resolve_backend_impl_and_device(backend_impl, device) impl, device = _resolve_impl_and_device(impl, device)
if isinstance(m, types.Model) and m.backend_impl != backend_impl: if isinstance(m, types.Model) and m.impl != impl:
raise ValueError( raise ValueError(
f'Model backend_impl {m.backend_impl} does not match make_data ' f'Model impl {m.impl} does not match make_data '
f'backend_impl {backend_impl}.' f'implementation {impl}.'
) )
if backend_impl == types.BackendImpl.JAX: if impl == types.Impl.JAX:
return _make_data_jax(m, device) return _make_data_jax(m, device)
elif backend_impl == types.BackendImpl.C: elif impl == types.Impl.C:
return _make_data_c(m, device) return _make_data_c(m, device)
raise NotImplementedError( raise NotImplementedError(
f'make_data for backend_impl "{backend_impl}" not implemented yet.' f'make_data for implementation "{impl}" not implemented yet.'
) )
@@ -951,7 +951,7 @@ def _put_data_c(
if hasattr(d, f.name) if hasattr(d, f.name)
} }
# TODO(stunya): support islanding via C backend impl. # TODO(stunya): support islanding via C impl.
impl_fields['solver_niter'] = impl_fields['solver_niter'][0] impl_fields['solver_niter'] = impl_fields['solver_niter'][0]
# TODO(btaba): remove dense actuator moment. # TODO(btaba): remove dense actuator moment.
@@ -1039,7 +1039,7 @@ def put_data(
m: mujoco.MjModel, m: mujoco.MjModel,
d: mujoco.MjData, d: mujoco.MjData,
device: Optional[jax.Device] = None, device: Optional[jax.Device] = None,
backend_impl: Optional[Union[str, types.BackendImpl]] = None, impl: Optional[Union[str, types.Impl]] = None,
_full_compat: bool = False, # pylint: disable=invalid-name _full_compat: bool = False, # pylint: disable=invalid-name
) -> types.Data: ) -> types.Data:
"""Puts mujoco.MjData onto a device, resulting in mjx.Data. """Puts mujoco.MjData onto a device, resulting in mjx.Data.
@@ -1048,7 +1048,7 @@ def put_data(
m: the model to use m: the model to use
d: the data to put on device d: the data to put on device
device: which device to use - if unspecified picks the default device device: which device to use - if unspecified picks the default device
backend_impl: backend implementation to use impl: implementation to use ('jax', 'warp')
_full_compat: put all MjModel fields onto device irrespective of MJX support _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 This is an experimental feature. Avoid using it for now. If using this
flag, also use _full_compat for put_model. flag, also use _full_compat for put_model.
@@ -1059,20 +1059,20 @@ def put_data(
if _full_compat: if _full_compat:
warnings.warn( warnings.warn(
'mjx.put_data(..., _full_compat=True) is deprecated. Use' 'mjx.put_data(..., _full_compat=True) is deprecated. Use'
' mjx.put_data(..., backend_impl=types.BackendImpl.C) instead.', ' mjx.put_data(..., impl=types.Impl.C) instead.',
DeprecationWarning, DeprecationWarning,
stacklevel=2, stacklevel=2,
) )
backend_impl = types.BackendImpl.C impl = types.Impl.C
backend_impl, device = _resolve_backend_impl_and_device(backend_impl, device) impl, device = _resolve_impl_and_device(impl, device)
if backend_impl == types.BackendImpl.JAX: if impl == types.Impl.JAX:
return _put_data_jax(m, d, device) return _put_data_jax(m, d, device)
elif backend_impl == types.BackendImpl.C: elif impl == types.Impl.C:
return _put_data_c(m, d, device) return _put_data_c(m, d, device)
raise NotImplementedError( raise NotImplementedError(
f'put_data for backend_impl "{backend_impl}" not implemented yet.' f'put_data for implementation "{impl}" not implemented yet.'
) )
@@ -1097,7 +1097,7 @@ def _get_data_into(
batch_size = d.qpos.shape[0] if batched else 1 batch_size = d.qpos.shape[0] if batched else 1
dof_i, dof_j = [], [] dof_i, dof_j = [], []
if d.backend_impl == types.BackendImpl.JAX: if d.impl == types.Impl.JAX:
for i in range(m.nv): for i in range(m.nv):
j = i j = i
while j > -1: while j > -1:
@@ -1116,13 +1116,13 @@ def _get_data_into(
if ncon != result_i.ncon or nefc != result_i.nefc or nj != result_i.nJ: if ncon != result_i.ncon or nefc != result_i.nefc or nj != result_i.nJ:
mujoco._functions._realloc_con_efc(result_i, ncon=ncon, nefc=nefc, nJ=nj) # pylint: disable=protected-access mujoco._functions._realloc_con_efc(result_i, ncon=ncon, nefc=nefc, nJ=nj) # pylint: disable=protected-access
if d.backend_impl == types.BackendImpl.JAX: if d.impl == types.Impl.JAX:
all_fields = types.Data.fields() + types.DataJAX.fields() all_fields = types.Data.fields() + types.DataJAX.fields()
elif d.backend_impl == types.BackendImpl.C: elif d.impl == types.Impl.C:
all_fields = types.Data.fields() + types.DataC.fields() all_fields = types.Data.fields() + types.DataC.fields()
else: else:
raise NotImplementedError( raise NotImplementedError(
f'get_data_into for backend_impl "{d.backend_impl}" not implemented' f'get_data_into for implementation "{d.impl}" not implemented'
' yet.' ' yet.'
) )
@@ -1188,7 +1188,7 @@ def _get_data_into(
value = value.reshape(-1) value = value.reshape(-1)
elif field.name.startswith('efc_'): elif field.name.startswith('efc_'):
value = value[efc_active] value = value[efc_active]
if d.backend_impl == types.BackendImpl.JAX: if d.impl == types.Impl.JAX:
if field.name == 'qM' and not support.is_sparse(m): if field.name == 'qM' and not support.is_sparse(m):
value = value[dof_i, dof_j] value = value[dof_i, dof_j]
elif field.name == 'qLD' and not support.is_sparse(m): elif field.name == 'qLD' and not support.is_sparse(m):
@@ -1226,12 +1226,12 @@ def get_data_into(
d = jax.device_get(d) d = jax.device_get(d)
if d.backend_impl in (types.BackendImpl.JAX, types.BackendImpl.C): if d.impl in (types.Impl.JAX, types.Impl.C):
# TODO(stunya): Split out _get_data_into once codepaths diverge enough. # TODO(stunya): Split out _get_data_into once codepaths diverge enough.
return _get_data_into(result, m, d) return _get_data_into(result, m, d)
raise NotImplementedError( raise NotImplementedError(
f'get_data_into for backend_impl "{d.backend_impl}" not implemented yet.' f'get_data_into for implementation "{d.impl}" not implemented yet.'
) )
mjx/mujoco/mjx/_src/io_test.py
+114 -114
View File
@@ -26,8 +26,8 @@ from mujoco.mjx._src import io as mjx_io
from mujoco.mjx._src import test_util from mujoco.mjx._src import test_util
# pylint: disable=g-importing-member # pylint: disable=g-importing-member
from mujoco.mjx._src.types import BackendImpl
from mujoco.mjx._src.types import ConeType from mujoco.mjx._src.types import ConeType
from mujoco.mjx._src.types import Impl
# pylint: enable=g-importing-member # pylint: enable=g-importing-member
import numpy as np import numpy as np
@@ -114,11 +114,11 @@ class ModelIOTest(parameterized.TestCase):
@parameterized.product( @parameterized.product(
xml=(_MULTIPLE_CONVEX_OBJECTS, _MULTIPLE_CONSTRAINTS), xml=(_MULTIPLE_CONVEX_OBJECTS, _MULTIPLE_CONSTRAINTS),
backend_impl=('jax', 'c'), impl=('jax', 'c'),
) )
def test_put_model(self, xml, backend_impl): def test_put_model(self, xml, impl):
m = mujoco.MjModel.from_xml_string(xml) m = mujoco.MjModel.from_xml_string(xml)
mx = mjx.put_model(m, backend_impl=backend_impl) mx = mjx.put_model(m, impl=impl)
def assert_not_weak_type(x): def assert_not_weak_type(x):
if isinstance(x, jax.Array): if isinstance(x, jax.Array):
@@ -140,10 +140,10 @@ class ModelIOTest(parameterized.TestCase):
self.assertEqual(mx.nM, m.nM) self.assertEqual(mx.nM, m.nM)
self.assertAlmostEqual(mx.opt.timestep, m.opt.timestep) self.assertAlmostEqual(mx.opt.timestep, m.opt.timestep)
if backend_impl == 'jax': if impl == 'jax':
# fields restricted to MuJoCo should not be populated # fields restricted to MuJoCo should not be populated
self.assertFalse(hasattr(mx, 'bvh_aabb')) self.assertFalse(hasattr(mx, 'bvh_aabb'))
elif backend_impl == 'c': elif impl == 'c':
# Options specific to C are populated. # Options specific to C are populated.
self.assertEqual(mx.opt.apirate, m.opt.apirate) self.assertEqual(mx.opt.apirate, m.opt.apirate)
# Fields private to C backend impl are populated. # Fields private to C backend impl are populated.
@@ -177,7 +177,7 @@ class ModelIOTest(parameterized.TestCase):
mujoco.MjModel.from_xml_string( mujoco.MjModel.from_xml_string(
'<mujoco><option viscosity="3.0"/><worldbody/></mujoco>' '<mujoco><option viscosity="3.0"/><worldbody/></mujoco>'
), ),
backend_impl='jax', impl='jax',
) )
self.assertTrue(m.opt.has_fluid_params) self.assertTrue(m.opt.has_fluid_params)
@@ -218,7 +218,7 @@ class ModelIOTest(parameterized.TestCase):
</body> </body>
</worldbody> </worldbody>
</mujoco>"""), </mujoco>"""),
backend_impl='jax', impl='jax',
) )
def test_implicitfast_fluid_not_implemented(self): def test_implicitfast_fluid_not_implemented(self):
@@ -229,18 +229,18 @@ class ModelIOTest(parameterized.TestCase):
<option viscosity="3.0" integrator="implicitfast"/> <option viscosity="3.0" integrator="implicitfast"/>
<worldbody/> <worldbody/>
</mujoco>"""), </mujoco>"""),
backend_impl='jax', impl='jax',
) )
def test_wrap_inside(self): def test_wrap_inside(self):
m = test_util.load_test_file('tendon/wrap_sidesite.xml') m = test_util.load_test_file('tendon/wrap_sidesite.xml')
mx0 = mjx.put_model(m, backend_impl='jax') mx0 = mjx.put_model(m, impl='jax')
np.testing.assert_equal( np.testing.assert_equal(
mx0._impl.is_wrap_inside, mx0._impl.is_wrap_inside,
np.array([1, 0, 1, 0, 1, 1, 0]), np.array([1, 0, 1, 0, 1, 1, 0]),
) )
m.site_pos[2] = m.site_pos[1] m.site_pos[2] = m.site_pos[1]
mx1 = mjx.put_model(m, backend_impl='jax') mx1 = mjx.put_model(m, impl='jax')
np.testing.assert_equal( np.testing.assert_equal(
mx1._impl.is_wrap_inside, mx1._impl.is_wrap_inside,
np.array([0, 0, 1, 0, 1, 0, 0]), np.array([0, 0, 1, 0, 1, 0, 0]),
@@ -251,10 +251,10 @@ class DataIOTest(parameterized.TestCase):
"""IO tests for mjx.Data.""" """IO tests for mjx.Data."""
@parameterized.parameters('jax', 'c') @parameterized.parameters('jax', 'c')
def test_make_data(self, backend_impl: str): def test_make_data(self, impl: str):
"""Test that make_data returns the correct shapes.""" """Test that make_data returns the correct shapes."""
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS) m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS)
d = mjx.make_data(m, backend_impl=backend_impl) d = mjx.make_data(m, impl=impl)
nq = 22 nq = 22
nbody = 5 nbody = 5
@@ -312,34 +312,34 @@ class DataIOTest(parameterized.TestCase):
self.assertEqual(d.qfrc_inverse.shape, (nv,)) self.assertEqual(d.qfrc_inverse.shape, (nv,))
self.assertEqual(d._impl.efc_force.shape, (nefc,)) self.assertEqual(d._impl.efc_force.shape, (nefc,))
if backend_impl == 'jax': if impl == 'jax':
self.assertEqual(d._impl.qM.shape, (nv, nv)) self.assertEqual(d._impl.qM.shape, (nv, nv))
self.assertEqual(d._impl.qLD.shape, (nv, nv)) self.assertEqual(d._impl.qLD.shape, (nv, nv))
self.assertEqual(d._impl.qLDiagInv.shape, (0,)) self.assertEqual(d._impl.qLDiagInv.shape, (0,))
elif backend_impl == 'c': elif impl == 'c':
self.assertEqual(d._impl.qM.shape, (nm,)) self.assertEqual(d._impl.qM.shape, (nm,))
self.assertEqual(d._impl.qLD.shape, (nm,)) self.assertEqual(d._impl.qLD.shape, (nm,))
self.assertEqual(d._impl.qLDiagInv.shape, (nv,)) self.assertEqual(d._impl.qLDiagInv.shape, (nv,))
# test sparse # test sparse
m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
d = mjx.make_data(m, backend_impl=backend_impl) d = mjx.make_data(m, impl=impl)
self.assertEqual(d._impl.qM.shape, (nm,)) self.assertEqual(d._impl.qM.shape, (nm,))
self.assertEqual(d._impl.qLD.shape, (nm,)) self.assertEqual(d._impl.qLD.shape, (nm,))
self.assertEqual(d._impl.qLDiagInv.shape, (nv,)) self.assertEqual(d._impl.qLDiagInv.shape, (nv,))
if backend_impl == 'c': if impl == 'c':
# check C specific fields # check C specific fields
self.assertEqual(d._impl.light_xpos.shape, (m.nlight, 3)) self.assertEqual(d._impl.light_xpos.shape, (m.nlight, 3))
self.assertEqual(d._impl.bvh_active.shape, (m.nbvh,)) self.assertEqual(d._impl.bvh_active.shape, (m.nbvh,))
@parameterized.parameters('jax', 'c') @parameterized.parameters('jax', 'c')
def test_put_data(self, backend_impl: str): def test_put_data(self, impl: str):
"""Test that put_data puts the correct data for dense and sparse.""" """Test that put_data puts the correct data for dense and sparse."""
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
d = mujoco.MjData(m) d = mujoco.MjData(m)
mujoco.mj_step(m, d, 2) mujoco.mj_step(m, d, 2)
dx = mjx.put_data(m, d, backend_impl=backend_impl) dx = mjx.put_data(m, d, impl=impl)
# check a few fields # check a few fields
np.testing.assert_allclose(dx.qpos, d.qpos) np.testing.assert_allclose(dx.qpos, d.qpos)
@@ -356,12 +356,12 @@ class DataIOTest(parameterized.TestCase):
) )
) )
if backend_impl == 'jax': if impl == 'jax':
# check that qM is transformed properly # check that qM is transformed properly
qm = np.zeros((m.nv, m.nv), dtype=np.float64) qm = np.zeros((m.nv, m.nv), dtype=np.float64)
mujoco.mj_fullM(m, qm, d.qM) mujoco.mj_fullM(m, qm, d.qM)
np.testing.assert_allclose(qm, mjx.full_m(mjx.put_model(m), dx)) np.testing.assert_allclose(qm, mjx.full_m(mjx.put_model(m), dx))
elif backend_impl == 'c': elif impl == 'c':
np.testing.assert_allclose(dx._impl.qM, d.qM) np.testing.assert_allclose(dx._impl.qM, d.qM)
np.testing.assert_allclose(dx._impl.qLD, d.qLD) np.testing.assert_allclose(dx._impl.qLD, d.qLD)
np.testing.assert_allclose(dx._impl.qLDiagInv, d.qLDiagInv) np.testing.assert_allclose(dx._impl.qLDiagInv, d.qLDiagInv)
@@ -417,7 +417,7 @@ class DataIOTest(parameterized.TestCase):
m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
d = mujoco.MjData(m) d = mujoco.MjData(m)
mujoco.mj_step(m, d, 2) mujoco.mj_step(m, d, 2)
dx_sparse = mjx.put_data(m, d, backend_impl=backend_impl) dx_sparse = mjx.put_data(m, d, impl=impl)
np.testing.assert_allclose(dx_sparse._impl.efc_J, dx._impl.efc_J, atol=1e-8) np.testing.assert_allclose(dx_sparse._impl.efc_J, dx._impl.efc_J, atol=1e-8)
# check sparse mass matrices are correct # check sparse mass matrices are correct
@@ -431,25 +431,25 @@ class DataIOTest(parameterized.TestCase):
m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE
d = mujoco.MjData(m) d = mujoco.MjData(m)
mujoco.mj_step(m, d, 2) mujoco.mj_step(m, d, 2)
dx_from_dense = mjx.put_data(m, d, backend_impl=backend_impl) dx_from_dense = mjx.put_data(m, d, impl=impl)
if backend_impl == 'jax': if impl == 'jax':
qm = np.zeros((m.nv, m.nv)) qm = np.zeros((m.nv, m.nv))
mujoco.mj_fullM(m, qm, d.qM) mujoco.mj_fullM(m, qm, d.qM)
np.testing.assert_allclose(dx_from_dense._impl.qM, qm, atol=1e-8) np.testing.assert_allclose(dx_from_dense._impl.qM, qm, atol=1e-8)
elif backend_impl == 'c': elif impl == 'c':
np.testing.assert_allclose(dx_from_dense._impl.qM, d.qM, atol=1e-8) np.testing.assert_allclose(dx_from_dense._impl.qM, d.qM, atol=1e-8)
@parameterized.parameters( @parameterized.parameters(
('jax', False), ('jax', True), ('c', False), ('c', True) ('jax', False), ('jax', True), ('c', False), ('c', True)
) )
def test_get_data(self, backend_impl: str, sparse: bool): def test_get_data(self, impl: str, sparse: bool):
"""Test that get_data makes correct MjData.""" """Test that get_data makes correct MjData."""
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
if sparse: if sparse:
m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
d = mujoco.MjData(m) d = mujoco.MjData(m)
mujoco.mj_step(m, d, 2) mujoco.mj_step(m, d, 2)
dx = mjx.put_data(m, d, backend_impl=backend_impl) dx = mjx.put_data(m, d, impl=impl)
d_2: mujoco.MjData = mjx.get_data(m, dx) d_2: mujoco.MjData = mjx.get_data(m, dx)
# check a few fields # check a few fields
@@ -505,7 +505,7 @@ class DataIOTest(parameterized.TestCase):
np.testing.assert_allclose(d_2.efc_aref, d.efc_aref) np.testing.assert_allclose(d_2.efc_aref, d.efc_aref)
np.testing.assert_allclose(d_2.contact.efc_address, d.contact.efc_address) np.testing.assert_allclose(d_2.contact.efc_address, d.contact.efc_address)
if backend_impl == 'c': if impl == 'c':
# check fields specific to the C implementation # check fields specific to the C implementation
np.testing.assert_allclose(d_2.bvh_active, d.bvh_active) np.testing.assert_allclose(d_2.bvh_active, d.bvh_active)
@@ -540,13 +540,13 @@ class DataIOTest(parameterized.TestCase):
mjx.get_data(m, dx) mjx.get_data(m, dx)
@parameterized.parameters('jax', 'c') @parameterized.parameters('jax', 'c')
def test_get_data_batched(self, backend_impl): def test_get_data_batched(self, impl):
"""Test that get_data makes correct List[MjData] for batched Data.""" """Test that get_data makes correct List[MjData] for batched Data."""
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
d = mujoco.MjData(m) d = mujoco.MjData(m)
mujoco.mj_step(m, d, 2) mujoco.mj_step(m, d, 2)
dx = mjx.put_data(m, d, backend_impl=backend_impl) dx = mjx.put_data(m, d, impl=impl)
# second data in batch has contact dist > 0, disables contact # second data in batch has contact dist > 0, disables contact
dx_b = jax.tree_util.tree_map(lambda x: jp.stack((x, x + 0.05)), dx) dx_b = jax.tree_util.tree_map(lambda x: jp.stack((x, x + 0.05)), dx)
ds = mjx.get_data(m, dx_b) ds = mjx.get_data(m, dx_b)
@@ -557,13 +557,13 @@ class DataIOTest(parameterized.TestCase):
self.assertEqual(ds[1].ncon, 0) self.assertEqual(ds[1].ncon, 0)
@parameterized.parameters('jax', 'c') @parameterized.parameters('jax', 'c')
def test_get_data_into(self, backend_impl): def test_get_data_into(self, impl):
"""Test that get_data_into correctly populates an MjData.""" """Test that get_data_into correctly populates an MjData."""
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
d = mujoco.MjData(m) d = mujoco.MjData(m)
mujoco.mj_step(m, d, 2) mujoco.mj_step(m, d, 2)
dx = mjx.put_data(m, d, backend_impl=backend_impl) dx = mjx.put_data(m, d, impl=impl)
d_2 = mujoco.MjData(m) d_2 = mujoco.MjData(m)
mjx.get_data_into(d_2, m, dx) mjx.get_data_into(d_2, m, dx)
@@ -580,32 +580,32 @@ class DataIOTest(parameterized.TestCase):
np.testing.assert_allclose(d_2.contact.frame, d.contact.frame) np.testing.assert_allclose(d_2.contact.frame, d.contact.frame)
@parameterized.parameters('jax', 'c') @parameterized.parameters('jax', 'c')
def test_get_data_into_wrong_shape(self, backend_impl): def test_get_data_into_wrong_shape(self, impl):
"""Tests that get_data_into throwsif input and output shapes don't match.""" """Tests that get_data_into throwsif input and output shapes don't match."""
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
d = mujoco.MjData(m) d = mujoco.MjData(m)
mujoco.mj_step(m, d, 2) mujoco.mj_step(m, d, 2)
dx = mjx.put_data(m, d, backend_impl=backend_impl) dx = mjx.put_data(m, d, impl=impl)
m_2 = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS) m_2 = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS)
d_2 = mujoco.MjData(m_2) d_2 = mujoco.MjData(m_2)
with self.assertRaisesRegex(ValueError, r'Input field.*has shape.*'): with self.assertRaisesRegex(ValueError, r'Input field.*has shape.*'):
mjx.get_data_into(d_2, m, dx) mjx.get_data_into(d_2, m, dx)
@parameterized.parameters('jax', 'c') @parameterized.parameters('jax', 'c')
def test_make_matches_put(self, backend_impl): def test_make_matches_put(self, impl):
"""Test that make_data produces a pytree that matches put_data.""" """Test that make_data produces a pytree that matches put_data."""
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
d = mujoco.MjData(m) d = mujoco.MjData(m)
mujoco.mj_step(m, d, 2) mujoco.mj_step(m, d, 2)
dx = mjx.put_data(m, d, backend_impl=backend_impl) dx = mjx.put_data(m, d, impl=impl)
step_fn = lambda d: d.replace(time=d.time + 1) step_fn = lambda d: d.replace(time=d.time + 1)
step_fn_jit = jax.jit(step_fn).lower(dx).compile() step_fn_jit = jax.jit(step_fn).lower(dx).compile()
# placing an MjData onto device should yield the same treedef mjx.Data as # placing an MjData onto device should yield the same treedef mjx.Data as
# calling make_data. they should be interchangeable for jax functions: # calling make_data. they should be interchangeable for jax functions:
step_fn_jit(mjx.make_data(m, backend_impl=backend_impl)) step_fn_jit(mjx.make_data(m, impl=impl))
def test_contact_elliptic_condim1(self): def test_contact_elliptic_condim1(self):
"""Test that condim=1 with ConeType.ELLIPTIC is not implemented.""" """Test that condim=1 with ConeType.ELLIPTIC is not implemented."""
@@ -654,7 +654,7 @@ class DataIOTest(parameterized.TestCase):
</mujoco> </mujoco>
""") """)
with self.assertRaises(NotImplementedError): with self.assertRaises(NotImplementedError):
mjx.put_model(m, backend_impl='jax') mjx.put_model(m, impl='jax')
class FullCompatTest(parameterized.TestCase): class FullCompatTest(parameterized.TestCase):
@@ -678,71 +678,71 @@ class FullCompatTest(parameterized.TestCase):
m = mujoco.MjModel.from_xml_string(xml) m = mujoco.MjModel.from_xml_string(xml)
with self.assertWarns(DeprecationWarning): with self.assertWarns(DeprecationWarning):
out = mjx_io.put_model(m, _full_compat=True) out = mjx_io.put_model(m, _full_compat=True)
self.assertEqual(out.backend_impl, BackendImpl.C) self.assertEqual(out.impl, Impl.C)
with self.assertWarns(DeprecationWarning): with self.assertWarns(DeprecationWarning):
out = mjx_io.make_data(m, _full_compat=True) out = mjx_io.make_data(m, _full_compat=True)
self.assertEqual(out.backend_impl, BackendImpl.C) self.assertEqual(out.impl, Impl.C)
# Test cases for `_resolve_backend_impl_and_device` where the device is # Test cases for `_resolve_impl_and_device` where the device is
# specified by the user and the device is available. # specified by the user and the device is available.
_DEVICE_TEST_CASES = [ _DEVICE_TEST_CASES = [
# Arguments use the following format: # Arguments use the following format:
# (device_type_str, backend_impl_str, # (device_type_str, impl_str,
# (expected_device, expected_backend_impl))) # (expected_device, expected_impl)))
# No backend specified. # No backend specified.
('cpu', None, ('cpu', BackendImpl.C)), ('cpu', None, ('cpu', Impl.C)),
('gpu-notnvidia', None, ('gpu', BackendImpl.JAX)), ('gpu-notnvidia', None, ('gpu', Impl.JAX)),
('gpu-nvidia', None, ('gpu', BackendImpl.WARP)), ('gpu-nvidia', None, ('gpu', Impl.WARP)),
('tpu', None, ('tpu', BackendImpl.JAX)), ('tpu', None, ('tpu', Impl.JAX)),
# JAX backend specified. # JAX backend specified.
('cpu', 'jax', ('cpu', BackendImpl.JAX)), ('cpu', 'jax', ('cpu', Impl.JAX)),
('gpu-notnvidia', 'jax', ('gpu', BackendImpl.JAX)), ('gpu-notnvidia', 'jax', ('gpu', Impl.JAX)),
('gpu-nvidia', 'jax', ('gpu', BackendImpl.JAX)), ('gpu-nvidia', 'jax', ('gpu', Impl.JAX)),
('tpu', 'jax', ('tpu', BackendImpl.JAX)), ('tpu', 'jax', ('tpu', Impl.JAX)),
# WARP backend specified. # WARP backend specified.
('cpu', 'warp', ('cpu', 'error')), ('cpu', 'warp', ('cpu', 'error')),
('gpu-notnvidia', 'warp', ('cpu', 'error')), ('gpu-notnvidia', 'warp', ('cpu', 'error')),
('gpu-nvidia', 'warp', ('gpu', BackendImpl.WARP)), ('gpu-nvidia', 'warp', ('gpu', Impl.WARP)),
('tpu', 'warp', ('tpu', 'error')), ('tpu', 'warp', ('tpu', 'error')),
# C backend specified. # C backend specified.
('cpu', 'c', ('cpu', BackendImpl.C)), ('cpu', 'c', ('cpu', Impl.C)),
('gpu-notnvidia', 'c', ('cpu', 'error')), ('gpu-notnvidia', 'c', ('cpu', 'error')),
('gpu-nvidia', 'c', ('cpu', 'error')), ('gpu-nvidia', 'c', ('cpu', 'error')),
('tpu', 'c', ('tpu', 'error')), ('tpu', 'c', ('tpu', 'error')),
] ]
# Test cases for `_resolve_backend_impl_and_device` where the user does NOT # Test cases for `_resolve_impl_and_device` where the user does NOT
# specify a device. We mock the JAX default device. # specify a device. We mock the JAX default device.
_DEFAULT_DEVICE_TEST_CASES = [ _DEFAULT_DEVICE_TEST_CASES = [
# Arguments use the following format: # Arguments use the following format:
# (jax.default_device, backend_impl_str, # (jax.default_device, impl_str,
# (expected_device, expected_backend_impl)) # (expected_device, expected_impl))
# No backend impl specified. # No backend impl specified.
('cpu', None, ('cpu', BackendImpl.C)), ('cpu', None, ('cpu', Impl.C)),
('gpu-notnvidia', None, ('gpu', BackendImpl.JAX)), ('gpu-notnvidia', None, ('gpu', Impl.JAX)),
('gpu-nvidia', None, ('gpu', BackendImpl.WARP)), ('gpu-nvidia', None, ('gpu', Impl.WARP)),
('tpu', None, ('tpu', BackendImpl.JAX)), ('tpu', None, ('tpu', Impl.JAX)),
# JAX backend impl specified. # JAX backend impl specified.
('cpu', 'jax', ('cpu', BackendImpl.JAX)), ('cpu', 'jax', ('cpu', Impl.JAX)),
('gpu-notnvidia', 'jax', ('gpu', BackendImpl.JAX)), ('gpu-notnvidia', 'jax', ('gpu', Impl.JAX)),
('gpu-nvidia', 'jax', ('gpu', BackendImpl.JAX)), ('gpu-nvidia', 'jax', ('gpu', Impl.JAX)),
('tpu', 'jax', ('tpu', BackendImpl.JAX)), ('tpu', 'jax', ('tpu', Impl.JAX)),
# WARP backend impl specified. # WARP backend impl specified.
('cpu', 'warp', ('cpu', 'error')), ('cpu', 'warp', ('cpu', 'error')),
('gpu-notnvidia', 'warp', ('cpu', 'error')), ('gpu-notnvidia', 'warp', ('cpu', 'error')),
('gpu-nvidia', 'warp', ('gpu', BackendImpl.WARP)), ('gpu-nvidia', 'warp', ('gpu', Impl.WARP)),
('tpu', 'warp', ('tpu', 'error')), ('tpu', 'warp', ('tpu', 'error')),
# C backend impl specified, CPU should always be available. # C backend impl specified, CPU should always be available.
('cpu', 'c', ('cpu', BackendImpl.C)), ('cpu', 'c', ('cpu', Impl.C)),
('gpu-notnvidia', 'c', ('cpu', BackendImpl.C)), ('gpu-notnvidia', 'c', ('cpu', Impl.C)),
('gpu-nvidia', 'c', ('cpu', BackendImpl.C)), ('gpu-nvidia', 'c', ('cpu', Impl.C)),
('tpu', 'c', ('cpu', BackendImpl.C)), ('tpu', 'c', ('cpu', Impl.C)),
] ]
class ResolveBackendImplAndDeviceTest(parameterized.TestCase): class ResolveImplAndDeviceTest(parameterized.TestCase):
"""Tests for the _resolve_backend_impl_and_device function.""" """Tests for the _resolve_impl_and_device function."""
def setUp(self): def setUp(self):
super().setUp() super().setUp()
@@ -783,7 +783,7 @@ class ResolveBackendImplAndDeviceTest(parameterized.TestCase):
def test_resolve_with_device( def test_resolve_with_device(
self, self,
device_type_str, device_type_str,
backend_impl_str, impl_str,
expected, expected,
): ):
"""Tests various combinations of device and backend impls.""" """Tests various combinations of device and backend impls."""
@@ -811,21 +811,21 @@ class ResolveBackendImplAndDeviceTest(parameterized.TestCase):
self.mock_jax_devices.side_effect = devices_side_effect self.mock_jax_devices.side_effect = devices_side_effect
expected_device, expected_backend_impl = expected expected_device, expected_impl = expected
if expected_backend_impl == 'error': if expected_impl == 'error':
with self.assertRaises(AssertionError): with self.assertRaises(AssertionError):
mjx_io._resolve_backend_impl_and_device( mjx_io._resolve_impl_and_device(
backend_impl=backend_impl_str, device=input_device impl=impl_str, device=input_device
) )
return return
actual_backend_impl, actual_device = ( actual_impl, actual_device = (
mjx_io._resolve_backend_impl_and_device( mjx_io._resolve_impl_and_device(
backend_impl=backend_impl_str, device=input_device impl=impl_str, device=input_device
) )
) )
self.assertEqual(actual_backend_impl, expected_backend_impl) self.assertEqual(actual_impl, expected_impl)
self.assertIsNotNone(actual_device) self.assertIsNotNone(actual_device)
self.assertEqual(actual_device.platform, expected_device) self.assertEqual(actual_device.platform, expected_device)
@@ -839,7 +839,7 @@ class ResolveBackendImplAndDeviceTest(parameterized.TestCase):
def test_resolve_without_device( def test_resolve_without_device(
self, self,
default_device_str, default_device_str,
backend_impl_str, impl_str,
expected, expected,
): ):
"""Tests various combinations of jax.default_device and backend impls.""" """Tests various combinations of jax.default_device and backend impls."""
@@ -878,32 +878,32 @@ class ResolveBackendImplAndDeviceTest(parameterized.TestCase):
lambda: default_device_side_effect_str lambda: default_device_side_effect_str
) )
expected_device, expected_backend_impl = expected expected_device, expected_impl = expected
if ( if (
expected_backend_impl == 'error' expected_impl == 'error'
and default_device_str != 'gpu-nvidia' and default_device_str != 'gpu-nvidia'
and backend_impl_str == 'warp' and impl_str == 'warp'
): ):
with self.assertRaisesRegex(RuntimeError, 'cuda backend not supported'): with self.assertRaisesRegex(RuntimeError, 'cuda backend not supported'):
mjx_io._resolve_backend_impl_and_device( mjx_io._resolve_impl_and_device(
backend_impl=backend_impl_str, device=None impl=impl_str, device=None
) )
return return
if expected_backend_impl == 'error': if expected_impl == 'error':
with self.assertRaises(AssertionError): with self.assertRaises(AssertionError):
mjx_io._resolve_backend_impl_and_device( mjx_io._resolve_impl_and_device(
backend_impl=backend_impl_str, device=None impl=impl_str, device=None
) )
return return
actual_backend_impl, actual_device = ( actual_impl, actual_device = (
mjx_io._resolve_backend_impl_and_device( mjx_io._resolve_impl_and_device(
backend_impl=backend_impl_str, device=None impl=impl_str, device=None
) )
) )
self.assertEqual(actual_backend_impl, expected_backend_impl) self.assertEqual(actual_impl, expected_impl)
self.assertIsNotNone(actual_device) self.assertIsNotNone(actual_device)
self.assertEqual(actual_device.platform, expected_device) self.assertEqual(actual_device.platform, expected_device)
@@ -918,26 +918,26 @@ class ResolveBackendImplAndDeviceTest(parameterized.TestCase):
self.mock_default_backend.side_effect = lambda: 'gpu' self.mock_default_backend.side_effect = lambda: 'gpu'
# Default to JAX instead of WARP on NVIDIA GPU. # Default to JAX instead of WARP on NVIDIA GPU.
backend_impl, device = mjx_io._resolve_backend_impl_and_device( impl, device = mjx_io._resolve_impl_and_device(
backend_impl=None, device=None impl=None, device=None
) )
self.assertEqual(backend_impl, BackendImpl.JAX) self.assertEqual(impl, Impl.JAX)
self.assertEqual(device.platform, 'gpu') self.assertEqual(device.platform, 'gpu')
# Specifying an NVIDIA GPU should still choose JAX. # Specifying an NVIDIA GPU should still choose JAX.
backend_impl, device = mjx_io._resolve_backend_impl_and_device( impl, device = mjx_io._resolve_impl_and_device(
backend_impl=None, device=self.mock_nvidia_gpu impl=None, device=self.mock_nvidia_gpu
) )
self.assertEqual(backend_impl, BackendImpl.JAX) self.assertEqual(impl, Impl.JAX)
self.assertEqual(device.platform, 'gpu') self.assertEqual(device.platform, 'gpu')
# Requesting warp explicitly should fail since it is disabled. # Requesting warp explicitly should fail since it is disabled.
with self.assertRaises(AssertionError): with self.assertRaises(AssertionError):
mjx_io._resolve_backend_impl_and_device( mjx_io._resolve_impl_and_device(
backend_impl='warp', device=self.mock_nvidia_gpu impl='warp', device=self.mock_nvidia_gpu
) )
with self.assertRaises(AssertionError): with self.assertRaises(AssertionError):
mjx_io._resolve_backend_impl_and_device(backend_impl='warp', device=None) mjx_io._resolve_impl_and_device(impl='warp', device=None)
@mock.patch.dict(os.environ, {'MJX_C_DEFAULT_ENABLED': 'false'}) @mock.patch.dict(os.environ, {'MJX_C_DEFAULT_ENABLED': 'false'})
def test_resolve_c_disabled(self): def test_resolve_c_disabled(self):
@@ -950,30 +950,30 @@ class ResolveBackendImplAndDeviceTest(parameterized.TestCase):
self.mock_default_backend.side_effect = lambda: 'cpu' self.mock_default_backend.side_effect = lambda: 'cpu'
# Default to JAX instead of C on CPU. # Default to JAX instead of C on CPU.
backend_impl, device = mjx_io._resolve_backend_impl_and_device( impl, device = mjx_io._resolve_impl_and_device(
backend_impl=None, device=None impl=None, device=None
) )
self.assertEqual(backend_impl, BackendImpl.JAX) self.assertEqual(impl, Impl.JAX)
self.assertEqual(device.platform, 'cpu') self.assertEqual(device.platform, 'cpu')
# Specifing CPU should still choose JAX. # Specifing CPU should still choose JAX.
backend_impl, device = mjx_io._resolve_backend_impl_and_device( impl, device = mjx_io._resolve_impl_and_device(
backend_impl=None, device=self.mock_cpu impl=None, device=self.mock_cpu
) )
self.assertEqual(backend_impl, BackendImpl.JAX) self.assertEqual(impl, Impl.JAX)
self.assertEqual(device.platform, 'cpu') self.assertEqual(device.platform, 'cpu')
# Specifying C should choose C! # Specifying C should choose C!
backend_impl, device = mjx_io._resolve_backend_impl_and_device( impl, device = mjx_io._resolve_impl_and_device(
backend_impl='c', device=None impl='c', device=None
) )
self.assertEqual(backend_impl, BackendImpl.C) self.assertEqual(impl, Impl.C)
self.assertEqual(device.platform, 'cpu') self.assertEqual(device.platform, 'cpu')
backend_impl, device = mjx_io._resolve_backend_impl_and_device( impl, device = mjx_io._resolve_impl_and_device(
backend_impl='c', device=self.mock_cpu impl='c', device=self.mock_cpu
) )
self.assertEqual(backend_impl, BackendImpl.C) self.assertEqual(impl, Impl.C)
self.assertEqual(device.platform, 'cpu') self.assertEqual(device.platform, 'cpu')
mjx/mujoco/mjx/_src/sensor.py
+1 -1
View File
@@ -22,7 +22,7 @@ from mujoco.mjx._src import math
from mujoco.mjx._src import ray from mujoco.mjx._src import ray
from mujoco.mjx._src import smooth from mujoco.mjx._src import smooth
from mujoco.mjx._src import support from mujoco.mjx._src import support
from mujoco.mjx._src.types import BackendImpl from mujoco.mjx._src.types import Impl
from mujoco.mjx._src.types import Data from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import DataJAX from mujoco.mjx._src.types import DataJAX
from mujoco.mjx._src.types import DisableBit from mujoco.mjx._src.types import DisableBit
mjx/mujoco/mjx/_src/types.py
+9 -9
View File
@@ -24,8 +24,8 @@ from mujoco.mjx._src.dataclasses import PyTreeNode # pylint: disable=g-importin
import numpy as np import numpy as np
class BackendImpl(enum.Enum): class Impl(enum.Enum):
"""Backend implementation to use.""" """Implementation to use."""
C = 'c' C = 'c'
JAX = 'jax' JAX = 'jax'
@@ -34,7 +34,7 @@ class BackendImpl(enum.Enum):
@classmethod @classmethod
def _missing_(cls, value): def _missing_(cls, value):
# This method is called only when lookup by value fails # This method is called only when lookup by value fails
# (e.g., BackendImpl('JAX') fails initially because 'JAX' != 'jax') # (e.g., Impl('JAX') fails initially because 'JAX' != 'jax')
if not isinstance(value, str): if not isinstance(value, str):
return None return None
for member in cls: for member in cls:
@@ -871,10 +871,10 @@ class Model(PyTreeNode):
_impl: Union[ModelC, ModelJAX] _impl: Union[ModelC, ModelJAX]
@property @property
def backend_impl(self) -> BackendImpl: def impl(self) -> Impl:
return { return {
ModelC: BackendImpl.C, ModelC: Impl.C,
ModelJAX: BackendImpl.JAX, ModelJAX: Impl.JAX,
}[type(self._impl)] }[type(self._impl)]
def __getattr__(self, name: str): def __getattr__(self, name: str):
@@ -1119,10 +1119,10 @@ class Data(PyTreeNode):
_impl: Union[DataC, DataJAX] _impl: Union[DataC, DataJAX]
@property @property
def backend_impl(self) -> BackendImpl: def impl(self) -> Impl:
return { return {
DataC: BackendImpl.C, DataC: Impl.C,
DataJAX: BackendImpl.JAX, DataJAX: Impl.JAX,
}[type(self._impl)] }[type(self._impl)]
def __getattr__(self, name: str): def __getattr__(self, name: str):