Import google-deepmind/mujoco_warp from GitHub.

PiperOrigin-RevId: 807766737
Change-Id: I1c22361ffa61f27fd8e0fd534b96488531761916

Author: thowell

mjx/cuda_requirements.txt

@@ -16,6 +16,6 @@ jax-cuda12-pjrt==0.5.3; python_version >= '3.10' \
 jax-cuda12-pjrt==0.4.30; python_version == '3.9' \
     --hash=sha256:895d0198ad99638fcaf976c47592e2a543eef79ea15fabd24a402d055390c328 \
     --hash=sha256:c36fb1e0c236563bf3a87e70f4d1ab28a31d7cf5d722c9ede30c4172116e8bcb
-warp-lang==1.8.1 \
-    --hash=sha256:cfc59e1070ad71531b5d83186de48162507277af344a102fa33d5df9cdb942f7 \
-    --hash=sha256:1db9ca92c46902b76bb99565c544347d1a32e9fb875ce902f1cafb94978d1ac3
+warp-lang==1.9.0 \
+    --hash=sha256:23165d3291eeecc5ac47b9a3de0b93d34b4c921414e5a761ab96d68861efa3cb \
+    --hash=sha256:7ea8057e5d6fdb9b0885d725b51954c5fc6dac62d6ac007828148b1cb13ee615

mjx/mujoco/mjx/_src/io_test.py

@@ -15,7 +15,9 @@
 """Tests for io functions."""
 
 import os
+import tempfile
 from unittest import mock
+
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
@@ -31,6 +33,7 @@
 # pylint: enable=g-importing-member
 import mujoco.mjx.warp as mjxw
 from mujoco.mjx.warp import types as mjxw_types
+from mujoco.mjx.warp import warp as wp  # pylint: disable=g-importing-member
 import numpy as np
 
 
@@ -125,6 +128,12 @@ def _get_name_from_path(path: jax.tree_util.KeyPath) -> str:
 class ModelIOTest(parameterized.TestCase):
   """IO tests for mjx.Model."""
 
+  def setUp(self):
+    super().setUp()
+    if mjxw.WARP_INSTALLED:
+      self.tempdir = tempfile.TemporaryDirectory()
+      wp.config.kernel_cache_dir = self.tempdir.name
+
   @parameterized.product(
       xml=(_MULTIPLE_CONVEX_OBJECTS, _MULTIPLE_CONSTRAINTS),
       impl=('jax', 'c', 'warp'),
@@ -326,6 +335,12 @@ def check_ndim(path, x):
 class DataIOTest(parameterized.TestCase):
   """IO tests for mjx.Data."""
 
+  def setUp(self):
+    super().setUp()
+    if mjxw.WARP_INSTALLED:
+      self.tempdir = tempfile.TemporaryDirectory()
+      wp.config.kernel_cache_dir = self.tempdir.name
+
   @parameterized.parameters('jax', 'c')
   def test_make_data(self, impl: str):
     """Test that make_data returns the correct shapes."""

mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py

@@ -21,6 +21,7 @@
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data as Data
 # isort: on
 
+from ._src import test_util as test_util  # used by viewer and testspeed, not meant for public consumption
 from mujoco.mjx.third_party.mujoco_warp._src.collision_driver import collision as collision
 from mujoco.mjx.third_party.mujoco_warp._src.collision_driver import nxn_broadphase as nxn_broadphase
 from mujoco.mjx.third_party.mujoco_warp._src.collision_driver import sap_broadphase as sap_broadphase
@@ -43,6 +44,7 @@
 from mujoco.mjx.third_party.mujoco_warp._src.io import make_data as make_data
 from mujoco.mjx.third_party.mujoco_warp._src.io import put_data as put_data
 from mujoco.mjx.third_party.mujoco_warp._src.io import put_model as put_model
+from mujoco.mjx.third_party.mujoco_warp._src.io import reset_data as reset_data
 from mujoco.mjx.third_party.mujoco_warp._src.passive import passive as passive
 from mujoco.mjx.third_party.mujoco_warp._src.ray import ray as ray
 from mujoco.mjx.third_party.mujoco_warp._src.sensor import energy_pos as energy_pos
@@ -66,8 +68,6 @@
 from mujoco.mjx.third_party.mujoco_warp._src.support import contact_force as contact_force
 from mujoco.mjx.third_party.mujoco_warp._src.support import mul_m as mul_m
 from mujoco.mjx.third_party.mujoco_warp._src.support import xfrc_accumulate as xfrc_accumulate
-from mujoco.mjx.third_party.mujoco_warp._src.test_util import BenchmarkSuite as BenchmarkSuite
-from mujoco.mjx.third_party.mujoco_warp._src.test_util import benchmark as benchmark
 from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseFilter as BroadphaseFilter
 from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseType as BroadphaseType
 from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType as ConeType

mjx/mujoco/mjx/third_party/mujoco_warp/_src/broadphase_test.py

@@ -326,9 +326,6 @@ def test_broadphase_filter(self):
     broadphase_caller(m, d)
     self.assertEqual(d.ncollision.numpy()[0], 0)
 
-  # TODO(team): test margin
-  # TODO(team): test DisableBit.FILTERPARENT
-
 
 if __name__ == "__main__":
   wp.init()