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: Optional[int] = None,
     ):
         """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 or 1000)
         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: Optional[int] = None,
+    ):
+        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: Optional[int] = None,
     ):
         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: Optional[int] = None,
     ):
         """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,
@@ -695,16 +740,17 @@ def _get_viewer(self, render_mode: str):
         self.viewer = self._viewers.get(render_mode)
         if self.viewer is None:
             if render_mode == "human":
-                self.viewer = WindowViewer(
-                    self.model, self.data, self.width, self.height
-                )
-
+                viewer_class = WindowViewer
             elif render_mode in {"rgb_array", "depth_array"}:
-                self.viewer = OffScreenViewer(self.model, self.data)
+                viewer_class = OffScreenViewer
             else:
                 raise AttributeError(
                     f"Unexpected mode: {render_mode}, expected modes: human, rgb_array, or depth_array"
                 )
+            self.viewer = viewer_class(
+                self.model, self.data, self.width, self.height, self.max_geom
+            )
+
             # Add default camera parameters
             self._set_cam_config()
             self._viewers[render_mode] = self.viewer