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' \
--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:6d43677f22f3be9544a205216cd6dac591335b1d9bbbed018cd17dbb1f3f4def \
--hash=sha256:5bb9ea0e68d72d44e57e4cb6a58a1a729fe3fe32e964f71e398d8a25c2103b19
--hash=sha256:21fec1b56c98783ea0569b747a56751f1f9ff2187b48acc11c700d3bfc5e1a31 \
--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' \
--hash=sha256:d8d196241b9253ecb1144a4409b5deacbb9771624f097b2bbf025da3c7d8f4f8
--hash=sha256:d8d196241b9253ecb1144a4409b5deacbb9771624f097b2bbf025da3c7d8f4f8 \
--hash=sha256:cb8edccdce358451205f689e3536272200761c625c8e8059ab10523984cf8b61
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' \
--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.
# ==============================================================================
"""Tests for collision driver."""
import os
from absl.testing import absltest
import jax
@@ -32,6 +33,8 @@ except ImportError:
collision_driver = None
smooth = None
_FORCE_TEST = os.environ.get('MJX_WARP_FORCE_TEST', '0') == '1'
class CollisionTest(absltest.TestCase):
@@ -58,10 +61,11 @@ class CollisionTest(absltest.TestCase):
def test_collision_nested_vmap(self):
"""Tests collision with batched data."""
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp not installed.')
if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
if not _FORCE_TEST:
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp not installed.')
if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
m = mujoco.MjModel.from_xml_string(self._SPHERE_SPHERE)
d = mujoco.MjData(m)
mjx/mujoco/mjx/warp/forward_test.py
+27 -40
View File
@@ -15,7 +15,7 @@
"""Tests for forward functions."""
import functools
import logging
import os
from absl.testing import absltest
from absl.testing import parameterized
@@ -35,6 +35,8 @@ try:
except ImportError:
forward = None
_FORCE_TEST = os.environ.get('MJX_WARP_FORCE_TEST', '0') == '1'
class ForwardTest(parameterized.TestCase):
@@ -50,10 +52,11 @@ class ForwardTest(parameterized.TestCase):
)
def test_jit_caching(self, xml):
"""Tests jit caching on the full step function."""
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp not installed.')
if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
if not _FORCE_TEST:
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp not installed.')
if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
batch_size = 7
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)
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
jax_logger = logging.getLogger('jax')
was_propagating = jax_logger.propagate
jax.config.update('jax_log_compiles', True)
jax_logger.propagate = False # do not print to stdout for this test
with self.assertLogs('jax', level='INFO') as log:
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.'
),
)
# 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)
self.assertEqual(step_fn._cache_size(), 1)
@parameterized.product(
xml=(
@@ -100,10 +85,11 @@ class ForwardTest(parameterized.TestCase):
batch_size=(1, 7),
)
def test_forward(self, xml: str, batch_size: int):
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp not installed.')
if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
if not _FORCE_TEST:
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp not installed.')
if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
m = test_util.load_test_file(xml)
m.opt.iterations = 10
@@ -228,10 +214,11 @@ class StepTest(parameterized.TestCase):
batch_size=(1, 7),
)
def test_step(self, xml: str, batch_size: int):
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp not installed.')
if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
if not _FORCE_TEST:
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp not installed.')
if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
m = test_util.load_test_file(xml)
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."""
import functools
import os
from absl.testing import absltest
import jax
@@ -33,6 +34,8 @@ try:
except ImportError:
smooth = None
_FORCE_TEST = os.environ.get('MJX_WARP_FORCE_TEST', '0') == '1'
class SmoothTest(absltest.TestCase):
@@ -44,10 +47,11 @@ class SmoothTest(absltest.TestCase):
def test_kinematics(self):
"""Tests kinematics with unbatched data."""
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp not installed.')
if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
if not _FORCE_TEST:
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp not installed.')
if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
m = tu.load_test_file('pendula.xml')
@@ -129,10 +133,11 @@ class SmoothTest(absltest.TestCase):
def test_kinematics_nested_vmap(self):
"""Tests kinematics with nested batch data."""
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp not installed.')
if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
if not _FORCE_TEST:
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp not installed.')
if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
m = tu.load_test_file('pendula.xml')
@@ -175,10 +180,11 @@ class SmoothTest(absltest.TestCase):
def test_kinematics_model_vmap(self):
"""Tests kinematics with vmap on model and data fields."""
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp not installed.')
if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
if not _FORCE_TEST:
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp not installed.')
if not io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')
m = tu.load_test_file('pendula.xml')
mjx/pyproject.toml
+1 -5
View File
@@ -32,11 +32,7 @@ dependencies = [
"mujoco>=3.3.5.dev0",
"scipy",
"trimesh",
]
[project.optional-dependencies]
warp = [
"warp-lang==1.8.0",
"warp-lang==1.8.1",
]
[project.scripts]
mjx/requirements.txt
-1
View File
@@ -103,7 +103,6 @@ trimesh==4.5.2 \
--hash=sha256:2e50f3a7fd135c3045da887a1b9f91230528f3ce11d2ec1ba44750d82d6b4f73
wheel==0.45.0 \
--hash=sha256:52f0baa5e6522155090a09c6bd95718cc46956d1b51d537ea5454249edb671c7
# TODO(robotics-simulation): add back warp-lang dependencies
# Transitive dependencies of etils[epath]
fsspec==2024.10.0 \