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

Remove disable_jit in test.

Browse Source 
PiperOrigin-RevId: 614824680
Change-Id: I15f1441dfb1a70b98e5a221c0e15b06572860748
This commit is contained in:
Baruch Tabanpour
2024-03-11 16:02:31 -07:00
committed by Copybara-Service
parent 8413a615d7
commit a49f3f63fc
1 changed files with 16 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/collision_driver_test.py
+16 -18
View File
@@ -297,8 +297,7 @@ class ConvexTest(absltest.TestCase):
def test_flat_box_plane(self):
"""Tests box collision with a plane."""
with jax.disable_jit():
d, dx = _collide(self._FLAT_BOX_PLANE)
d, dx = _collide(self._FLAT_BOX_PLANE)
np.testing.assert_array_less(dx.contact.dist, 0)
@@ -586,25 +585,24 @@ class TopKContactTest(absltest.TestCase):
def test_max_pair(self):
"""Tests contact culling before the collision functions were dispatched."""
with jax.disable_jit():
m = mujoco.MjModel.from_xml_string(self._CAPSULES_MAX_PAIR)
mx_top_k = mjx.put_model(m)
mx_all = mx_top_k.replace(
nnumeric=0, name_numericadr=np.array([]), numeric_data=np.array([])
)
d = mujoco.MjData(m)
dx = mjx.put_data(m, d)
m = mujoco.MjModel.from_xml_string(self._CAPSULES_MAX_PAIR)
mx_top_k = mjx.put_model(m)
mx_all = mx_top_k.replace(
nnumeric=0, name_numericadr=np.array([]), numeric_data=np.array([])
)
d = mujoco.MjData(m)
dx = mjx.put_data(m, d)
collision_jit_fn = jax.jit(mjx.collision)
kinematics_jit_fn = jax.jit(mjx.kinematics)
dx = kinematics_jit_fn(mx_all, dx)
collision_jit_fn = jax.jit(mjx.collision)
kinematics_jit_fn = jax.jit(mjx.kinematics)
dx = kinematics_jit_fn(mx_all, dx)
dx_all = collision_jit_fn(mx_all, dx)
dx_top_k = collision_jit_fn(mx_top_k, dx)
dx_all = collision_jit_fn(mx_all, dx)
dx_top_k = collision_jit_fn(mx_top_k, dx)
self.assertEqual(dx_all.contact.dist.shape, (6,))
self.assertEqual(dx_top_k.contact.dist.shape, (2,))
self.assertTrue((dx_top_k.contact.dist < 0).all())
self.assertEqual(dx_all.contact.dist.shape, (6,))
self.assertEqual(dx_top_k.contact.dist.shape, (2,))
self.assertTrue((dx_top_k.contact.dist < 0).all())
if __name__ == '__main__':