MJX put_data with impl='warp'
PiperOrigin-RevId: 874017611 Change-Id: I23a25727d41cd888d90f2ea15b2b534eba5583be
This commit is contained in:
Taylor Howell
committed by
Copybara-Service
Copybara-Service
parent
ba0a5fe4db
commit
2bf9539234
+42
-14
@@ -903,17 +903,6 @@ def _make_data_warp(
|
||||
|
||||
data = jax.device_put(data, device=device)
|
||||
|
||||
with wp.ScopedDevice('cuda:0'): # pylint: disable=undefined-variable
|
||||
# Warm-up the warp kernel cache.
|
||||
# TODO(robotics-simulation): remove this warmup compilation once warp
|
||||
# stops unloading modules during XLA graph capture for tile kernels.
|
||||
# pylint: disable=undefined-variable
|
||||
dw = mjwp.make_data(m, nworld=1, naconmax=naconmax, njmax=njmax)
|
||||
mw = mjwp.put_model(m)
|
||||
_ = mjwp.step(mw, dw)
|
||||
# pylint: enable=undefined-variable
|
||||
del dw, mw
|
||||
|
||||
return data
|
||||
|
||||
|
||||
@@ -1363,6 +1352,44 @@ def _put_data_cpp(
|
||||
return _strip_weak_type(data)
|
||||
|
||||
|
||||
def _put_data_warp(
|
||||
m: mujoco.MjModel,
|
||||
d: mujoco.MjData,
|
||||
device: Optional[jax.Device] = None,
|
||||
naconmax: Optional[int] = None,
|
||||
njmax: Optional[int] = None,
|
||||
) -> types.Data:
|
||||
"""Puts mujoco.MjData onto a device, resulting in mjx.Data."""
|
||||
|
||||
with wp.ScopedDevice('cpu'): # pylint: disable=undefined-variable
|
||||
dw = mjwp.put_data(m, d, nworld=1, naconmax=naconmax, njmax=njmax) # pylint: disable=undefined-variable
|
||||
|
||||
fields = _put_data_public_fields(d)
|
||||
for k in fields:
|
||||
if not hasattr(dw, k):
|
||||
continue
|
||||
field = _wp_to_np_type(getattr(dw, k))
|
||||
if mjxw.types._BATCH_DIM['Data'][k]: # pylint: disable=protected-access
|
||||
field = field.reshape(field.shape[1:])
|
||||
fields[k] = field
|
||||
|
||||
impl_fields = {}
|
||||
for k in mjxw.types.DataWarp.__annotations__.keys():
|
||||
field = _get_nested_attr(dw, k, split='__')
|
||||
field = _wp_to_np_type(field)
|
||||
if mjxw.types._BATCH_DIM['Data'][k]: # pylint: disable=protected-access
|
||||
field = field.reshape(field.shape[1:])
|
||||
impl_fields[k] = field
|
||||
|
||||
data = types.Data(
|
||||
**fields,
|
||||
_impl=mjxw.types.DataWarp(**impl_fields),
|
||||
)
|
||||
|
||||
data = jax.device_put(data, device=device)
|
||||
return data
|
||||
|
||||
|
||||
def put_data(
|
||||
m: mujoco.MjModel,
|
||||
d: mujoco.MjData,
|
||||
@@ -1393,7 +1420,6 @@ def put_data(
|
||||
an mjx.Data placed on device
|
||||
DeprecationWarning: if nconmax is used
|
||||
"""
|
||||
del njmax
|
||||
if nconmax is not None:
|
||||
warnings.warn(
|
||||
'nconmax will be deprecated in mujoco-mjx>=3.5. Use naconmax instead.',
|
||||
@@ -1410,8 +1436,10 @@ def put_data(
|
||||
return _put_data_cpp(
|
||||
m, d, device, dummy_arg_for_batching=dummy_arg_for_batching
|
||||
)
|
||||
|
||||
# TODO(robotics-team): implement put_data_warp
|
||||
elif impl == types.Impl.WARP:
|
||||
_check_warp_installed()
|
||||
naconmax = nconmax if naconmax is None else naconmax
|
||||
return _put_data_warp(m, d, device, naconmax, njmax)
|
||||
|
||||
raise NotImplementedError(
|
||||
f'put_data for implementation "{impl}" not implemented yet.'
|
||||
|
||||
@@ -494,9 +494,15 @@ class DataIOTest(parameterized.TestCase):
|
||||
self.assertEqual(d._impl.contact__dist.shape[0], 9)
|
||||
self.assertEqual(d._impl.efc__pos.shape[0], 23)
|
||||
|
||||
@parameterized.parameters('jax', 'c', 'cpp')
|
||||
@parameterized.parameters('jax', 'c', 'cpp', 'warp')
|
||||
def test_put_data(self, impl: str):
|
||||
"""Test that put_data puts the correct data for dense and sparse."""
|
||||
if impl == 'warp':
|
||||
if not mjxw.WARP_INSTALLED:
|
||||
self.skipTest('Warp is not installed.')
|
||||
if not mjx_io.has_cuda_gpu_device():
|
||||
self.skipTest('No CUDA GPU device.')
|
||||
|
||||
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
|
||||
d = mujoco.MjData(m)
|
||||
mujoco.mj_step(m, d, 2)
|
||||
@@ -517,6 +523,15 @@ class DataIOTest(parameterized.TestCase):
|
||||
)
|
||||
)
|
||||
|
||||
# xmat, ximat, geom_xmat are all shape transformed
|
||||
np.testing.assert_allclose(dx.xmat.reshape((-1, 9)), d.xmat)
|
||||
np.testing.assert_allclose(dx.ximat.reshape((-1, 9)), d.ximat)
|
||||
np.testing.assert_allclose(dx.geom_xmat.reshape((-1, 9)), d.geom_xmat)
|
||||
np.testing.assert_allclose(dx.site_xmat.reshape((-1, 9)), d.site_xmat)
|
||||
|
||||
# tendon length is correct
|
||||
np.testing.assert_allclose(dx.ten_length, d.ten_length)
|
||||
|
||||
if impl == 'jax':
|
||||
# check that qM is transformed properly
|
||||
qm = np.zeros((m.nv, m.nv), dtype=np.float64)
|
||||
@@ -530,6 +545,21 @@ class DataIOTest(parameterized.TestCase):
|
||||
self.assertTrue(hasattr(dx._impl, 'pointer_lo'))
|
||||
self.assertTrue(hasattr(dx._impl, 'pointer_hi'))
|
||||
return # cpp does not populate other fields in _impl
|
||||
elif impl == 'warp':
|
||||
qm = np.zeros((m.nv, m.nv), dtype=np.float64)
|
||||
mujoco.mj_fullM(m, qm, d.qM)
|
||||
np.testing.assert_allclose(dx._impl.qM, qm)
|
||||
# TODO(taylorhowell): test efc__J
|
||||
np.testing.assert_allclose(dx._impl.efc__aref[:3], d.efc_aref[:3])
|
||||
|
||||
# tendon impl data is correct
|
||||
np.testing.assert_equal(dx._impl.ten_wrapadr, np.zeros((1,)))
|
||||
np.testing.assert_equal(dx._impl.ten_wrapnum, np.zeros((1,)))
|
||||
np.testing.assert_equal(dx._impl.wrap_obj, np.zeros((2, 2)))
|
||||
np.testing.assert_equal(dx._impl.wrap_xpos, np.zeros((2, 6)))
|
||||
|
||||
if impl == 'warp':
|
||||
return
|
||||
|
||||
# 4 contacts, 2 for each capsule against the plane
|
||||
self.assertEqual(dx._impl.contact.dist.shape, (4,))
|
||||
@@ -542,23 +572,6 @@ class DataIOTest(parameterized.TestCase):
|
||||
)
|
||||
np.testing.assert_allclose(dx._impl.contact.frame[1:], 0)
|
||||
|
||||
# xmat, ximat, geom_xmat are all shape transformed
|
||||
self.assertEqual(dx.xmat.shape, (3, 3, 3))
|
||||
self.assertEqual(dx.ximat.shape, (3, 3, 3))
|
||||
self.assertEqual(dx.geom_xmat.shape, (3, 3, 3))
|
||||
self.assertEqual(dx.site_xmat.shape, (1, 3, 3))
|
||||
np.testing.assert_allclose(dx.xmat.reshape((3, 9)), d.xmat)
|
||||
np.testing.assert_allclose(dx.ximat.reshape((3, 9)), d.ximat)
|
||||
np.testing.assert_allclose(dx.geom_xmat.reshape((3, 9)), d.geom_xmat)
|
||||
np.testing.assert_allclose(dx.site_xmat.reshape((1, 9)), d.site_xmat)
|
||||
|
||||
# tendon data is correct
|
||||
np.testing.assert_allclose(dx.ten_length, d.ten_length)
|
||||
np.testing.assert_equal(dx._impl.ten_wrapadr, np.zeros((1,)))
|
||||
np.testing.assert_equal(dx._impl.ten_wrapnum, np.zeros((1,)))
|
||||
np.testing.assert_equal(dx._impl.wrap_obj, np.zeros((2, 2)))
|
||||
np.testing.assert_equal(dx._impl.wrap_xpos, np.zeros((2, 6)))
|
||||
|
||||
# efc_ are also shape transformed and padded
|
||||
self.assertEqual(dx._impl.efc_J.shape, (45, 8)) # nefc, nv
|
||||
d_efc_j = d.efc_J.reshape((-1, 8))
|
||||
@@ -583,7 +596,9 @@ class DataIOTest(parameterized.TestCase):
|
||||
d = mujoco.MjData(m)
|
||||
mujoco.mj_step(m, d, 2)
|
||||
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
|
||||
np.testing.assert_allclose(dx_sparse._impl.qM, d.qM, atol=1e-8)
|
||||
@@ -604,6 +619,31 @@ class DataIOTest(parameterized.TestCase):
|
||||
elif impl == 'c':
|
||||
np.testing.assert_allclose(dx_from_dense._impl.qM, d.qM, atol=1e-8)
|
||||
|
||||
def test_put_data_warp_ndim(self):
|
||||
"""Tests that put_data produces expected dimensions for Warp fields."""
|
||||
if not mjxw.WARP_INSTALLED:
|
||||
self.skipTest('Warp is not installed.')
|
||||
if not mjx_io.has_cuda_gpu_device():
|
||||
self.skipTest('No CUDA GPU device.')
|
||||
|
||||
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
|
||||
d = mujoco.MjData(m)
|
||||
mujoco.mj_step(m, d, 2)
|
||||
dx = mjx.put_data(m, d, impl='warp')
|
||||
|
||||
def check_ndim(path, x):
|
||||
k = _get_name_from_path(path)
|
||||
if k not in mjxw_types._NDIM['Data']:
|
||||
return
|
||||
is_batched = mjxw_types._BATCH_DIM['Data'][k]
|
||||
expected_ndim = mjxw_types._NDIM['Data'][k] - is_batched
|
||||
if not hasattr(x, 'ndim'):
|
||||
return
|
||||
msg = f'Field {k} has ndim {x.ndim} but expected {expected_ndim}'
|
||||
self.assertEqual(x.ndim, expected_ndim, msg)
|
||||
|
||||
_ = jax.tree.map_with_path(check_ndim, dx)
|
||||
|
||||
@parameterized.parameters(
|
||||
('jax', False), ('jax', True), ('c', False), ('c', True)
|
||||
)
|
||||
|
||||
@@ -88,15 +88,9 @@ def _main(argv: Sequence[str]) -> None:
|
||||
m = mujoco.MjModel.from_xml_path(_MODEL_PATH.value)
|
||||
d = mujoco.MjData(m)
|
||||
mx = mjx.put_model(m, impl=_IMPL.value)
|
||||
if _IMPL.value == 'warp':
|
||||
# TODO(btaba): use put_data.
|
||||
dx = mjx.make_data(
|
||||
m, impl=_IMPL.value, naconmax=_NACONMAX.value, njmax=_NJMAX.value
|
||||
)
|
||||
else:
|
||||
dx = mjx.put_data(
|
||||
m, d, impl=_IMPL.value, naconmax=_NACONMAX.value, njmax=_NJMAX.value
|
||||
)
|
||||
dx = mjx.put_data(
|
||||
m, d, impl=_IMPL.value, naconmax=_NACONMAX.value, njmax=_NJMAX.value
|
||||
)
|
||||
|
||||
print(f'Default backend: {jax.default_backend()}')
|
||||
step_fn = mjx.step
|
||||
|
||||
Reference in New Issue
Block a user