diff --git a/python/mujoco/experimental/studio/native_viewer.py b/python/mujoco/experimental/studio/native_viewer.py index ab884015..cd50d8cf 100644 --- a/python/mujoco/experimental/studio/native_viewer.py +++ b/python/mujoco/experimental/studio/native_viewer.py @@ -23,25 +23,21 @@ to use these classes. import mujoco from mujoco.experimental.studio import native_viewer_cc as _viewer from mujoco.experimental.studio import ux -from mujoco.experimental.studio import viewer_protocol +from mujoco.experimental.studio import viewer_protocol as vp from mujoco.experimental.dear_imgui import dear_imgui as imgui -class NativeViewer(viewer_protocol.Viewer): +class NativeViewer(vp.Viewer): """Simulation-agnostic native viewer for MuJoCo models.""" def __init__( self, - model: mujoco.MjModel, + config: vp.ViewerConfig, *, camera: mujoco.MjvCamera | None = None, vis_options: mujoco.MjvOption | None = None, perturb: mujoco.MjvPerturb | None = None, render_flags: ux.RenderFlags | None = None, - title: str = '', - width: int = 1200, - height: int = 800, - gfx: str | None = None, ) -> None: """Initializes the NativeViewer. @@ -49,25 +45,23 @@ class NativeViewer(viewer_protocol.Viewer): visualization option objects unless they are provided. Args: - model: The MuJoCo model, used to initialize the renderer. + config: Viewer window configuration. camera: Camera parameters. Internal object is created if None. vis_options: Visualization options. Internal object is created if None. perturb: Perturbation parameters. Internal object is created if None. render_flags: Render flags. Internal object is created if None. - title: Title of the viewer window. - width: Initial width of the viewer window. - height: Initial height of the viewer window. - gfx: Graphics mode. If None, uses the platform default. """ + self.config = config self.camera = camera or mujoco.MjvCamera() self.perturb = perturb or mujoco.MjvPerturb() self.vis_options = vis_options or mujoco.MjvOption() - self._viewer = _viewer.Viewer(title, width, height, gfx or '') - self._viewer.InitRenderer(model) + self._viewer = _viewer.Viewer( + config.title, config.width, config.height, config.gfx or '' + ) # This class does not own the model but we need to know if the model being # rendered has changed, so we store the unique python object id here so we # can use it to detect model changes. - self._renderer_model_id = id(model) + self._renderer_model_id = id(None) self._is_running = True if render_flags is not None: self.render_flags = render_flags @@ -114,9 +108,14 @@ class NativeViewer(viewer_protocol.Viewer): self.render_flags.flags, ) + def close(self) -> None: + """Close the viewer.""" + self._is_running = False + + # TODO(matijak): Remove stop() and rename callers to close(). def stop(self) -> None: """Stop the viewer.""" - self._is_running = False + self.close() def get_drop_file(self) -> str: """Returns the path of the file dropped into the window, or empty string.""" diff --git a/python/mujoco/experimental/studio/studio.py b/python/mujoco/experimental/studio/studio.py index 92953a50..8d4b4754 100644 --- a/python/mujoco/experimental/studio/studio.py +++ b/python/mujoco/experimental/studio/studio.py @@ -30,12 +30,12 @@ def main(argv: list[str]) -> None: app = studio_app.StudioApp.from_argv(argv) # Initialize the viewer. - viewer = native_viewer.NativeViewer( - app.model, + config = viewer_protocol.ViewerConfig( width=_WIDTH.value, height=_HEIGHT.value, gfx=_GFX.value, ) + viewer = native_viewer.NativeViewer(config) # Main viewer loop. while viewer.is_running(): diff --git a/python/mujoco/experimental/studio/studio_app_events.py b/python/mujoco/experimental/studio/studio_app_events.py index 45397c0e..f3265f89 100644 --- a/python/mujoco/experimental/studio/studio_app_events.py +++ b/python/mujoco/experimental/studio/studio_app_events.py @@ -175,8 +175,9 @@ def handle_step_control_keyboard_events( step_control.set_pause_state(sim.PauseState.UNPAUSED) return True elif pressed(imgui.Key.Backspace): - mujoco.mj_resetData(model, data) - mujoco.mj_forward(model, data) + if model is not None and data is not None: + mujoco.mj_resetData(model, data) + mujoco.mj_forward(model, data) return True elif pressed(imgui.Key.Minus): ux.set_speed_index(step_control, ux_state, ux_state.speed_index + 1) diff --git a/python/mujoco/experimental/studio/viewer_protocol.py b/python/mujoco/experimental/studio/viewer_protocol.py index d2de18d8..fa1e7138 100644 --- a/python/mujoco/experimental/studio/viewer_protocol.py +++ b/python/mujoco/experimental/studio/viewer_protocol.py @@ -13,12 +13,12 @@ # limitations under the License. """Structural protocol defining the common viewer interface. -StudioApp uses the protocol for convenience methods that accept any viewer. +ViewerApp uses the protocol for convenience methods that accept any viewer. """ import dataclasses -from typing import Any -from typing import Protocol +import enum +from typing import Any, Protocol import mujoco from mujoco.experimental.studio import ux import numpy as np @@ -35,10 +35,36 @@ GFX_MODES = ( ) +class ViewerMode(enum.StrEnum): + """Determines where the viewer is rendered.""" + + NATIVE = 'native' + WEB = 'web' + + +# ----------------------------------------------------------------------------- +# Viewer configuration. +# ----------------------------------------------------------------------------- + + +@dataclasses.dataclass +class ViewerConfig: + """Common configuration for creating a viewer window.""" + + title: str = '' + width: int = 1200 + height: int = 800 + gfx: str = '' + viewer_mode: ViewerMode = ViewerMode.NATIVE + + +# Legacy message types kept for backward compatibility. +# Will be removed when callers are migrated. + + @dataclasses.dataclass class SimToView: """A message sent from the simulation to the viewer.""" - model: mujoco.MjModel | None = None state: np.ndarray | None = None state_sig: int = 0 @@ -48,7 +74,6 @@ class SimToView: @dataclasses.dataclass class ViewToSim: """A message sent from the viewer to the simulation.""" - state: np.ndarray | None = None state_sig: int = 0 reset: bool = False @@ -56,6 +81,11 @@ class ViewToSim: user_data: dict[str, Any] = dataclasses.field(default_factory=dict) +# ----------------------------------------------------------------------------- +# Structural interface for any viewer. +# ----------------------------------------------------------------------------- + + class Viewer(Protocol): """Structural interface for any viewer.""" @@ -70,8 +100,13 @@ class Viewer(Protocol): def sync(self, model: mujoco.MjModel, data: mujoco.MjData) -> None: ... - def stop(self) -> None: + def close(self) -> None: ... def get_drop_file(self) -> str: ... + + def upload_image( + self, tex_id: int, img: str | bytes, width: int, height: int, bpp: int + ) -> int: + ...