Skip to content

Commit

Permalink
added native mujoco viewer
Browse files Browse the repository at this point in the history
  • Loading branch information
oliverweissl committed Apr 30, 2024
1 parent 51b2d28 commit dd9050b
Show file tree
Hide file tree
Showing 11 changed files with 301 additions and 40 deletions.
5 changes: 4 additions & 1 deletion examples/1_simulator_basics/1a_simulate_single_robot/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down
3 changes: 2 additions & 1 deletion simulation/revolve2/simulation/simulator/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
53 changes: 53 additions & 0 deletions simulation/revolve2/simulation/simulator/_viewer.py
Original file line number Diff line number Diff line change
@@ -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(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.
"""
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -18,6 +19,7 @@ class LocalSimulator(Simulator):
_cast_shadows: bool
_fast_sim: bool
_manual_control: bool
_viewer_type: ViewerType

def __init__(
self,
Expand All @@ -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.
Expand All @@ -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
Expand All @@ -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]]:
"""
Expand Down Expand Up @@ -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)
]
Expand All @@ -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)
]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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]:
"""
Expand All @@ -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")
Expand All @@ -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)
Expand All @@ -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
):
Expand Down Expand Up @@ -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()
Expand All @@ -164,20 +187,21 @@ 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()

Expand All @@ -193,5 +217,4 @@ def simulate_scene(
)

logging.info(f"Scene {scene_id} done.")

return simulation_states
Original file line number Diff line number Diff line change
@@ -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",
]
Loading

0 comments on commit dd9050b

Please sign in to comment.