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

Add warp-lang requirements back.

Browse Source 
PiperOrigin-RevId: 791872106
Change-Id: I74bc5069f0f356538b6ce4dde4bac0a47e1cf5f5
This commit is contained in:
Baruch Tabanpour
2025-08-06 15:40:45 -07:00
committed by Copybara-Service
parent 132f90b17f
commit 9b149f334f
6 changed files with 70 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/cuda_requirements.txt
+16 -13
View File
@@ -1,18 +1,21 @@
jax-cuda12-plugin==0.5.3; python_version >= '3.10' \ jax-cuda12-plugin==0.5.3; python_version >= '3.10' \
--hash=sha256:6171aed2f4b3bdd5fc13782de1072c6a634fce13731b75d0cb0a6ab8f4e6e650 \
--hash=sha256:ba2555967f9b6c381c8b4ef9fb03d05bc55ec25ecfee5cfe45c5ace34f7d4152 \
--hash=sha256:298d2d768f1029b74a0b1d01270e549349d2c37dc07658796542cda967eb7bd3 \
--hash=sha256:aaa704a5ef547595d022db1c1e4878a0677116412a9360c115d67ff4b64e1596 \
--hash=sha256:c2517a7c2186f8708894696e26cf96ebd60b7879ceca398b2c46abb28d2c96c8 \
--hash=sha256:2030cf1208ce4ea70ee56cac61ddd239f9798695fc39bb7739c50a25d6e9da44 \
--hash=sha256:21fec1b56c98783ea0569b747a56751f1f9ff2187b48acc11c700d3bfc5e1a31 \
--hash=sha256:1862595b2b6d815679d11e0e889e523185ee54a46d46e022689f70fc4554dd91 \ --hash=sha256:1862595b2b6d815679d11e0e889e523185ee54a46d46e022689f70fc4554dd91 \
--hash=sha256:6d43677f22f3be9544a205216cd6dac591335b1d9bbbed018cd17dbb1f3f4def \ --hash=sha256:21fec1b56c98783ea0569b747a56751f1f9ff2187b48acc11c700d3bfc5e1a31 \
--hash=sha256:5bb9ea0e68d72d44e57e4cb6a58a1a729fe3fe32e964f71e398d8a25c2103b19 --hash=sha256:2030cf1208ce4ea70ee56cac61ddd239f9798695fc39bb7739c50a25d6e9da44 \
--hash=sha256:c2517a7c2186f8708894696e26cf96ebd60b7879ceca398b2c46abb28d2c96c8 \
--hash=sha256:aaa704a5ef547595d022db1c1e4878a0677116412a9360c115d67ff4b64e1596 \
--hash=sha256:298d2d768f1029b74a0b1d01270e549349d2c37dc07658796542cda967eb7bd3 \
--hash=sha256:ba2555967f9b6c381c8b4ef9fb03d05bc55ec25ecfee5cfe45c5ace34f7d4152 \
--hash=sha256:6171aed2f4b3bdd5fc13782de1072c6a634fce13731b75d0cb0a6ab8f4e6e650
jax-cuda12-plugin==0.4.30; python_version == '3.9' \ jax-cuda12-plugin==0.4.30; python_version == '3.9' \
--hash=sha256:d8d196241b9253ecb1144a4409b5deacbb9771624f097b2bbf025da3c7d8f4f8 --hash=sha256:d8d196241b9253ecb1144a4409b5deacbb9771624f097b2bbf025da3c7d8f4f8 \
--hash=sha256:cb8edccdce358451205f689e3536272200761c625c8e8059ab10523984cf8b61
jax-cuda12-pjrt==0.5.3; python_version >= '3.10' \ jax-cuda12-pjrt==0.5.3; python_version >= '3.10' \
--hash=sha256:04ee111eaf5fc2692978ad4a5c84d5925e42eb05c1701849ba3a53f6515400cc \ --hash=sha256:c5378306568ba0c81b230a779dd3194c9dd10339ab6360ae80928108d37e7f75 \
--hash=sha256:c5378306568ba0c81b230a779dd3194c9dd10339ab6360ae80928108d37e7f75 --hash=sha256:04ee111eaf5fc2692978ad4a5c84d5925e42eb05c1701849ba3a53f6515400cc
jax-cuda12-pjrt==0.4.30; python_version == '3.9' \ jax-cuda12-pjrt==0.4.30; python_version == '3.9' \
--hash=sha256:895d0198ad99638fcaf976c47592e2a543eef79ea15fabd24a402d055390c328 --hash=sha256:895d0198ad99638fcaf976c47592e2a543eef79ea15fabd24a402d055390c328 \
--hash=sha256:c36fb1e0c236563bf3a87e70f4d1ab28a31d7cf5d722c9ede30c4172116e8bcb
warp-lang==1.8.1 \
--hash=sha256:cfc59e1070ad71531b5d83186de48162507277af344a102fa33d5df9cdb942f7 \
--hash=sha256:1db9ca92c46902b76bb99565c544347d1a32e9fb875ce902f1cafb94978d1ac3
mjx/mujoco/mjx/warp/collision_driver_test.py
+8 -4
View File
@@ -13,6 +13,7 @@
# limitations under the License. # limitations under the License.
# ============================================================================== # ==============================================================================
"""Tests for collision driver.""" """Tests for collision driver."""
import os
from absl.testing import absltest from absl.testing import absltest
import jax import jax
@@ -32,6 +33,8 @@ except ImportError:
collision_driver = None collision_driver = None
smooth = None smooth = None
_FORCE_TEST = os.environ.get('MJX_WARP_FORCE_TEST', '0') == '1'
class CollisionTest(absltest.TestCase): class CollisionTest(absltest.TestCase):
@@ -58,10 +61,11 @@ class CollisionTest(absltest.TestCase):
def test_collision_nested_vmap(self): def test_collision_nested_vmap(self):
"""Tests collision with batched data.""" """Tests collision with batched data."""
if not mjxw.WARP_INSTALLED: if not _FORCE_TEST:
self.skipTest('Warp not installed.') if not mjxw.WARP_INSTALLED:
if not io.has_cuda_gpu_device(): self.skipTest('Warp not installed.')
self.skipTest('No CUDA GPU device available.') if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
m = mujoco.MjModel.from_xml_string(self._SPHERE_SPHERE) m = mujoco.MjModel.from_xml_string(self._SPHERE_SPHERE)
d = mujoco.MjData(m) d = mujoco.MjData(m)
mjx/mujoco/mjx/warp/forward_test.py
+27 -40
View File
@@ -15,7 +15,7 @@
"""Tests for forward functions.""" """Tests for forward functions."""
import functools import functools
import logging import os
from absl.testing import absltest from absl.testing import absltest
from absl.testing import parameterized from absl.testing import parameterized
@@ -35,6 +35,8 @@ try:
except ImportError: except ImportError:
forward = None forward = None
_FORCE_TEST = os.environ.get('MJX_WARP_FORCE_TEST', '0') == '1'
class ForwardTest(parameterized.TestCase): class ForwardTest(parameterized.TestCase):
@@ -50,10 +52,11 @@ class ForwardTest(parameterized.TestCase):
) )
def test_jit_caching(self, xml): def test_jit_caching(self, xml):
"""Tests jit caching on the full step function.""" """Tests jit caching on the full step function."""
if not mjxw.WARP_INSTALLED: if not _FORCE_TEST:
self.skipTest('Warp not installed.') if not mjxw.WARP_INSTALLED:
if not io.has_cuda_gpu_device(): self.skipTest('Warp not installed.')
self.skipTest('No CUDA GPU device available.') if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
batch_size = 7 batch_size = 7
m = test_util.load_test_file(xml) m = test_util.load_test_file(xml)
@@ -63,34 +66,16 @@ class ForwardTest(parameterized.TestCase):
dx_batch = jax.vmap(functools.partial(tu.make_data, m))(keys) dx_batch = jax.vmap(functools.partial(tu.make_data, m))(keys)
step_fn = jax.jit(jax.vmap(forward.step, in_axes=(None, 0))) step_fn = jax.jit(jax.vmap(forward.step, in_axes=(None, 0)))
dx_batch1 = step_fn(mx, dx_batch)
jax.tree_util.tree_map(lambda x: x.block_until_ready(), dx_batch1)
self.assertEqual(step_fn._cache_size(), 1)
was_logging_compiles = jax.config.jax_log_compiles # Re-generate data and run step_fn again to test jit caching.
jax_logger = logging.getLogger('jax') keys = jp.arange(batch_size, batch_size * 2)
was_propagating = jax_logger.propagate dx_batch = jax.vmap(functools.partial(tu.make_data, m))(keys)
jax.config.update('jax_log_compiles', True) dx_batch2 = step_fn(mx, dx_batch)
jax_logger.propagate = False # do not print to stdout for this test jax.tree_util.tree_map(lambda x: x.block_until_ready(), dx_batch2)
with self.assertLogs('jax', level='INFO') as log: self.assertEqual(step_fn._cache_size(), 1)
dx_batch1 = step_fn(mx, dx_batch)
jax.tree_util.tree_map(lambda x: x.block_until_ready(), dx_batch1)
# Re-generate data and run step_fn again to test jit caching.
keys = jp.arange(batch_size, batch_size * 2)
dx_batch = jax.vmap(functools.partial(tu.make_data, m))(keys)
dx_batch2 = step_fn(mx, dx_batch)
jax.tree_util.tree_map(lambda x: x.block_until_ready(), dx_batch2)
jax.config.update('jax_log_compiles', was_logging_compiles)
jax_logger.propagate = was_propagating
compilation_logs = [
r for r in log.records if 'Compiling jit(step)' in r.getMessage()
]
self.assertLen(
compilation_logs,
1,
msg=(
f'Expected 1 compilation, got {len(compilation_logs)} compilations.'
),
)
@parameterized.product( @parameterized.product(
xml=( xml=(
@@ -100,10 +85,11 @@ class ForwardTest(parameterized.TestCase):
batch_size=(1, 7), batch_size=(1, 7),
) )
def test_forward(self, xml: str, batch_size: int): def test_forward(self, xml: str, batch_size: int):
if not mjxw.WARP_INSTALLED: if not _FORCE_TEST:
self.skipTest('Warp not installed.') if not mjxw.WARP_INSTALLED:
if not io.has_cuda_gpu_device(): self.skipTest('Warp not installed.')
self.skipTest('No CUDA GPU device available.') if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
m = test_util.load_test_file(xml) m = test_util.load_test_file(xml)
m.opt.iterations = 10 m.opt.iterations = 10
@@ -228,10 +214,11 @@ class StepTest(parameterized.TestCase):
batch_size=(1, 7), batch_size=(1, 7),
) )
def test_step(self, xml: str, batch_size: int): def test_step(self, xml: str, batch_size: int):
if not mjxw.WARP_INSTALLED: if not _FORCE_TEST:
self.skipTest('Warp not installed.') if not mjxw.WARP_INSTALLED:
if not io.has_cuda_gpu_device(): self.skipTest('Warp not installed.')
self.skipTest('No CUDA GPU device available.') if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
m = test_util.load_test_file(xml) m = test_util.load_test_file(xml)
m.opt.iterations = 10 m.opt.iterations = 10
mjx/mujoco/mjx/warp/smooth_test.py
+18 -12
View File
@@ -15,6 +15,7 @@
"""Tests for codegen'd smooth functions.""" """Tests for codegen'd smooth functions."""
import functools import functools
import os
from absl.testing import absltest from absl.testing import absltest
import jax import jax
@@ -33,6 +34,8 @@ try:
except ImportError: except ImportError:
smooth = None smooth = None
_FORCE_TEST = os.environ.get('MJX_WARP_FORCE_TEST', '0') == '1'
class SmoothTest(absltest.TestCase): class SmoothTest(absltest.TestCase):
@@ -44,10 +47,11 @@ class SmoothTest(absltest.TestCase):
def test_kinematics(self): def test_kinematics(self):
"""Tests kinematics with unbatched data.""" """Tests kinematics with unbatched data."""
if not mjxw.WARP_INSTALLED: if not _FORCE_TEST:
self.skipTest('Warp not installed.') if not mjxw.WARP_INSTALLED:
if not io.has_cuda_gpu_device(): self.skipTest('Warp not installed.')
self.skipTest('No CUDA GPU device available.') if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
m = tu.load_test_file('pendula.xml') m = tu.load_test_file('pendula.xml')
@@ -129,10 +133,11 @@ class SmoothTest(absltest.TestCase):
def test_kinematics_nested_vmap(self): def test_kinematics_nested_vmap(self):
"""Tests kinematics with nested batch data.""" """Tests kinematics with nested batch data."""
if not mjxw.WARP_INSTALLED: if not _FORCE_TEST:
self.skipTest('Warp not installed.') if not mjxw.WARP_INSTALLED:
if not io.has_cuda_gpu_device(): self.skipTest('Warp not installed.')
self.skipTest('No CUDA GPU device available.') if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
m = tu.load_test_file('pendula.xml') m = tu.load_test_file('pendula.xml')
@@ -175,10 +180,11 @@ class SmoothTest(absltest.TestCase):
def test_kinematics_model_vmap(self): def test_kinematics_model_vmap(self):
"""Tests kinematics with vmap on model and data fields.""" """Tests kinematics with vmap on model and data fields."""
if not mjxw.WARP_INSTALLED: if not _FORCE_TEST:
self.skipTest('Warp not installed.') if not mjxw.WARP_INSTALLED:
if not io.has_cuda_gpu_device(): self.skipTest('Warp not installed.')
self.skipTest('No CUDA GPU device available.') if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
m = tu.load_test_file('pendula.xml') m = tu.load_test_file('pendula.xml')
mjx/pyproject.toml
+1 -5
View File
@@ -32,11 +32,7 @@ dependencies = [
"mujoco>=3.3.5.dev0", "mujoco>=3.3.5.dev0",
"scipy", "scipy",
"trimesh", "trimesh",
] "warp-lang==1.8.1",
[project.optional-dependencies]
warp = [
"warp-lang==1.8.0",
] ]
[project.scripts] [project.scripts]
mjx/requirements.txt
-1
View File
@@ -103,7 +103,6 @@ trimesh==4.5.2 \
--hash=sha256:2e50f3a7fd135c3045da887a1b9f91230528f3ce11d2ec1ba44750d82d6b4f73 --hash=sha256:2e50f3a7fd135c3045da887a1b9f91230528f3ce11d2ec1ba44750d82d6b4f73
wheel==0.45.0 \ wheel==0.45.0 \
--hash=sha256:52f0baa5e6522155090a09c6bd95718cc46956d1b51d537ea5454249edb671c7 --hash=sha256:52f0baa5e6522155090a09c6bd95718cc46956d1b51d537ea5454249edb671c7
# TODO(robotics-simulation): add back warp-lang dependencies
# Transitive dependencies of etils[epath] # Transitive dependencies of etils[epath]
fsspec==2024.10.0 \ fsspec==2024.10.0 \