Refactor Viewer to be a base class owning the communication with the simulation, handler registry, core visualization objects and the render function

ViewerApp is simplified, interacting with the endpoint and handlers through the Viewer instance. This change makes it possible to write simulation viewers without using the UI/UX provided by ViewerApp (which should be renamed StudioApp)

PiperOrigin-RevId: 943642660
Change-Id: Id8eb3d4e6a27ceda93833367fec4699fa6bc838e

python/mujoco/experimental/studio/viewer_protocol.py

@@ -11,15 +11,17 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Structural protocol defining the common viewer interface.
-
-ViewerApp uses the protocol for convenience methods that accept any viewer.
-"""
+"""Base class and configuration for any viewer."""
 
+import abc
+import copy
 import dataclasses
 import enum
-from typing import Any, Protocol
+from typing import Any
 import mujoco
+from mujoco.experimental.studio import endpoints
+from mujoco.experimental.studio import handler_registry
+from mujoco.experimental.studio import messages
 from mujoco.experimental.studio import ux
 import numpy as np
 
@@ -65,6 +67,7 @@ class ViewerConfig:
 @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
@@ -74,6 +77,7 @@ 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
@@ -82,32 +86,183 @@ class ViewToSim:
 
 
 # -----------------------------------------------------------------------------
-# Structural interface for any viewer.
+# Base class for any viewer.
 # -----------------------------------------------------------------------------
 
 
-class Viewer(Protocol):
-  """Structural interface for any viewer."""
+@dataclasses.dataclass(frozen=True)
+class ViewerInitEvent(messages.Event):
+  """Lifecycle event dispatched once when the concrete Viewer is initialized.
 
