"""Scene utilities for synchronous Python MuJoCo simulations.
This module provides the `mujoco_scene` helper used by RobotBlockSet synchronous
MuJoCo backends to load models, manage the viewer, render auxiliary camera
windows, and control simulation stepping and resets.
Copyright (c) 2025 Jozef Stefan Institute
Authors: Leon Zlajpah.
"""
import numpy as np
import threading
from typing import Optional, Sequence, Union
try:
import mujoco
import mujoco.viewer
except Exception as e:
raise ModuleNotFoundError(f"{e}\nMuJoCo is not installed.\nInstall it with:\n pip install mujoco") from None
try:
import glfw
except Exception as e:
raise ModuleNotFoundError(f"{e}\nGLFW is not installed.\nInstall it with:\n pip install glfw") from None
from robotblockset.tools import rbs_object
[docs]
class mujoco_scene(rbs_object):
"""
MuJoCo scene manager with viewer support for synchronous backends.
Attributes
----------
spec : Optional[mujoco.MjSpec]
MuJoCo specification object when the scene was loaded from XML.
model : mujoco.MjModel
Compiled MuJoCo model used by the scene.
data : mujoco.MjData
Runtime MuJoCo data associated with `model`.
pause : bool
Flag indicating whether simulation stepping is paused.
viewer : Optional[mujoco.viewer.Handle]
Passive MuJoCo viewer used for the main scene window.
show_camera : list[Union[str, int]]
Camera names or IDs rendered in additional windows.
cam_windows : list[dict[str, object]]
Metadata for additional GLFW camera windows.
"""
[docs]
def __init__(self, model_xml_file: Optional[str] = None, model: Optional[mujoco.MjModel] = None, show_viewer: bool = True, show_camera: Optional[Sequence[Union[str, int]]] = None, verbose: int = 0) -> None:
"""Create a synchronous MuJoCo scene manager.
Parameters
----------
model_xml_file : str, optional
Path to the MuJoCo XML model file to load.
model : mujoco.MjModel, optional
Existing MuJoCo model to use instead of loading from XML.
show_viewer : bool, optional
If `True`, open the main passive MuJoCo viewer.
show_camera : sequence of str or int, optional
Camera names or IDs to render in auxiliary windows.
verbose : int, optional
Verbosity level used for status messages.
Returns
-------
None
This constructor initializes the synchronous MuJoCo scene object in place.
"""
rbs_object.__init__(self)
self._verbose = verbose
self.Name = "pyMuJoCo_sim"
if model is not None:
self.spec = None
self.model = model
elif model_xml_file is not None:
self.spec = mujoco.MjSpec.from_file(model_xml_file)
self.model = self.spec.compile()
else:
raise ValueError("Either model_xml_file or model must be provided.")
self.data = mujoco.MjData(self.model)
mujoco.mj_forward(self.model, self.data)
self.pause = False
self.DebugMessage("Model loaded successfully.")
self._connected = False
self.viewer = None
self._viewer_active = show_viewer
self.cam_windows = []
self.show_camera = list(show_camera) if show_camera is not None else []
self.scene = None
self.opt = None
self._camera_thread_id: Optional[int] = None
self._glfw_initialized_by_scene = False
self.start_simulation()
[docs]
def start_simulation(self) -> None:
"""
Start the viewer and auxiliary camera windows.
Returns
-------
None
This method creates the passive viewer and any requested camera windows.
Raises
------
RuntimeError
If viewer or camera-window initialization fails, or auxiliary
windows are requested from a non-main thread.
ValueError
If a camera identifier is invalid.
"""
viewer_ready = not self._viewer_active or (self.viewer is not None and self.viewer.is_running())
cameras_ready = not self.show_camera or bool(self.cam_windows)
if self._connected and viewer_ready and cameras_ready:
return
if self.viewer is not None or self.cam_windows:
self.stop_simulation()
try:
if self._viewer_active:
self.viewer = mujoco.viewer.launch_passive(self.model, self.data)
if self.show_camera:
self._require_camera_main_thread()
self._camera_thread_id = threading.get_ident()
if not self._viewer_active:
if not glfw.init():
raise RuntimeError("Failed to initialize GLFW for auxiliary camera windows")
self._glfw_initialized_by_scene = True
default_w, default_h = 400, 300
for camera_id in self.show_camera:
cam_id = self._resolve_camera_id(camera_id)
win = glfw.create_window(default_w, default_h, f"Camera {cam_id}", None, None)
if not win:
raise RuntimeError(f"Failed to create window for camera {cam_id}")
glfw.make_context_current(win)
ctx = mujoco.MjrContext(self.model, mujoco.mjtFontScale.mjFONTSCALE_100)
cam = mujoco.MjvCamera()
cam.type = mujoco.mjtCamera.mjCAMERA_FIXED
cam.fixedcamid = cam_id
self.cam_windows.append({"id": cam_id, "win": win, "ctx": ctx, "cam": cam})
self.scene = mujoco.MjvScene(self.model, maxgeom=20_000)
self.opt = mujoco.MjvOption()
self.opt.frame = mujoco.mjtFrame.mjFRAME_NONE
self.opt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT.value] = 0
self.opt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTFORCE.value] = 1
except BaseException:
self._close_camera_windows()
if self.viewer is not None:
self.viewer.close()
self.viewer = None
self._connected = False
raise
self._connected = True
self.DebugMessage("Viewer started.")
[docs]
def stop_simulation(self) -> None:
"""
Stop rendering and close all viewer resources.
Returns
-------
None
This method closes the main viewer and all auxiliary windows.
"""
self._close_camera_windows()
if self.viewer is not None:
self.viewer.close()
self.viewer = None
self._connected = False
self.DebugMessage("Viewer stopped.")
[docs]
def restart_simulation(self) -> None:
"""
Restart rendering and recompile the model when possible.
Returns
-------
None
This method refreshes the scene after stopping the viewer.
"""
self.stop_simulation()
if self.spec is not None:
self.model, self.data = self.spec.recompile(self.model, self.data)
mujoco.mj_forward(self.model, self.data)
self.start_simulation()
def _require_camera_main_thread(self) -> None:
"""Require auxiliary GLFW operations to run on their owning main thread."""
if threading.current_thread() is not threading.main_thread():
raise RuntimeError("Auxiliary GLFW camera windows must be managed from the main thread")
if self._camera_thread_id is not None and threading.get_ident() != self._camera_thread_id:
raise RuntimeError("Auxiliary camera windows must be managed by the thread that created them")
def _resolve_camera_id(self, camera: Union[str, int]) -> int:
"""Resolve and validate a MuJoCo camera name or ID."""
if isinstance(camera, str):
try:
camera_id = int(self.model.camera(camera).id)
except Exception as exc:
raise ValueError(f"Unknown MuJoCo camera '{camera}'") from exc
elif isinstance(camera, (int, np.integer)) and not isinstance(camera, (bool, np.bool_)):
camera_id = int(camera)
else:
raise TypeError("Camera identifiers must be names or integer IDs")
if camera_id < 0 or camera_id >= self.model.ncam:
raise ValueError(f"Camera ID must be in the range [0, {self.model.ncam})")
return camera_id
def _close_camera_window(self, camera_window: dict[str, object]) -> None:
"""Release one auxiliary camera context and destroy its GLFW window."""
win = camera_window["win"]
glfw.make_context_current(win)
context = camera_window["ctx"]
if hasattr(context, "free"):
context.free()
glfw.destroy_window(win)
def _close_camera_windows(self) -> None:
"""Close all auxiliary camera windows and release their contexts."""
if self.cam_windows:
self._require_camera_main_thread()
for camera_window in self.cam_windows:
self._close_camera_window(camera_window)
self.cam_windows = []
self.scene = None
self.opt = None
self._camera_thread_id = None
if self._glfw_initialized_by_scene:
glfw.terminate()
self._glfw_initialized_by_scene = False
[docs]
def update_cam_windows(self) -> None:
"""Render all active auxiliary camera windows.
Returns
-------
None
This method updates and prunes additional GLFW camera windows.
"""
if not self.cam_windows:
return
self._require_camera_main_thread()
alive = []
for cw in self.cam_windows:
win = cw["win"]
if glfw.window_should_close(win):
self._close_camera_window(cw)
self.DebugMessage(f"Camera {cw['id']} view closed.")
continue
alive.append(cw)
# Get framebuffer viewport
glfw.make_context_current(win)
fb_w, fb_h = glfw.get_framebuffer_size(win)
viewport = mujoco.MjrRect(0, 0, fb_w, fb_h)
# Update scene and render
mujoco.mjv_updateScene(self.model, self.data, self.opt, None, cw["cam"], mujoco.mjtCatBit.mjCAT_ALL.value, self.scene)
mujoco.mjr_render(viewport, self.scene, cw["ctx"])
# Swap OpenGL buffers (blocking vsync)
glfw.swap_buffers(win)
self.cam_windows = alive
glfw.poll_events()
[docs]
def mj_step(self) -> None:
"""
Advances the simulation by one time step.
Returns
-------
None
This method steps the MuJoCo simulation and refreshes viewers.
"""
if not self.pause:
mujoco.mj_step(self.model, self.data)
if self._viewer_active and self.viewer is not None and self.viewer.is_running():
self.viewer.sync()
self.update_cam_windows()
[docs]
def mj_forward(self) -> None:
"""
Advances the simulation by using forward dynamics.
Returns
-------
None
This method recomputes forward dynamics and refreshes viewers.
"""
mujoco.mj_forward(self.model, self.data)
if self._viewer_active and self.viewer is not None and self.viewer.is_running():
self.viewer.sync()
self.update_cam_windows()
[docs]
def mj_pause(self) -> None:
"""Pause simulation stepping.
Returns
-------
None
This method sets the internal pause flag.
"""
self.pause = True
[docs]
def mj_run(self) -> None:
"""Resume simulation stepping.
Returns
-------
None
This method clears the internal pause flag.
"""
self.pause = False
[docs]
def mj_wait(self, wait: float = 0) -> None:
"""
Advance the MuJoCo simulation for a specified duration.
This method performs repeated simulation steps until the internal
simulation time (`self.data.time`) has advanced by the given
amount. Unlike a passive delay (e.g., ``time.sleep``), this
function actively progresses the MuJoCo physics simulation.
Parameters
----------
wait : float, optional
Amount of simulated time (in seconds) to advance.
A value of ``0`` (default) performs no additional steps.
Notes
-----
- Simulation progresses by repeatedly calling ``self.mj_step``.
- The function returns only after the MuJoCo model time exceeds
the initial time plus ``wait``.
- This does *not* block real time; the speed of advancement
depends on computation speed and simulation complexity.
Returns
-------
None
This method blocks until the requested amount of simulated time has elapsed.
"""
if isinstance(wait, (bool, np.bool_)) or not np.isscalar(wait):
raise TypeError("wait must be a real scalar")
try:
wait = float(wait)
except (TypeError, ValueError) as exc:
raise TypeError("wait must be a real scalar") from exc
if not np.isfinite(wait) or wait < 0:
raise ValueError("wait must be finite and nonnegative")
if wait == 0:
return
if self.pause:
raise RuntimeError("Cannot advance simulation time while the scene is paused")
start_time = float(self.data.time)
target_time = start_time + wait
while self.data.time < target_time:
previous_time = float(self.data.time)
self.mj_step()
current_time = float(self.data.time)
if not np.isfinite(current_time):
raise RuntimeError("MuJoCo returned invalid simulation time")
if current_time <= previous_time:
raise RuntimeError("MuJoCo simulation time did not advance")
[docs]
def mj_reset(self, keyframe: Optional[int] = None) -> None:
"""
Resets the simulation state to its initial state, or to a specified keyframe.
Parameters
----------
keyframe : int, optional
The index of the keyframe to reset to. If None, the simulation is reset to the initial state.
Returns
-------
None
"""
if keyframe is None:
mujoco.mj_resetData(self.model, self.data)
else:
mujoco.mj_resetDataKeyframe(self.model, self.data, keyframe)
mujoco.mj_forward(self.model, self.data)
[docs]
def mj_capture_camera(self, camera: Union[int, str] = -1, scene_option: Optional[mujoco.MjvOption] = None, **kwargs: object) -> np.ndarray:
"""
Captures an image from a specified camera in the simulation.
Parameters
----------
camera : Union[int, str], optional
The ID or name of the camera to capture. Default is -1, which captures the default camera.
scene_option : mujoco.MjvOption, optional
Rendering options to customize the scene appearance. Default is None.
**kwargs : object
Additional keyword arguments passed to `mujoco.Renderer`.
Returns
-------
np.ndarray
The rendered image as a NumPy array (RGB format).
"""
with mujoco.Renderer(self.model, **kwargs) as renderer:
renderer.update_scene(self.data, camera=camera, scene_option=scene_option)
frame = renderer.render()
return frame