Fix MuJoCo add_markers for mujoco>=3.2 (#1329)

Author: amacati

Diff for: gymnasium/envs/mujoco/mujoco_rendering.py

@@ -6,10 +6,16 @@
 import imageio
 import mujoco
 import numpy as np
+from packaging.version import Version
 
 from gymnasium.logger import warn
 
 
+# The marker API changed in MuJoCo 3.2.0, so we check the mujoco version and set a flag that
+# determines which function we use when adding markers to the scene.
+_MUJOCO_MARKER_LEGACY_MODE = Version(mujoco.__version__) < Version("3.2.0")
+
+
 def _import_egl(width, height):
     from mujoco.egl import GLContext
 
@@ -88,8 +94,37 @@ def add_marker(self, **marker_params):
 
     def _add_marker_to_scene(self, marker: dict):
         if self.scn.ngeom >= self.scn.maxgeom:
-            raise RuntimeError("Ran out of geoms. maxgeom: %d" % self.scn.maxgeom)
+            raise RuntimeError(f"Ran out of geoms. maxgeom: {self.scn.maxgeom}")
+
+        if _MUJOCO_MARKER_LEGACY_MODE:  # Old API for markers requires special handling
+            self._legacy_add_marker_to_scene(marker)
+        else:
+            geom_type = marker.get("type", mujoco.mjtGeom.mjGEOM_SPHERE)
+            size = marker.get("size", np.array([0.01, 0.01, 0.01]))
+            pos = marker.get("pos", np.array([0.0, 0.0, 0.0]))
+            mat = marker.get("mat", np.eye(3).flatten())
+            rgba = marker.get("rgba", np.array([1.0, 1.0, 1.0, 1.0]))
+            mujoco.mjv_initGeom(
+                self.scn.geoms[self.scn.ngeom],
+                geom_type,
+                size=size,
+                pos=pos,
+                mat=mat,
+                rgba=rgba,
+            )
+
+        self.scn.ngeom += 1
 
+    def _legacy_add_marker_to_scene(self, marker: dict):
+        """Add a marker to the scene compatible with older versions of MuJoCo.
+
+        MuJoCo 3.2 introduced breaking changes to the visual geometries API. To maintain
+        compatibility with older versions, we use the legacy API when an older version of MuJoCo is
+        detected.
+
+        Args:
+            marker: A dictionary containing the marker parameters.
+        """
         g = self.scn.geoms[self.scn.ngeom]
         # default values.
         g.dataid = -1
@@ -130,8 +165,6 @@ def _add_marker_to_scene(self, marker: dict):
             else:
                 raise ValueError("mjtGeom doesn't have field %s" % key)
 
-        self.scn.ngeom += 1
-
     def close(self):
         """Override close in your rendering subclass to perform any necessary cleanup
         after env.close() is called.

Diff for: pyproject.toml

@@ -41,7 +41,7 @@ classic-control = ["pygame >=2.1.3"]
 classic_control = ["pygame >=2.1.3"]  # kept for backward compatibility
 mujoco-py = ["mujoco-py >=2.1,<2.2", "cython<3"]
 mujoco_py = ["mujoco-py >=2.1,<2.2", "cython<3"]  # kept for backward compatibility
-mujoco = ["mujoco >=2.1.5", "imageio >=2.14.1"]
+mujoco = ["mujoco >=2.1.5", "imageio >=2.14.1", "packaging >=23.0"]
 toy-text = ["pygame >=2.1.3"]
 toy_text = ["pygame >=2.1.3"]  # kept for backward compatibility
 jax = ["jax >=0.4.16", "jaxlib >=0.4.16", "flax >=0.5.0"]
@@ -64,6 +64,7 @@ all = [
     # mujoco
     "mujoco >=2.1.5",
     "imageio >=2.14.1",
+    "packaging >=23.0",
     # toy-text
     "pygame >=2.1.3",
     # jax

Diff for: tests/envs/mujoco/test_mujoco_rendering.py

@@ -117,6 +117,28 @@ def test_max_geom_attribute(
     viewer.close()
 
 
+@pytest.mark.parametrize(
+    "render_mode", ["human", "rgb_array", "depth_array", "rgbd_tuple"]
+)
+def test_add_markers(model: mujoco.MjModel, data: mujoco.MjData, render_mode: str):
+    """Test that the add_markers function works correctly."""
+    # initialize renderer
+    renderer = ExposedViewerRenderer(
+        model, data, width=DEFAULT_SIZE, height=DEFAULT_SIZE, max_geom=10
+    )
+    # initialize viewer via render
+    viewer = renderer.get_viewer(render_mode)
+    viewer.add_marker(
+        pos=np.array([0, 0, 0]),
+        size=np.array([1, 1, 1]),
+        rgba=np.array([1, 0, 0, 1]),
+    )
+    args = tuple() if render_mode == "human" else (render_mode,)
+    viewer.render(*args)  # We need to render to trigger the marker addition in MuJoCo
+    # close viewer after usage
+    viewer.close()
+
+
 @pytest.mark.parametrize(
     "render_mode", ["human", "rgb_array", "depth_array", "rgbd_tuple"]
 )