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

Fix #2306.

Browse Source 
PiperOrigin-RevId: 712575320
Change-Id: Ia74ef0b31b4e7098647340760572a1f6368eab64
This commit is contained in:
Baruch Tabanpour
2025-01-06 10:27:16 -08:00
committed by Copybara-Service
parent f607d9554e
commit 0be7e6a1f6
2 changed files with 11 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/io.py
+2 -1
View File
@@ -681,4 +681,5 @@ def put_data(
# copy because device_put is async: # copy because device_put is async:
data = types.Data(**{k: copy.copy(v) for k, v in fields.items()}) data = types.Data(**{k: copy.copy(v) for k, v in fields.items()})
return jax.device_put(data, device=device) data = jax.device_put(data, device=device)
return _strip_weak_type(data)
mjx/mujoco/mjx/_src/io_test.py
+9
View File
@@ -310,6 +310,15 @@ class DataIOTest(parameterized.TestCase):
np.testing.assert_allclose(dx.cvel, d.cvel) np.testing.assert_allclose(dx.cvel, d.cvel)
np.testing.assert_allclose(dx.cdof_dot, d.cdof_dot) np.testing.assert_allclose(dx.cdof_dot, d.cdof_dot)
# check that there are no weak types
self.assertFalse(
any(
jax.tree_util.tree_flatten(
jax.tree_util.tree_map(lambda x: x.weak_type, dx)
)[0]
)
)
# 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)