Add an async example using multi threading

PiperOrigin-RevId: 932400523 Change-Id: I6649dfd3a657ce11b7dc6cf130f01be8a85a095c
2026-06-15 05:32:31 -07:00
parent 9e45ce9801
commit 81bc303a8f
1 changed files with 185 additions and 0 deletions
@@ -0,0 +1,185 @@
+# 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.
+"""This script runs a simulation and viewer in separate threads communicating asynchronously.
+
+In this example, we will run the viewer in an independent thread communicating
+via thread-safe queues.
+
+You must provide a mjcf model file via the first command-line argument.
+"""
+
+import copy
+import os
+import queue
+import sys
+import threading
+
+from absl import app as absl_app
+from absl import flags as absl_flags
+import mujoco
+from mujoco.experimental.studio import native_viewer as _viewer
+from mujoco.experimental.studio import parser
+from mujoco.experimental.studio import sim as _sim
+from mujoco.experimental.studio import studio_app
+from mujoco.experimental.studio import ux
+from mujoco.experimental.studio import viewer_protocol as vp
+import numpy as np
+
+from mujoco.experimental.dear_imgui import dear_imgui as imgui
+
+_GFX = absl_flags.DEFINE_enum(
+    'gfx', None, vp.GFX_MODES, 'Rendering graphics mode.'
+)
+_WIDTH = absl_flags.DEFINE_integer('width', 1200, 'Width of the output image.')
+_HEIGHT = absl_flags.DEFINE_integer('height', 800, 'Height of the output image')
+
+
+def view(
+    sim_to_view: queue.Queue[vp.SimToView],
+    view_to_sim: queue.Queue[vp.ViewToSim],
+) -> None:
+  """Entry-point for thread that renders the simulation."""
+  # Block until the first message (containing the model) arrives.
+  first_msg = sim_to_view.get()
+  assert (
+      first_msg.model is not None
+  ), 'First message must contain the MuJoCo model.'
+
+  title = os.path.basename(sys.argv[0])
+  xfrc_sig = int(mujoco.mjtState.mjSTATE_XFRC_APPLIED)
+  xfrc_size = mujoco.mj_stateSize(first_msg.model, xfrc_sig)
+  xfrc_state = np.zeros(xfrc_size, np.float64)
+
+  data = mujoco.MjData(first_msg.model)
+  app = studio_app.StudioApp(first_msg.model, data)
+  viewer = _viewer.NativeViewer(
+      app.model,
+      title=title,
+      width=_WIDTH.value,
+      height=_HEIGHT.value,
+      gfx=_GFX.value,
+  )
+
+  while viewer.is_running() and app.is_running():
+
+    # Drain the queue, keeping only the latest message.
+    msg = None
+    while not sim_to_view.empty():
+      msg = sim_to_view.get()
+
+    # Update the camera and compute the perturbation.
+    app.handle_mouse_events(viewer.camera, viewer.vis_options, viewer.perturb)
+
+    # Sync state from the backend if a new payload actually arrived.
+    if msg is not None and msg.state is not None:
+      mujoco.mj_setState(app.model, app.data, msg.state, msg.state_sig)
+      mujoco.mj_forward(app.model, app.data)
+
+    # Always apply the perturbation forces from the viewer.
+    app.apply_perturb(viewer.perturb)
+
+    # Transmit user interaction when we get a new state
+    if msg is not None:
+      mujoco.mj_getState(app.model, app.data, xfrc_state, xfrc_sig)
+      view_to_sim.put(vp.ViewToSim(state=xfrc_state, state_sig=xfrc_sig))
+
+    # Build the UI.
+    ux.setup_theme(app.theme)
+    if imgui.Begin(
+        'Settings',
+        flags=int(imgui.WindowFlags.AlwaysAutoResize)
+        | int(imgui.WindowFlags.NoTitleBar)
+        | int(imgui.WindowFlags.NoCollapse),
+    ):
+      imgui.PushItemWidth(200.0)
+
+      imgui.SetNextItemWidth(-1)
+      if imgui.Button('Reset Simulation'):
+        view_to_sim.put(vp.ViewToSim(reset=True))
+
+      imgui.PopItemWidth()
+    imgui.End()
+
+    viewer.sync(app.model, app.data)
+
+
+def sim(
+    data: mujoco.MjData,
+    model: mujoco.MjModel,
+    sim_to_view: queue.Queue[vp.SimToView],
+    view_to_sim: queue.Queue[vp.ViewToSim],
+    view_thread: threading.Thread,
+) -> None:
+  """Entry-point for thread that runs the simulation."""
+
+  sim_to_view.put(vp.SimToView(model=copy.copy(model)))
+
+  step_control = _sim.StepControl()
+  integration_sig = int(mujoco.mjtState.mjSTATE_INTEGRATION)
+  integration_size = mujoco.mj_stateSize(model, integration_sig)
+  integration_state = np.empty(integration_size, np.float64)
+
+  msg = vp.ViewToSim()
+
+  while view_thread.is_alive():
+    while not view_to_sim.empty():
+      msg = view_to_sim.get()
+
+    if msg.reset:
+      mujoco.mj_resetData(model, data)
+      mujoco.mj_forward(model, data)
+      msg.reset = False
+
+    # Apply perturbation forces received from the viewer thread.
+    if msg.state is not None:
+      mujoco.mj_setState(model, data, msg.state, msg.state_sig)
+
+    # Advance the simulation keeping up with real-time.
+    step_control.advance(model, data)
+
+    # Send the simulation state every iteration. The in-process queue is
+    # essentially free and the viewer drains to keep only the latest.
+    mujoco.mj_getState(model, data, integration_state, integration_sig)
+    sim_to_view.put(
+        vp.SimToView(state=integration_state, state_sig=integration_sig)
+    )
+
+
+def main(argv: list[str]) -> None:
+  if len(argv) < 2:
+    print('Usage: async_thread <model_path.xml>')
+    sys.exit(1)
+
+  data = parser.parse(argv[1])
+  model = data.model
+
+  # Queues for communication between the simulation and viewer threads.
+  sim_to_view = queue.Queue()
+  view_to_sim = queue.Queue()
+
+  # Start the viewer thread.
+  view_thread = threading.Thread(
+      target=view, args=(sim_to_view, view_to_sim), name='Viewer'
+  )
+  view_thread.start()
+
+  # Start the simulation in the main thread.
+  sim(data, model, sim_to_view, view_to_sim, view_thread)
+
+  # Wait for the viewer thread to fully exit.
+  view_thread.join(timeout=5.0)
+
+
+if __name__ == '__main__':
+  absl_app.run(main)