Native Viewer Support (#551)

Author: oliverweissl

examples/1_simulator_basics/1a_simulate_single_robot/main.py

@@ -73,12 +73,15 @@ def main() -> None:
     """
     After we have the scene ready we create a simulator that will perform the simulation.
     This tutorial chooses to use Mujoco, but your version of revolve might contain other simulators as well.
+    
+    For mujoco we can select either the `native` mujoco viewer (more performance) or our `custom` viewer (which is more flexible for adjustments).
     """
-    simulator = LocalSimulator()
+    simulator = LocalSimulator(viewer_type="native")
 
     # `batch_parameters` are important parameters for simulation.
     # Here, we use the parameters that are standard in CI Group.
     batch_parameters = make_standard_batch_parameters()
+    batch_parameters.simulation_time = 60  # Here we update our simulation time.
 
     # Simulate the scene.
     # A simulator can run multiple sets of scenes sequentially; it can be reused.

examples/2_modular_robot_basics/2b_brain_with_feedback/main.py

@@ -136,7 +136,7 @@ def main() -> None:
     scene.add_robot(robot)
 
     # Simulate the scene.
-    simulator = LocalSimulator()
+    simulator = LocalSimulator(viewer_type="native")
     simulate_scenes(
         simulator=simulator,
         batch_parameters=make_standard_batch_parameters(),

simulation/revolve2/simulation/simulator/__init__.py

@@ -4,5 +4,6 @@
 from ._batch_parameters import BatchParameters
 from ._record_settings import RecordSettings
 from ._simulator import Simulator
+from ._viewer import Viewer
 
-__all__ = ["Batch", "BatchParameters", "RecordSettings", "Simulator"]
+__all__ = ["Batch", "BatchParameters", "RecordSettings", "Simulator", "Viewer"]

simulation/revolve2/simulation/simulator/_viewer.py

@@ -0,0 +1,53 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class Viewer(ABC):
+    """An abstract viewer class, enabling the rendering of simulations."""
+
+    @abstractmethod
+    def close_viewer(self) -> None:
+        """Close the viewer."""
+
+    @abstractmethod
+    def render(self) -> Any | None:
+        """
+        Render the scene on the viewer.
+
+        :returns: Nothing or feedback.
+        """
+
+    @abstractmethod
+    def current_viewport_size(self) -> tuple[int, int]:
+        """
+        Get the current viewport size.
+
+        :returns: The size as a tuple.
+        """
+
+    @property
+    @abstractmethod
+    def view_port(self) -> Any:
+        """
+        Get the viewport.
+
+        :returns: The viewport object.
+        """
+
+    @property
+    @abstractmethod
+    def context(self) -> Any:
+        """
+        Return the viewer context.
+
+        :returns: The context object.
+        """
+
+    @property
+    @abstractmethod
+    def can_record(self) -> bool:
+        """
+        Check if this viewer can record.
+
+        :returns: The truth.
+        """

simulators/mujoco_simulator/revolve2/simulators/mujoco_simulator/_local_simulator.py

@@ -7,6 +7,7 @@
 
 from ._simulate_manual_scene import simulate_manual_scene
 from ._simulate_scene import simulate_scene
+from .viewers import ViewerType
 
 
 class LocalSimulator(Simulator):
@@ -18,6 +19,7 @@ class LocalSimulator(Simulator):
     _cast_shadows: bool
     _fast_sim: bool
     _manual_control: bool
+    _viewer_type: ViewerType
 
     def __init__(
         self,
@@ -27,6 +29,7 @@ def __init__(
         cast_shadows: bool = False,
         fast_sim: bool = False,
         manual_control: bool = False,
+        viewer_type: ViewerType | str = ViewerType.CUSTOM,
     ):
         """
         Initialize this object.
@@ -37,6 +40,7 @@ def __init__(
         :param cast_shadows: Whether shadows are cast in the simulation.
         :param fast_sim: Whether more complex rendering prohibited.
         :param manual_control: Whether the simulation should be controlled manually.
+        :param viewer_type: The viewer-implementation to use in the local simulator.
         """
         assert (
             headless or num_simulators == 1
@@ -52,6 +56,11 @@ def __init__(
         self._cast_shadows = cast_shadows
         self._fast_sim = fast_sim
         self._manual_control = manual_control
+        self._viewer_type = (
+            ViewerType.from_string(viewer_type)
+            if isinstance(viewer_type, str)
+            else viewer_type
+        )
 
     def simulate_batch(self, batch: Batch) -> list[list[SimulationState]]:
         """
@@ -101,6 +110,7 @@ def simulate_batch(self, batch: Batch) -> list[list[SimulationState]]:
                         batch.parameters.simulation_timestep,
                         self._cast_shadows,
                         self._fast_sim,
+                        self._viewer_type,
                     )
                     for scene_index, scene in enumerate(batch.scenes)
                 ]
@@ -119,6 +129,7 @@ def simulate_batch(self, batch: Batch) -> list[list[SimulationState]]:
                     batch.parameters.simulation_timestep,
                     self._cast_shadows,
                     self._fast_sim,
+                    self._viewer_type,
                 )
                 for scene_index, scene in enumerate(batch.scenes)
             ]

simulators/mujoco_simulator/revolve2/simulators/mujoco_simulator/_simulate_manual_scene.py

@@ -11,10 +11,10 @@
 from revolve2.simulation.scene import Scene
 
 from ._control_interface_impl import ControlInterfaceImpl
-from ._custom_mujoco_viewer import CustomMujocoViewer, CustomMujocoViewerMode
 from ._render_backend import RenderBackend
 from ._scene_to_model import scene_to_model
 from ._simulation_state_impl import SimulationStateImpl
+from .viewers import CustomMujocoViewer, CustomMujocoViewerMode
 
 
 def simulate_manual_scene(
@@ -94,5 +94,5 @@ def simulate_manual_scene(
         """If we press ctrl-C this script will end with the finally clause."""
         pass
     finally:
-        viewer.close()
+        viewer.close_viewer()
         logging.info("Testing done.")

simulators/mujoco_simulator/revolve2/simulators/mujoco_simulator/_simulate_scene.py

@@ -10,11 +10,11 @@
 from revolve2.simulation.simulator import RecordSettings
 
 from ._control_interface_impl import ControlInterfaceImpl
-from ._custom_mujoco_viewer import CustomMujocoViewer
 from ._open_gl_vision import OpenGLVision
 from ._render_backend import RenderBackend
 from ._scene_to_model import scene_to_model
 from ._simulation_state_impl import SimulationStateImpl
+from .viewers import CustomMujocoViewer, NativeMujocoViewer, ViewerType
 
 
 def simulate_scene(
@@ -29,6 +29,7 @@ def simulate_scene(
     simulation_timestep: float,
     cast_shadows: bool,
     fast_sim: bool,
+    viewer_type: ViewerType,
     render_backend: RenderBackend = RenderBackend.EGL,
 ) -> list[SimulationState]:
     """
@@ -45,45 +46,73 @@ def simulate_scene(
     :param simulation_timestep: The duration to integrate over during each step of the simulation. In seconds.
     :param cast_shadows: If shadows are cast.
     :param fast_sim: If fancy rendering is disabled.
+    :param viewer_type: The type of viewer used for the rendering in a window.
     :param render_backend: The backend to be used for rendering (EGL by default and switches to GLFW if no cameras are on the robot).
     :returns: The results of simulation. The number of returned states depends on `sample_step`.
+    :raises ValueError: If the viewer is not able to record.
     """
     logging.info(f"Simulating scene {scene_id}")
 
+    """Define mujoco data and model objects for simuating."""
     model, mapping = scene_to_model(
         scene, simulation_timestep, cast_shadows=cast_shadows, fast_sim=fast_sim
     )
+    data = mujoco.MjData(model)
+
+    """Define a control interface for the mujoco simulation (used to control robots)."""
+    control_interface = ControlInterfaceImpl(
+        data=data, abstraction_to_mujoco_mapping=mapping
+    )
+    """Make separate viewer for camera sensors."""
+    camera_viewers = {
+        camera.camera_id: OpenGLVision(
+            model=model, camera=camera, headless=headless, open_gl_lib=render_backend
+        )
+        for camera in mapping.camera_sensor.values()
+    }
+
+    """Define some additional control variables."""
+    last_control_time = 0.0
+    last_sample_time = 0.0
+    last_video_time = 0.0  # time at which last video frame was saved
+
+    simulation_states: list[SimulationState] = (
+        []
+    )  # The measured states of the simulation
 
     """If we dont have cameras and the backend is not set we go to the default GLFW."""
     if len(mapping.camera_sensor.values()) == 0:
         render_backend = RenderBackend.GLFW
-    data = mujoco.MjData(model)
 
+    """Initialize viewer object if we need to render the scene."""
     if not headless or record_settings is not None:
-        width, height = (
-            (record_settings.width, record_settings.height)
-            if record_settings is not None
-            else (None, None)
-        )
-        viewer = CustomMujocoViewer(
+        match viewer_type:
+            case viewer_type.CUSTOM:
+                viewer = CustomMujocoViewer
+            case viewer_type.NATIVE:
+                viewer = NativeMujocoViewer
+            case _:
+                raise ValueError(
+                    f"Viewer of type {viewer_type} not defined in _simulate_scene."
+                )
+
+        viewer = viewer(
             model,
             data,
-            width=width,
-            height=height,
+            width=None if record_settings is None else record_settings.width,
+            height=None if record_settings is None else record_settings.height,
             backend=render_backend,
             start_paused=start_paused,
             render_every_frame=False,
             hide_menus=(record_settings is not None),
         )
 
-    camera_viewers = {
-        camera.camera_id: OpenGLVision(
-            model=model, camera=camera, headless=headless, open_gl_lib=render_backend
-        )
-        for camera in mapping.camera_sensor.values()
-    }
-
+    """Record the scene if we want to record."""
     if record_settings is not None:
+        if not viewer.can_record:
+            raise ValueError(
+                f"Selected Viewer {type(viewer).__name__} has no functionality to record."
+            )
         video_step = 1 / record_settings.fps
         video_file_path = f"{record_settings.video_directory}/{scene_id}.mp4"
         fourcc = cv2.VideoWriter.fourcc(*"mp4v")
@@ -94,15 +123,10 @@ def simulate_scene(
             viewer.current_viewport_size(),
         )
 
-    last_control_time = 0.0
-    last_sample_time = 0.0
-    last_video_time = 0.0  # time at which last video frame was saved
-
-    # The measured states of the simulation
-    simulation_states: list[SimulationState] = []
-
-    # Compute forward dynamics without actually stepping forward in time.
-    # This updates the data so we can read out the initial state.
+    """
+    Compute forward dynamics without actually stepping forward in time.
+    This updates the data so we can read out the initial state.
+    """
     mujoco.mj_forward(model, data)
     images = {
         camera_id: camera_viewer.process(model, data)
@@ -117,9 +141,7 @@ def simulate_scene(
             )
         )
 
-    control_interface = ControlInterfaceImpl(
-        data=data, abstraction_to_mujoco_mapping=mapping
-    )
+    """After rendering the initial state, we enter the rendering loop."""
     while (time := data.time) < (
         float("inf") if simulation_time is None else simulation_time
     ):
@@ -147,6 +169,7 @@ def simulate_scene(
 
         # step simulation
         mujoco.mj_step(model, data)
+        # extract images from camera sensors.
         images = {
             camera_id: camera_viewer.process(model, data)
             for camera_id, camera_viewer in camera_viewers.items()
@@ -164,22 +187,23 @@ def simulate_scene(
 
             # https://github.com/deepmind/mujoco/issues/285 (see also record.cc)
             img: npt.NDArray[np.uint8] = np.empty(
-                (viewer.viewport.height, viewer.viewport.width, 3),
+                (*viewer.current_viewport_size(), 3),
                 dtype=np.uint8,
             )
 
             mujoco.mjr_readPixels(
                 rgb=img,
                 depth=None,
-                viewport=viewer.viewport,
-                con=viewer.ctx,
+                viewport=viewer.view_port,
+                con=viewer.context,
             )
             # Flip the image and map to OpenCV colormap (BGR -> RGB)
             img = np.flipud(img)[:, :, ::-1]
             video.write(img)
 
+    """Once simulation is done we close the potential viewer and release the potential video."""
     if not headless or record_settings is not None:
-        viewer.close()
+        viewer.close_viewer()
 
     if record_settings is not None:
         video.release()
@@ -193,5 +217,4 @@ def simulate_scene(
         )
 
     logging.info(f"Scene {scene_id} done.")
-
     return simulation_states

simulators/mujoco_simulator/revolve2/simulators/mujoco_simulator/viewers/__init__.py

@@ -0,0 +1,12 @@
+"""Different viewer implementations for mujoco."""
+
+from ._custom_mujoco_viewer import CustomMujocoViewer, CustomMujocoViewerMode
+from ._native_mujoco_viewer import NativeMujocoViewer
+from ._viewer_type import ViewerType
+
+__all__ = [
+    "CustomMujocoViewer",
+    "CustomMujocoViewerMode",
+    "NativeMujocoViewer",
+    "ViewerType",
+]