"""Scene utilities for interactive Python MuJoCo simulations.
This module provides the `mujoco_scene` helper used by RobotBlockSet MuJoCo
backends to load models, manage the passive viewer, advance simulation in a
worker thread, capture camera images, and reset simulation state.
Copyright (c) 2025 Jozef Stefan Institute
Authors: Leon Zlajpah.
"""
from time import perf_counter, sleep
import threading
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
import numpy as np
from typing import ContextManager, Optional, Sequence, Union
from robotblockset.tools import rbs_object
[docs]
class mujoco_scene(rbs_object):
"""
MuJoCo scene manager with a passive viewer and background physics loop.
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.
synchronized : bool
Whether simulation stepping is synchronized to the model timestep.
viewer : Optional[mujoco.viewer.Handle]
Passive MuJoCo viewer used for the main scene window.
show_camera : list[Union[str, int]]
Requested auxiliary cameras. Auxiliary camera windows are not supported
by this threaded backend.
visual_thread : Optional[threading.Thread]
Background thread that advances the simulation.
"""
[docs]
def __init__(self, model_xml_file: Optional[str] = None, model: Optional[mujoco.MjModel] = None, show_camera: Optional[Sequence[Union[str, int]]] = None, synchronized: bool = True, verbose: int = 0) -> None:
"""Create a 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_camera : sequence of str or int, optional
Auxiliary camera names or IDs. Additional GLFW windows are not
supported by this threaded backend; use ``mj_capture_camera()`` or
:mod:`robotblockset.mujoco.scene_pymujoco_sim` instead.
synchronized : bool, optional
If `True`, synchronize stepping to the MuJoCo model timestep.
verbose : int, optional
Verbosity level used for status messages.
Returns
-------
None
This constructor initializes the MuJoCo scene object in place.
"""
rbs_object.__init__(self)
self._verbose = verbose
self.Name = "pyMuJoCo"
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.synchronized = synchronized
self._data_lock = threading.RLock()
self._viewer_ready = threading.Event()
self._viewer_error: Optional[BaseException] = None
self.abort_viewer = False
self._connected = False
self.DebugMessage("Model loaded successfully.")
self.show_camera = list(show_camera) if show_camera is not None else []
if self.show_camera:
raise ValueError(
"Auxiliary GLFW camera windows are not supported by the threaded "
"scene. Use mj_capture_camera() or scene_pymujoco_sim instead."
)
self.viewer = None
self.visual_thread: Optional[threading.Thread] = None
self._last_step_time = perf_counter() - self.model.opt.timestep
self.start_simulation()
def _synchro_simulation(self) -> None:
"""
Synchronize simulation stepping to the configured model timestep.
Returns
-------
None
This method delays the caller so the next step matches the target timestep.
"""
next_time = self._last_step_time + self.model.opt.timestep
self._remaining = next_time - perf_counter()
if self._remaining > 0:
sleep(self._remaining)
self._last_step_time = perf_counter()
[docs]
def start_simulation(self) -> None:
"""
Start the passive viewer and simulation thread.
Returns
-------
None
This method launches the viewer on the caller thread and starts the
background physics loop.
Raises
------
RuntimeError
If the viewer or simulation thread cannot be started.
"""
if self.visual_thread is not None and self.visual_thread.is_alive():
return
self.abort_viewer = False
self._connected = False
self._viewer_error = None
self._viewer_ready.clear()
self._last_step_time = perf_counter() - self.model.opt.timestep
def key_callback(keycode: int) -> None:
if keycode == ord("Q"):
self.abort_viewer = True
elif keycode == ord(" "):
self.pause = not self.pause
try:
# MuJoCo performs platform-specific GLFW initialization here, so
# launch the passive viewer on the caller thread.
self.viewer = mujoco.viewer.launch_passive(self.model, self.data, key_callback=key_callback)
except BaseException as exc:
self.viewer = None
raise RuntimeError("Failed to start the MuJoCo viewer") from exc
self.visual_thread = threading.Thread(target=self.render, daemon=True)
self.visual_thread.start()
if not self._viewer_ready.wait(timeout=10.0):
self.abort_viewer = True
self.viewer.close()
self.visual_thread.join(timeout=1.0)
raise RuntimeError("Timed out while starting the MuJoCo simulation thread")
if self._viewer_error is not None:
error = self._viewer_error
self.viewer.close()
self.visual_thread.join(timeout=1.0)
raise RuntimeError("Failed to start the MuJoCo simulation thread") from error
if not self.visual_thread.is_alive() or not self.viewer.is_running():
self.viewer.close()
self.visual_thread.join(timeout=1.0)
raise RuntimeError("MuJoCo viewer closed during startup")
self._connected = True
self.DebugMessage("Viewer started.")
[docs]
def stop_simulation(self) -> None:
"""
Stop the visualization thread and close viewer windows.
Returns
-------
None
This method stops the background render loop.
"""
self.abort_viewer = True
if self.visual_thread is not None and self.visual_thread.is_alive():
self.visual_thread.join(timeout=5.0)
if self.visual_thread.is_alive():
if self.viewer is not None:
self.viewer.close()
self.visual_thread.join(timeout=1.0)
if self.visual_thread.is_alive():
raise RuntimeError("MuJoCo simulation thread did not stop")
self.DebugMessage("Viewer stopped.")
if self.viewer is not None:
self.viewer.close()
self._connected = False
[docs]
def restart_simulation(self) -> None:
"""
Restart the viewer thread and recompile the model when possible.
Returns
-------
None
This method restarts visualization using the current scene state.
"""
self.stop_simulation()
if self.spec is not None:
with self._data_lock:
self.model, self.data = self.spec.recompile(self.model, self.data)
mujoco.mj_forward(self.model, self.data)
self.start_simulation()
[docs]
def render(self) -> None:
"""
Advance the simulation and synchronize the passive viewer.
Viewer initialization is performed by :meth:`start_simulation` on the
caller thread. This worker owns physics stepping but never calls GLFW
directly.
Returns
-------
None
"""
try:
with self._data_lock:
self.viewer.sync()
self._viewer_ready.set()
self.DebugMessage("Viewer launched, running simulation...")
while self.viewer.is_running() and not self.abort_viewer:
if self.synchronized:
self._synchro_simulation()
if self.abort_viewer:
break
with self._data_lock:
if not self.pause:
mujoco.mj_step(self.model, self.data)
self.viewer.sync()
if self.pause and not self.synchronized:
sleep(self.model.opt.timestep)
except BaseException as exc:
self._viewer_error = exc
self._viewer_ready.set()
self.ErrorMessage(f"MuJoCo simulation thread stopped with an error: {exc}")
finally:
self._connected = False
self._viewer_ready.set()
self.DebugMessage("Viewer closed.")
[docs]
def lock(self) -> ContextManager[object]:
"""Return the lock protecting the shared MuJoCo model and data.
Code that directly reads or modifies ``model`` or ``data`` while the
simulation thread is running should hold this lock.
Returns
-------
ContextManager[object]
Reentrant scene-data lock.
"""
return self._data_lock
[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:
"""
Wait for a specified duration.
Parameters
----------
wait : float, optional
Amount of simulated time (in seconds) to advance.
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 wait for simulation time while the scene is paused")
with self._data_lock:
start_time = float(self.data.time)
target_time = start_time + wait
last_time = start_time
last_progress = perf_counter()
stall_timeout = max(1.0, 2.0 * float(self.model.opt.timestep))
while True:
if self.pause:
raise RuntimeError("Simulation was paused while waiting")
if self.visual_thread is None or not self.visual_thread.is_alive():
error = RuntimeError("Simulation time cannot advance because the simulation thread is not running")
if self._viewer_error is not None:
raise error from self._viewer_error
raise error
with self._data_lock:
current_time = float(self.data.time)
if not np.isfinite(current_time):
raise RuntimeError("MuJoCo returned invalid simulation time")
if current_time >= target_time:
return
if current_time > last_time:
last_time = current_time
last_progress = perf_counter()
elif perf_counter() - last_progress >= stall_timeout:
raise RuntimeError("MuJoCo simulation time did not advance")
sleep(min(max(float(self.model.opt.timestep), 1e-4), 0.01))
[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
"""
with self._data_lock:
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 self._data_lock:
with mujoco.Renderer(self.model, **kwargs) as renderer:
renderer.update_scene(self.data, camera=camera, scene_option=scene_option)
frame = renderer.render()
return frame