More control over offscreen dimensions and scene geometries

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,48 @@ 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: Optional[int] = None,
+        height: Optional[int] = None,
+        max_geom: int = 1000,
+    ):
+        buffer_width = model.vis.global_.offwidth
+        buffer_height = model.vis.global_.offheight
+
+        width = width or buffer_width
+        height = height or buffer_height
+
+        # check if the framebuffer is large enough to handle the requested image dimensions.
+        # same check as in `mujoco.Renderer` class
+        if width > buffer_width:
+            raise ValueError(
+                f"""
+Image width {width} > framebuffer width {buffer_width}. Either reduce the image
+width or specify a larger offscreen framebuffer in the model XML using the
+clause:
+<visual>
+<global offwidth="my_width"/>
+</visual>""".lstrip()
+            )
+
+        if height > buffer_height:
+            raise ValueError(
+                f"""
+Image height {height} > framebuffer height {buffer_height}. Either reduce the
+image height or specify a larger offscreen framebuffer in the model XML using
+the clause:
+<visual>
+<global offheight="my_height"/>
+</visual>""".lstrip()
+            )
 
         # 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 +326,7 @@ def __init__(
         data: "mujoco.MjData",
         width: Optional[int] = None,
         height: Optional[int] = None,
+        max_geom: int = 1000,
     ):
         glfw.init()
 
@@ -325,7 +365,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 +665,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 +684,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 +741,14 @@ 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,86 @@
+import os
+
+import mujoco
+import pytest
+
+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_FRAMEBUFFER_WIDTH = 640
+DEFAULT_FRAMEBUFFER_HEIGHT = 480
+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_FRAMEBUFFER_WIDTH
+    model.vis.global_.offheight = DEFAULT_FRAMEBUFFER_HEIGHT
+    return model
+
+
+@pytest.fixture(scope="module")
+def data(model):
+    """Initialize data."""
+    return mujoco.MjData(model)
+
+
+@pytest.mark.parametrize("width", [10, 100, 1000, None])
+@pytest.mark.parametrize("height", [10, 100, 1000, None])
+@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."""
+
+    # set default buffer dimensions if no dims are given
+    check_width = width or DEFAULT_FRAMEBUFFER_WIDTH
+    check_height = height or DEFAULT_FRAMEBUFFER_HEIGHT
+
+    # check for "dimensions too big" error
+    if (
+        check_width > DEFAULT_FRAMEBUFFER_WIDTH
+        or check_height > DEFAULT_FRAMEBUFFER_HEIGHT
+    ):
+        # after ValueError, AttributeError is raised on call to __del__
+        with pytest.raises((ValueError, AttributeError)):
+            viewer = OffScreenViewer(model, data, width=width, height=height)
+        return
+
+    # initialize viewer
+    viewer = OffScreenViewer(model, data, width=width, height=height)
+
+    # assert viewer dimensions
+    assert viewer.viewport.width == check_width
+    assert viewer.viewport.height == check_height
+
+
+@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, 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