Refactor viewer initialization and protocol.

Introduced a ViewerConfig dataclass to encapsulate common viewer window parameters. The NativeViewer now accepts a ViewerConfig and no longer requires a model at initialization. The Viewer protocol has been updated to include a close method (renamed from stop) and an upload_image method. Added checks for None models/data in studio event handling.

PiperOrigin-RevId: 939631378
Change-Id: Ibc07d0e775efa0d79ce5552ee52102e943a9b290

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."""

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():

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)

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:
+    ...