Add MuJoCo Studio threaded/passive launch function

PiperOrigin-RevId: 941043766
Change-Id: I5b75d0a2b183d402e8ce221883c2734255e87f0f

python/mujoco/experimental/studio/launch_passive.py

@@ -0,0 +1,163 @@
+# Copyright 2026 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+"""Non-blocking Studio launcher.
+
+Launches the full Studio GUI in a daemon thread and returns a ``ViewerHandle``
+that the calling thread uses to push simulation state into the viewer.
+"""
+
+import queue
+import threading
+
+from mujoco.experimental.studio import endpoints
+from mujoco.experimental.studio import messages
+from mujoco.experimental.studio import sim_app
+from mujoco.experimental.studio import viewer_app
+from mujoco.experimental.studio import viewer_protocol
+
+sa = sim_app
+va = viewer_app
+vp = viewer_protocol
+
+
+class _PassiveSnapshotChannel:
+  """SnapshotChannel for passive mode: per-type snapshot by reference, wakes the viewer thread on put."""
+
+  def __init__(self):
+    self._pending_snapshots: dict[
+        type[messages.Snapshot], messages.Snapshot
+    ] = {}
+    self._lock = threading.Lock()
+    self._wakeup = threading.Event()
+
+  def put(self, snapshot: messages.Snapshot) -> None:
+    with self._lock:
+      self._pending_snapshots[type(snapshot)] = snapshot
+    self._wakeup.set()
+
+  def get(self) -> list[messages.Snapshot]:
+    self._wakeup.clear()
+    with self._lock:
+      out = list(self._pending_snapshots.values())
+      self._pending_snapshots.clear()
+      return out
+
+  def close(self) -> None:
+    pass
+
+
+class _PassiveEventChannel(messages.EventChannel):
+  """EventChannel for passive mode: a thread-safe queue.Queue of event messages."""
+
+  def __init__(self):
+    self._events: queue.Queue[messages.Event] = queue.Queue()
+
+  def put(self, event: messages.Event) -> None:
+    self._events.put(event)
+
+  def get(self) -> list[messages.Event]:
+    out = []
+    while True:
+      try:
+        out.append(self._events.get_nowait())
+      except queue.Empty:
+        break
+    return out
+
+  def close(self) -> None:
+    pass
+
+
+def _run_viewer_target(
+    config: viewer_protocol.ViewerConfig,
+    viewer_endpoint: endpoints.ViewerEndpoint,
+    viewer_gui_hook: va.ViewerGuiHook | None = None,
+    viewer_update_hook: va.ViewerUpdateHook | None = None,
+    viewer_event_handler: va.ViewerEventHandler | None = None,
+    viewer_snapshot_handler: va.ViewerSnapshotHandler | None = None,
+) -> None:
+  """Creates the appropriate viewer and runs the viewer loop."""
+
+  if config.viewer_mode == viewer_protocol.ViewerMode.NATIVE:
+    from mujoco.experimental.studio import native_viewer  # pylint: disable=g-import-not-at-top
+
+    viewer = native_viewer.NativeViewer(config)
+  elif config.viewer_mode == viewer_protocol.ViewerMode.WEB:
+    raise NotImplementedError('Web viewer not implemented yet')
+  else:
+    raise ValueError(f'Unknown viewer mode: {config.viewer_mode!r}')
+
+  va.run_viewer(
+      viewer,
+      viewer_endpoint,
+      viewer_gui_hook=viewer_gui_hook,
+      viewer_update_hook=viewer_update_hook,
+      viewer_event_handler=viewer_event_handler,
+      viewer_snapshot_handler=viewer_snapshot_handler,
+  )
+
+
+def launch_passive(
+    config: vp.ViewerConfig,
+    *,
+    viewer_gui_hook: va.ViewerGuiHook | None = None,
+    viewer_update_hook: va.ViewerUpdateHook | None = None,
+    viewer_event_handler: va.ViewerEventHandler | None = None,
+    viewer_snapshot_handler: va.ViewerSnapshotHandler | None = None,
+    sim_event_handler: sa.SimEventHandler | None = None,
+) -> sa.ViewerHandle:
+  """Launches the Studio GUI in a daemon thread without blocking.
+
+  The viewer runs the full Studio UI (toolbar, options, inspector) on the
+  rendering thread.  The caller keeps running and pushes state via
+  ``handle.sync()``.
+
+  Args:
+    config: Viewer window configuration.
+    viewer_gui_hook: Optional hook to draw custom ImGui panels.
+    viewer_update_hook: Optional hook called once per frame before GUI.
+    viewer_event_handler: Optional handler for sim-to-viewer events.
+    viewer_snapshot_handler: Optional handler for sim-to-viewer snapshots.
+    sim_event_handler: Optional handler for viewer-to-sim events.
+
+  Returns:
+    A ViewerHandle for interacting with the viewer.
+  """
+  viewer_endpoint, sim_endpoint = endpoints.make_endpoints(
+      s2v_snapshot=_PassiveSnapshotChannel(),
+      s2v_events=_PassiveEventChannel(),
+      v2s_events=_PassiveEventChannel(),
+      v2s_snapshot=_PassiveSnapshotChannel(),
+  )
+
+  thread = threading.Thread(
+      target=_run_viewer_target,
+      args=(
+          config,
+          viewer_endpoint,
+          viewer_gui_hook,
+          viewer_update_hook,
+          viewer_event_handler,
+          viewer_snapshot_handler,
+      ),
+      daemon=True,
+  )
+  thread.start()
+
+  handle = sa.ViewerHandle(
+      sim_endpoint,
+      is_alive_fn=thread.is_alive,
+      sim_event_handler=sim_event_handler,
+  )
+  return handle

python/mujoco/experimental/studio/messages.py

@@ -58,7 +58,14 @@ class SnapshotChannel(Protocol):
   """Channel to transport Snapshot messages with latest-wins semantics."""
 
   def put(self, value: Snapshot) -> None:
-    """Overwrites the latest snapshot of the same type."""
+    """Overwrites the latest snapshot of the same type.
+
+    After calling put() it is the caller's responsibility to ensure it no longer
+    holds a reference to the snapshot.
+
+    Args:
+      value: The snapshot to put into the channel.
+    """
     ...
 
   def get(self) -> list[Snapshot]:

python/mujoco/experimental/studio/viewer.py

@@ -0,0 +1,69 @@
+# Copyright 2026 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+"""Interactive Studio GUI viewer for MuJoCo."""
+
+from absl import app as _app
+from absl import flags
+from mujoco.experimental.studio import launch_passive
+from mujoco.experimental.studio import messages
+from mujoco.experimental.studio import parser
+from mujoco.experimental.studio import sim
+from mujoco.experimental.studio import viewer_protocol
+
+vp = viewer_protocol
+
+_GFX = flags.DEFINE_enum('gfx', None, vp.GFX_MODES, 'Graphics mode.')
+_WIDTH = flags.DEFINE_integer('width', 1200, 'Width of output window.')
+_HEIGHT = flags.DEFINE_integer('height', 800, 'Height of output window.')
+_MJCF_PATH = flags.DEFINE_string('mjcf', None, 'Path to MJCF file.')
+_VIEWER = flags.DEFINE_enum_class(
+    'viewer', vp.ViewerMode.NATIVE, vp.ViewerMode, 'Viewer mode.'
+)
+
+
+def main(argv: list[str]) -> None:
+  config = vp.ViewerConfig(
+      width=_WIDTH.value,
+      height=_HEIGHT.value,
+      gfx=_GFX.value or '',
+      viewer_mode=_VIEWER.value,
+  )
+
+  # Get the model path, if provided.
+  model_path = None
+  if _MJCF_PATH.value is not None:
+    model_path = _MJCF_PATH.value
+  elif len(argv) > 1 and not argv[1].startswith('--'):
+    model_path = argv[1]
+
+  with launch_passive.launch_passive(config) as handle:
+    # Load the model if we have a path.
+    model, data = None, None
+    if model_path is not None:
+      data = parser.parse(model_path)
+      if data is not None:
+        model = data.model
+
+    # Send the model to the viewer, if we have a model.
+    if model is not None:
+      handle.send_to_viewer(messages.ModelEvent(model=model))
+
+    # Run the simulation.
+    step_control = sim.StepControl()
+    while handle.is_running():
+      step_control.advance(model, data)
+      model, data, step_control = handle.sync(model, data, step_control)
+
+
+_app.run(main)