-  camera: mujoco.MjvCamera
-  perturb: mujoco.MjvPerturb
-  vis_options: mujoco.MjvOption
-  render_flags: ux.RenderFlags
-  extra_geoms: list[mujoco.MjvGeom]
+  Handlers that need access to the Viewer should handle this event
+  and cache the reference.
+  """
 
-  def is_running(self) -> bool:
-    ...
+  viewer: 'Viewer'
 
-  def sync(self, model: mujoco.MjModel, data: mujoco.MjData) -> None:
-    ...
+
+class Viewer(abc.ABC):
+  """Base class for any viewer.
+
+  Owns the communication endpoint, handler registry and core visualization
+  objects. The application is rendered by calling ``sync()``.
+  """
+
+  def __init__(
+      self,
+      config: ViewerConfig,
+      endpoint: endpoints.ViewerEndpoint,
+      *,
+      model: mujoco.MjModel | None = None,
+      model_path: str = '',
+      handlers: list[Any] | None = None,
+      camera: mujoco.MjvCamera | None = None,
+      vis_options: mujoco.MjvOption | None = None,
+      perturb: mujoco.MjvPerturb | None = None,
+      render_flags: ux.RenderFlags | None = None,
+      extra_geoms: list[mujoco.MjvGeom] | None = None,
+  ) -> None:
+    """Initializes the Viewer.
+
+    Args:
+      config: Viewer window configuration.
+      endpoint: The viewer endpoint for communication with the sim side.
+      model: Optional initial MjModel. If None, an empty model is created from
+        an empty MjSpec. The Viewer deep-copies this model and creates its own
+        MjData.
+      model_path: Optional path to the model file.
+      handlers: Optional list of handler instances for viewer-side processing.
+      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.
+      extra_geoms: List of extra geoms. Internal list is created if None.
+    """
+    self.config = config
+    self._endpoint = endpoint
+    self._is_running = True
+
+    # Viewer-owned model and data.
+    if model is None:
+      model = mujoco.MjSpec().compile()
+    self.model: mujoco.MjModel
+    self.data: mujoco.MjData
+    self.model_path: str = ''
+    self.load_model(model, model_path)
+
+    # Visual state.
+    self.camera = camera or mujoco.MjvCamera()
+    self.cam_speed = 0.001
+    self.perturb = perturb or mujoco.MjvPerturb()
+    self.vis_options = vis_options or mujoco.MjvOption()
+    self.extra_geoms = extra_geoms or []
+    if render_flags is not None:
+      self.render_flags = render_flags
+    else:
+      self.render_flags = ux.RenderFlags()
+      # Initted to match mujoco/src/engine/engine_vis_init.c
+      self.render_flags.flags = [1, 0, 1, 0, 1, 0, 1, 0, 0, 0, 1]
+
+    # Handler infrastructure.
+    all_handlers = [self] + list(handlers or [])
+    self.handlers = handler_registry.HandlerRegistry(all_handlers)
 
   def close(self) -> None:
+    """Closes the viewer, sends an exit event and shuts down the endpoint."""
+    if self._is_running:
+      self._is_running = False
+      try:
+        self.send_to_sim(messages.ExitEvent())
+      except Exception:  # pylint: disable=broad-exception-caught
+        pass  # Ignore exceptions, the sim may have already closed.
+      self._endpoint.close()
+
+  def is_running(self) -> bool:
+    """Returns True while the viewer has not been closed."""
+    return self._is_running
+
+  def send_to_sim(self, message: messages.Message) -> None:
+    """Sends a message to the simulation process."""
+    self._endpoint.send_to_sim(message)
+
+  def get_sim_events(self) -> list[messages.Event]:
+    """Returns all pending events from the simulation."""
+    return self._endpoint.get_sim_events()
+
+  def get_sim_snapshots(self) -> list[messages.Snapshot]:
+    """Returns all pending latest snapshots from the simulation, one per type."""
+    return self._endpoint.get_sim_snapshots()
+
+  def dispatch(self, message: messages.Message) -> None:
+    """Dispatches a message to registered handlers in priority order."""
+    self.handlers.dispatch(message)
+
+  def load_model(self, model: mujoco.MjModel, model_path: str = '') -> None:
+    """Deep-copies a model and creates fresh data for the viewer."""
+    self.model_path = model_path
+    self.model = copy.deepcopy(model)
+    self.data = mujoco.MjData(self.model)
+    assert id(self.model) != id(model)
+    mujoco.mj_forward(self.model, self.data)
+
+  @messages.handler(priority=messages.Priority.CRITICAL)
+  def _on_model(self, event: messages.ModelEvent) -> bool:
+    """Deep-copies the incoming model so the Viewer owns its data."""
+    self.load_model(event.model, event.path)
+    self.extra_geoms.clear()
+    return False  # Do not consume; let other handlers see the event.
+
+  @messages.handler(priority=messages.Priority.CRITICAL)
+  def _on_state(self, event: messages.StateSnapshot) -> bool:
+    """Applies incoming simulation state to the viewer's model/data."""
+    state_size = mujoco.mj_stateSize(self.model, event.state_sig)
+    if len(event.state) == state_size:
+      mujoco.mj_setState(self.model, self.data, event.state, event.state_sig)
+      mujoco.mj_forward(self.model, self.data)
+    return False  # Do not consume; let other handlers see the event.
+
+  @abc.abstractmethod
+  def sync(self) -> None:
+    """Renders the scene using the viewer's current model and data."""
     ...
 
+  @abc.abstractmethod
   def get_drop_file(self) -> str:
     ...
 
+  @abc.abstractmethod
   def upload_image(
       self, tex_id: int, img: str | bytes, width: int, height: int, bpp: int
   ) -> int:
     ...
+
+
+# -----------------------------------------------------------------------------
+# Standalone viewer loop.
+# -----------------------------------------------------------------------------
+
+
+def run_viewer_loop(viewer: Viewer) -> None:
+  """Minimal viewer loop: process sim messages, dispatch lifecycle events, sync.
+
+  Runs until the viewer window is closed or an exit event is received.
+  On exit, closes the viewer (which sends an ExitEvent to the sim side).
+
+  Args:
+    viewer: A Viewer that owns the endpoint and handler registry.
+  """
+  while viewer.is_running():
+    # Process incoming messages.
+    for event in viewer.get_sim_events():
+      viewer.dispatch(event)
+    for snapshot in viewer.get_sim_snapshots():
+      viewer.dispatch(snapshot)
+
+    # Dispatch lifecycle events.
+    viewer.dispatch(messages.UpdateEvent())
+    viewer.dispatch(messages.BuildGuiEvent())
+
+    # Render the scene.
+    viewer.sync()
+
+  viewer.close()