upgrade to mujoco 3.5.0

Author: thowell

mujoco_warp/_src/broadphase_test.py

@@ -26,7 +26,6 @@
 from mujoco_warp import DisableBit
 from mujoco_warp import test_data
 from mujoco_warp._src import collision_driver
-from mujoco_warp._src.io import BLEEDING_EDGE_MUJOCO
 
 
 def broadphase_caller(m, d):
@@ -186,7 +185,7 @@ def test_broadphase(self, broadphase, filter):
     (0.011, 0, 1),
     (0.00999, 0, 0),
     (0, 0.00999, 0),
-    (0.00999, 0.00999, 1 if BLEEDING_EDGE_MUJOCO else 0),
+    (0.00999, 0.00999, 1),
   )
   def test_broadphase_margin(self, margin1, margin2, ncollision):
     _MJCF = f"""

mujoco_warp/_src/collision_convex.py

@@ -24,7 +24,6 @@
 from mujoco_warp._src.collision_primitive import contact_params
 from mujoco_warp._src.collision_primitive import geom_collision_pair
 from mujoco_warp._src.collision_primitive import write_contact
-from mujoco_warp._src.io import BLEEDING_EDGE_MUJOCO
 from mujoco_warp._src.math import make_frame
 from mujoco_warp._src.math import upper_trid_index
 from mujoco_warp._src.types import MJ_MAX_EPAFACES
@@ -92,10 +91,7 @@ def _hfield_filter(
   r2 = geom_rbound[rbound_id, g2]
 
   # TODO(team): margin?
-  if BLEEDING_EDGE_MUJOCO:
-    margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2]
-  else:
-    margin = wp.max(geom_margin[margin_id, g1], geom_margin[margin_id, g2])
+  margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2]
 
   # box-sphere test: horizontal plane
   for i in range(2):

mujoco_warp/_src/collision_driver.py

@@ -20,7 +20,6 @@
 from mujoco_warp._src.collision_convex import convex_narrowphase
 from mujoco_warp._src.collision_primitive import primitive_narrowphase
 from mujoco_warp._src.collision_sdf import sdf_narrowphase
-from mujoco_warp._src.io import BLEEDING_EDGE_MUJOCO
 from mujoco_warp._src.math import upper_tri_index
 from mujoco_warp._src.types import MJ_MAXVAL
 from mujoco_warp._src.types import BroadphaseFilter
@@ -102,27 +101,18 @@ def _plane_filter(
   if size1 == 0.0:
     # geom1 is a plane
     dist = wp.dot(xpos2 - xpos1, wp.vec3(xmat1[0, 2], xmat1[1, 2], xmat1[2, 2]))
-    if BLEEDING_EDGE_MUJOCO:
-      return dist <= size2 + margin1 + margin2
-    else:
-      return dist <= size2 + wp.max(margin1, margin2)
+    return dist <= size2 + margin1 + margin2
   elif size2 == 0.0:
     # geom2 is a plane
     dist = wp.dot(xpos1 - xpos2, wp.vec3(xmat2[0, 2], xmat2[1, 2], xmat2[2, 2]))
-    if BLEEDING_EDGE_MUJOCO:
-      return dist <= size1 + margin1 + margin2
-    else:
-      return dist <= size1 + wp.max(margin1, margin2)
+    return dist <= size1 + margin1 + margin2
 
   return True
 
 
 @wp.func
 def _sphere_filter(size1: float, size2: float, margin1: float, margin2: float, xpos1: wp.vec3, xpos2: wp.vec3) -> bool:
-  if BLEEDING_EDGE_MUJOCO:
-    bound = size1 + size2 + margin1 + margin2
-  else:
-    bound = size1 + size2 + wp.max(margin1, margin2)
+  bound = size1 + size2 + margin1 + margin2
   dif = xpos2 - xpos1
   dist_sq = wp.dot(dif, dif)
   return dist_sq <= bound * bound
@@ -151,10 +141,7 @@ def _aabb_filter(
   center1 = xmat1 @ center1 + xpos1
   center2 = xmat2 @ center2 + xpos2
 
-  if BLEEDING_EDGE_MUJOCO:
-    margin = margin1 + margin2
-  else:
-    margin = wp.max(margin1, margin2)
+  margin = margin1 + margin2
 
   max_x1 = -MJ_MAXVAL
   max_y1 = -MJ_MAXVAL
@@ -249,10 +236,7 @@ def _obb_filter(
   xmat2: wp.mat33,
 ) -> bool:
   """Oriented bounding boxes collision (see Gottschalk et al.), see mj_collideOBB."""
-  if BLEEDING_EDGE_MUJOCO:
-    margin = margin1 + margin2
-  else:
-    margin = wp.max(margin1, margin2)
+  margin = margin1 + margin2
 
   xcenter = mat23()
   normal = mat63()

mujoco_warp/_src/collision_primitive.py

@@ -29,7 +29,6 @@
 from mujoco_warp._src.collision_primitive_core import sphere_capsule
 from mujoco_warp._src.collision_primitive_core import sphere_cylinder
 from mujoco_warp._src.collision_primitive_core import sphere_sphere
-from mujoco_warp._src.io import BLEEDING_EDGE_MUJOCO
 from mujoco_warp._src.math import make_frame
 from mujoco_warp._src.math import safe_div
 from mujoco_warp._src.math import upper_trid_index
@@ -553,12 +552,8 @@ def contact_params(
     solreffriction = wp.vec2(0.0, 0.0)
     solimp = mix * geom_solimp[solimp_id, g1] + (1.0 - mix) * geom_solimp[solimp_id, g2]
     # geom priority is ignored
-    if BLEEDING_EDGE_MUJOCO:
-      margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2]
-      gap = geom_gap[gap_id, g1] + geom_gap[gap_id, g2]
-    else:
-      margin = wp.max(geom_margin[margin_id, g1], geom_margin[margin_id, g2])
-      gap = wp.max(geom_gap[gap_id, g1], geom_gap[gap_id, g2])
+    margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2]
+    gap = geom_gap[gap_id, g1] + geom_gap[gap_id, g2]
 
   friction = vec5(
     wp.max(MJ_MINMU, friction[0]),

mujoco_warp/_src/io.py

@@ -30,11 +30,11 @@
 
 
 def _is_mujoco_fresh() -> bool:
-  """Checks if mujoco version is > 3.4.0."""
+  """Checks if mujoco version is > 3.5.0."""
   version = importlib.metadata.version("mujoco")
   version = version.split(".")
   version = tuple(map(int, version[:3])) + tuple(version[3:])
-  return version > (3, 4, 0)
+  return version > (3, 5, 0)
 
 
 BLEEDING_EDGE_MUJOCO = _is_mujoco_fresh()
@@ -583,17 +583,9 @@ def geom_trid_index(i, j):
       m.qM_mulm_madr.append(madr)
     m.qM_mulm_rowadr.append(len(m.qM_mulm_col))
 
-  # TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml
-  if BLEEDING_EDGE_MUJOCO:
-    m.flexedge_J_rownnz = mjm.flexedge_J_rownnz
-    m.flexedge_J_rowadr = mjm.flexedge_J_rowadr
-    m.flexedge_J_colind = mjm.flexedge_J_colind.reshape(-1)
-  else:
-    mjd = mujoco.MjData(mjm)
-    mujoco.mj_forward(mjm, mjd)
-    m.flexedge_J_rownnz = mjd.flexedge_J_rownnz
-    m.flexedge_J_rowadr = mjd.flexedge_J_rowadr
-    m.flexedge_J_colind = mjd.flexedge_J_colind.reshape(-1)
+  m.flexedge_J_rownnz = mjm.flexedge_J_rownnz
+  m.flexedge_J_rowadr = mjm.flexedge_J_rowadr
+  m.flexedge_J_colind = mjm.flexedge_J_colind.reshape(-1)
 
   # place m on device
   sizes = dict({"*": 1}, **{f.name: getattr(m, f.name) for f in dataclasses.fields(types.Model) if f.type is int})
@@ -1065,21 +1057,7 @@ def get_data_into(
   result.cinert[:] = d.cinert.numpy()[world_id]
   result.flexvert_xpos[:] = d.flexvert_xpos.numpy()[world_id]
   if mjm.nflexedge > 0:
-    # TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml
-    if not BLEEDING_EDGE_MUJOCO:
-      m = put_model(mjm)
-      result.flexedge_J_rownnz[:] = m.flexedge_J_rownnz.numpy()
-      result.flexedge_J_rowadr[:] = m.flexedge_J_rowadr.numpy()
-      result.flexedge_J_colind[:, :] = m.flexedge_J_colind.numpy().reshape((mjm.nflexedge, mjm.nv))
-      mujoco.mju_sparse2dense(
-        result.flexedge_J,
-        d.flexedge_J.numpy()[world_id].reshape(-1),
-        m.flexedge_J_rownnz.numpy(),
-        m.flexedge_J_rowadr.numpy(),
-        m.flexedge_J_colind.numpy(),
-      )
-    else:
-      result.flexedge_J[:] = d.flexedge_J.numpy()[world_id].reshape(-1)
+    result.flexedge_J[:] = d.flexedge_J.numpy()[world_id].reshape(-1)
   result.flexedge_length[:] = d.flexedge_length.numpy()[world_id]
   result.flexedge_velocity[:] = d.flexedge_velocity.numpy()[world_id]
   result.actuator_length[:] = d.actuator_length.numpy()[world_id]
@@ -2312,20 +2290,12 @@ def create_render_context(
   if render_rgb and isinstance(render_rgb, bool):
     render_rgb = [render_rgb] * ncam
   elif render_rgb is None:
-    # TODO: remove after mjwarp depends on mujoco >= 3.4.1 in pyproject.toml
-    if BLEEDING_EDGE_MUJOCO:
-      render_rgb = [mjm.cam_output[i] & mujoco.mjtCamOutBit.mjCAMOUT_RGB for i in active_cam_indices]
-    else:
-      render_rgb = [True] * ncam
+    render_rgb = [mjm.cam_output[i] & mujoco.mjtCamOutBit.mjCAMOUT_RGB for i in active_cam_indices]
 
   if render_depth and isinstance(render_depth, bool):
     render_depth = [render_depth] * ncam
   elif render_depth is None:
-    # TODO: remove after mjwarp depends on mujoco >= 3.4.1 in pyproject.toml
-    if BLEEDING_EDGE_MUJOCO:
-      render_depth = [mjm.cam_output[i] & mujoco.mjtCamOutBit.mjCAMOUT_DEPTH for i in active_cam_indices]
-    else:
-      render_depth = [True] * ncam
+    render_depth = [mjm.cam_output[i] & mujoco.mjtCamOutBit.mjCAMOUT_DEPTH for i in active_cam_indices]
 
   assert len(render_rgb) == ncam and len(render_depth) == ncam, (
     f"Render RGB and depth must be provided for all active cameras (got {len(render_rgb)}, {len(render_depth)}, expected {ncam})"
@@ -2352,10 +2322,7 @@ def create_render_context(
 
   ray = wp.zeros(int(total), dtype=wp.vec3)
 
-  # TODO: remove after mjwarp depends on mujoco >= 3.4.1 in pyproject.toml
-  cam_projection = np.zeros(mjm.ncam, dtype=int)
-  if BLEEDING_EDGE_MUJOCO:
-    cam_projection = mjm.cam_projection
+  cam_projection = mjm.cam_projection
 
   offset = 0
   for idx, cam_id in enumerate(active_cam_indices):

mujoco_warp/_src/io_test.py

@@ -306,44 +306,11 @@ def test_get_data_into_io_test_models(self, xml):
 
     # flexedge_J
     if xml == "flex/floppy.xml":
-      from mujoco_warp._src.io import BLEEDING_EDGE_MUJOCO
-
-      if BLEEDING_EDGE_MUJOCO:
-        _assert_eq(
-          mjd_result.flexedge_J.reshape(-1),
-          d.flexedge_J.numpy()[0].reshape(-1),
-          "flexedge_J",
-        )
-      else:
-        m = mjwarp.put_model(mjm)
-        _assert_eq(
-          mjd_result.flexedge_J_rownnz,
-          m.flexedge_J_rownnz.numpy(),
-          "flexedge_J_rownnz",
-        )
-        _assert_eq(
-          mjd_result.flexedge_J_rowadr,
-          m.flexedge_J_rowadr.numpy(),
-          "flexedge_J_rowadr",
-        )
-        _assert_eq(
-          mjd_result.flexedge_J_colind,
-          m.flexedge_J_colind.numpy().reshape((mjm.nflexedge, mjm.nv)),
-          "flexedge_J_colind",
-        )
-        flexedge_J = np.zeros((mjm.nflexedge, mjm.nv))
-        mujoco.mju_sparse2dense(
-          flexedge_J,
-          d.flexedge_J.numpy().reshape(-1),
-          m.flexedge_J_rownnz.numpy(),
-          m.flexedge_J_rowadr.numpy(),
-          m.flexedge_J_colind.numpy().reshape(-1),
-        )
-        _assert_eq(
-          mjd_result.flexedge_J,
-          flexedge_J,
-          "flexedge_J",
-        )
+      _assert_eq(
+        mjd_result.flexedge_J.reshape(-1),
+        d.flexedge_J.numpy()[0].reshape(-1),
+        "flexedge_J",
+      )
 
   def test_ellipsoid_fluid_model(self):
     mjm = mujoco.MjModel.from_xml_string(
@@ -1265,13 +1232,6 @@ def test_bvh_creation(self, nworld):
 
   def test_output_buffers(self):
     """Test that the output rgb and depth buffers have correct shapes and addresses."""
-    # TODO: remove after mjwarp depends on mujoco >= 3.4.1 in pyproject.toml
-    from mujoco_warp._src.io import BLEEDING_EDGE_MUJOCO
-
-    if not BLEEDING_EDGE_MUJOCO:
-      self.skipTest("Skipping test that requires mujoco >= 3.4.1")
-      return
-
     mjm, mjd, m, d = test_data.fixture(xml=_CAMERA_TEST_XML)
     width, height = 32, 24
     rc = mjwarp.create_render_context(mjm, cam_res=(width, height), render_rgb=True, render_depth=True)
@@ -1288,13 +1248,6 @@ def test_output_buffers(self):
     _assert_eq(depth_adr, [0, width * height, 2 * width * height], "depth_adr")
 
   def test_heterogeneous_camera(self):
-    # TODO: remove after mjwarp depends on mujoco >= 3.4.1 in pyproject.toml
-    from mujoco_warp._src.io import BLEEDING_EDGE_MUJOCO
-
-    if not BLEEDING_EDGE_MUJOCO:
-      self.skipTest("Skipping test that requires mujoco >= 3.4.1")
-      return
-
     """Tests render context with different resolutions and output."""
     mjm, mjd, m, d = test_data.fixture(xml=_CAMERA_TEST_XML)
     cam_res = [(64, 64), (32, 32), (16, 16)]
@@ -1320,13 +1273,6 @@ def test_heterogeneous_camera(self):
     _assert_eq(rc.depth_adr.numpy(), rc_xml.depth_adr.numpy(), "depth_adr")
 
   def test_cam_active_filtering(self):
-    # TODO: remove after mjwarp depends on mujoco >= 3.4.1 in pyproject.toml
-    from mujoco_warp._src.io import BLEEDING_EDGE_MUJOCO
-
-    if not BLEEDING_EDGE_MUJOCO:
-      self.skipTest("Skipping test that requires mujoco >= 3.4.1")
-      return
-
     mjm, mjd, m, d = test_data.fixture(xml=_CAMERA_TEST_XML)
     width, height = 32, 32
 
@@ -1339,13 +1285,6 @@ def test_cam_active_filtering(self):
 
   def test_rgb_only_and_depth_only(self):
     """Test that disabling rgb or depth correctly reduces the shape and invalidates the address."""
-    # TODO: remove after mjwarp depends on mujoco >= 3.4.1 in pyproject.toml
-    from mujoco_warp._src.io import BLEEDING_EDGE_MUJOCO
-
-    if not BLEEDING_EDGE_MUJOCO:
-      self.skipTest("Skipping test that requires mujoco >= 3.4.1")
-      return
-
     mjm, mjd, m, d = test_data.fixture(xml=_CAMERA_TEST_XML)
     width, height = 32, 32
     pixels = width * height

mujoco_warp/_src/smooth_test.py

@@ -491,17 +491,9 @@ def test_flex(self):
     mujoco.mj_comPos(mjm, mjd)
     mujoco.mj_flex(mjm, mjd)
 
-    # TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml
-    from mujoco_warp._src.io import BLEEDING_EDGE_MUJOCO
-
-    if BLEEDING_EDGE_MUJOCO:
-      rownnz = mjm.flexedge_J_rownnz
-      rowadr = mjm.flexedge_J_rowadr
-      colind = mjm.flexedge_J_colind.reshape(-1)
-    else:
-      rownnz = mjd.flexedge_J_rownnz
-      rowadr = mjd.flexedge_J_rowadr
-      colind = mjd.flexedge_J_colind.reshape(-1)
+    rownnz = mjm.flexedge_J_rownnz
+    rowadr = mjm.flexedge_J_rowadr
+    colind = mjm.flexedge_J_colind.reshape(-1)
 
     mj_flexedge_J = np.zeros((mjm.nflexedge, mjm.nv), dtype=float)
     mujoco.mju_sparse2dense(mj_flexedge_J, mjd.flexedge_J.ravel(), rownnz, rowadr, colind)

mujoco_warp/_src/types_test.py

@@ -55,32 +55,6 @@ def test_field_order(self, mj_class, mjw_class):
 
     actual_fields = [f.name for f in dataclasses.fields(mjw_class)]
 
-    # TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml
-    from mujoco_warp._src.io import BLEEDING_EDGE_MUJOCO
-
-    _FLEXEDGE_J_FIELDS = ("flexedge_J_rownnz", "flexedge_J_rowadr", "flexedge_J_colind")
-
-    def _remove_fields(fields, to_remove):
-      for field in to_remove:
-        if field in fields:
-          fields.remove(field)
-
-    if not BLEEDING_EDGE_MUJOCO:
-      _remove_fields(actual_fields, _FLEXEDGE_J_FIELDS)
-      _remove_fields(desired_fields, _FLEXEDGE_J_FIELDS)
-      _remove_fields(
-        actual_fields,
-        [
-          "cam_projection",
-        ],
-      )
-      _remove_fields(
-        desired_fields,
-        [
-          "cam_projection",
-        ],
-      )
-
     self.assertListEqual(actual_fields, desired_fields)
 
   @parameterized.parameters(Option, Model, Data)

pyproject.toml

@@ -27,7 +27,7 @@ requires-python = ">=3.10"
 dependencies = [
     "absl-py",
     "etils[epath]",
-    "mujoco>=3.4.0",
+    "mujoco>=3.5.0",
     "numpy",
     "warp-lang>=1.11.0",
 ]
@@ -54,7 +54,7 @@ dev = [
     "ruff",
     "pygls>=1.0.0,<2.0.0",
     "lsprotocol>=2023.0.1,<2024.0.0",
-    "mujoco>=3.4.1.dev0",
+    "mujoco>=3.5.0.dev0",
     "warp-lang>=1.11.0.dev0",
 ]
 # TODO(team): cpu and cuda JAX optional dependencies are temporary, remove after we land MJX:Warp