From f255122ec76ee2cb3215692810a924d630bea1df Mon Sep 17 00:00:00 2001 From: guyazran Date: Wed, 11 Oct 2023 12:44:16 +0300 Subject: [PATCH] More control over offscreen dimensions and scene geometries (#731) --- gymnasium/envs/mujoco/mujoco_env.py | 8 ++- gymnasium/envs/mujoco/mujoco_rendering.py | 38 +++++++--- tests/envs/mujoco/test_mujoco_rendering.py | 83 ++++++++++++++++++++++ 3 files changed, 117 insertions(+), 12 deletions(-) create mode 100644 tests/envs/mujoco/test_mujoco_rendering.py diff --git a/gymnasium/envs/mujoco/mujoco_env.py b/gymnasium/envs/mujoco/mujoco_env.py index 7b374c84c..1c1a986fe 100644 --- a/gymnasium/envs/mujoco/mujoco_env.py +++ b/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( diff --git a/gymnasium/envs/mujoco/mujoco_rendering.py b/gymnasium/envs/mujoco/mujoco_rendering.py index 78abbbeec..a5681a68d 100644 --- a/gymnasium/envs/mujoco/mujoco_rendering.py +++ b/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,6 +635,7 @@ 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. @@ -632,6 +643,9 @@ def __init__( 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" diff --git a/tests/envs/mujoco/test_mujoco_rendering.py b/tests/envs/mujoco/test_mujoco_rendering.py new file mode 100644 index 000000000..14cf441c2 --- /dev/null +++ b/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()