More control over offscreen dimensions and scene geometries (#731)

gymnasium/envs/mujoco/mujoco_env.py

@@ -354,6 +354,7 @@ def __init__(
         camera_id: Optional[int] = None,
         camera_name: Optional[str] = None,
         default_camera_config: Optional[Dict[str, Union[float, int]]] = None,
+        max_geom: int = 1000,
     ):
         if MUJOCO_IMPORT_ERROR is not None:
             raise error.DependencyNotInstalled(
@@ -375,7 +376,12 @@ def __init__(
         from gymnasium.envs.mujoco.mujoco_rendering import MujocoRenderer
 
         self.mujoco_renderer = MujocoRenderer(
-            self.model, self.data, default_camera_config, self.width, self.height
+            self.model,
+            self.data,
+            default_camera_config,
+            self.width,
+            self.height,
+            max_geom,
         )
 
     def _initialize_simulation(

gymnasium/envs/mujoco/mujoco_rendering.py

@@ -38,7 +38,12 @@ def _import_osmesa(width, height):
 
 class BaseRender:
     def __init__(
-        self, model: "mujoco.MjModel", data: "mujoco.MjData", width: int, height: int
+        self,
+        model: "mujoco.MjModel",
+        data: "mujoco.MjData",
+        width: int,
+        height: int,
+        max_geom: int = 1000,
     ):
         """Render context superclass for offscreen and window rendering."""
         self.model = model
@@ -50,7 +55,7 @@ def __init__(
         self.viewport = mujoco.MjrRect(0, 0, width, height)
 
         # This goes to specific visualizer
-        self.scn = mujoco.MjvScene(self.model, 1000)
+        self.scn = mujoco.MjvScene(self.model, max_geom)
         self.cam = mujoco.MjvCamera()
         self.vopt = mujoco.MjvOption()
         self.pert = mujoco.MjvPerturb()
@@ -134,14 +139,18 @@ def close(self):
 class OffScreenViewer(BaseRender):
     """Offscreen rendering class with opengl context."""
 
-    def __init__(self, model: "mujoco.MjMujoco", data: "mujoco.MjData"):
-        width = model.vis.global_.offwidth
-        height = model.vis.global_.offheight
-
+    def __init__(
+        self,
+        model: "mujoco.MjMujoco",
+        data: "mujoco.MjData",
+        width: int,
+        height: int,
+        max_geom: int = 1000,
+    ):
         # We must make GLContext before MjrContext
         self._get_opengl_backend(width, height)
 
-        super().__init__(model, data, width, height)
+        super().__init__(model, data, width, height, max_geom)
 
         self._init_camera()
 
@@ -287,6 +296,7 @@ def __init__(
         data: "mujoco.MjData",
         width: Optional[int] = None,
         height: Optional[int] = None,
+        max_geom: int = 1000,
     ):
         glfw.init()
 
@@ -325,7 +335,7 @@ def __init__(
         glfw.set_scroll_callback(self.window, self._scroll_callback)
         glfw.set_key_callback(self.window, self._key_callback)
 
-        super().__init__(model, data, width, height)
+        super().__init__(model, data, width, height, max_geom)
         glfw.swap_interval(1)
 
     def _set_mujoco_buffer(self):
@@ -625,13 +635,17 @@ def __init__(
         default_cam_config: Optional[dict] = None,
         width: Optional[int] = None,
         height: Optional[int] = None,
+        max_geom: int = 1000,
     ):
         """A wrapper for clipping continuous actions within the valid bound.
 
         Args:
             model: MjModel data structure of the MuJoCo simulation
             data: MjData data structure of the MuJoCo simulation
             default_cam_config: dictionary with attribute values of the viewer's default camera, https://mujoco.readthedocs.io/en/latest/XMLreference.html?highlight=camera#visual-global
+            width: width of the OpenGL rendering context
+            height: height of the OpenGL rendering context
+            max_geom: maximum number of geometries to render
         """
         self.model = model
         self.data = data
@@ -640,6 +654,7 @@ def __init__(
         self.default_cam_config = default_cam_config
         self.width = width
         self.height = height
+        self.max_geom = max_geom
 
     def render(
         self,
@@ -696,11 +711,12 @@ def _get_viewer(self, render_mode: str):
         if self.viewer is None:
             if render_mode == "human":
                 self.viewer = WindowViewer(
-                    self.model, self.data, self.width, self.height
+                    self.model, self.data, self.width, self.height, self.max_geom
                 )
-
             elif render_mode in {"rgb_array", "depth_array"}:
-                self.viewer = OffScreenViewer(self.model, self.data)
+                self.viewer = OffScreenViewer(
+                    self.model, self.data, self.width, self.height, self.max_geom
+                )
             else:
                 raise AttributeError(
                     f"Unexpected mode: {render_mode}, expected modes: human, rgb_array, or depth_array"

tests/envs/mujoco/test_mujoco_rendering.py

@@ -0,0 +1,83 @@
+import os
+
+import mujoco
+import pytest
+
+from gymnasium.envs.mujoco.mujoco_env import DEFAULT_SIZE
+from gymnasium.envs.mujoco.mujoco_rendering import MujocoRenderer, OffScreenViewer
+
+
+ASSET_PATH = os.path.join(
+    os.path.dirname(__file__), "assets", "walker2d_v5_uneven_feet.xml"
+)
+DEFAULT_MAX_GEOMS = 1000
+
+
+class ExposedViewerRenderer(MujocoRenderer):
+    """Expose the viewer for testing to avoid warnings."""
+
+    def get_viewer(self, render_mode: str):
+        return self._get_viewer(render_mode)
+
+
+@pytest.fixture(scope="module")
+def model():
+    """Initialize a model."""
+    model = mujoco.MjModel.from_xml_path(ASSET_PATH)
+    model.vis.global_.offwidth = DEFAULT_SIZE
+    model.vis.global_.offheight = DEFAULT_SIZE
+    return model
+
+
+@pytest.fixture(scope="module")
+def data(model):
+    """Initialize data."""
+    return mujoco.MjData(model)
+
+
+@pytest.mark.parametrize("width", [10, 100, 200, 480])
+@pytest.mark.parametrize("height", [10, 100, 200, 480])
+@pytest.mark.filterwarnings("ignore::UserWarning")
+def test_offscreen_viewer_custom_dimensions(
+    model: mujoco.MjModel, data: mujoco.MjData, width: int, height: int
+):
+    """Test that the offscreen viewer has the correct dimensions."""
+
+    # initialize viewer
+    viewer = OffScreenViewer(model, data, width=width, height=height)
+
+    # assert viewer dimensions
+    assert viewer.viewport.width == width
+    assert viewer.viewport.height == height
+
+    # check that the render method returns an image of the correct shape
+    img = viewer.render(render_mode="rgb_array")
+    assert img.shape == (height, width, 3)
+
+    # close viewer after usage
+    viewer.close()
+
+
+@pytest.mark.parametrize("render_mode", ["human", "rgb_array", "depth_array"])
+@pytest.mark.parametrize("max_geom", [10, 100, 1000, 10000])
+def test_max_geom_attribute(
+    model: mujoco.MjModel, data: mujoco.MjData, render_mode: str, max_geom: int
+):
+    """Test that the max_geom attribute is set correctly."""
+
+    # initialize renderer
+    renderer = ExposedViewerRenderer(
+        model, data, width=DEFAULT_SIZE, height=DEFAULT_SIZE, max_geom=max_geom
+    )
+
+    # assert max_geom attribute
+    assert renderer.max_geom == max_geom
+
+    # initialize viewer via render
+    viewer = renderer.get_viewer(render_mode)
+
+    # assert that max_geom is set correctly in the viewer scene
+    assert viewer.scn.maxgeom == max_geom
+
+    # close viewer after usage
+    viewer.close()