Improve async example and remove Network demo code

- Move SimToView and ViewToSim dataclasses from async.py to viewer_protocol.py (adding a user_data field) to consolidate IPC protocol definitions.
  - Remove the simulated Network class and latency UI controls from async.py to simplify the example and focus on core multiprocessing communication.
  - Use parser.parse() in main() for cleaner and more direct MuJoCo model loading.
  - Add cancel_join_thread() to multiprocessing queues and a join timeout to prevent atexit handlers from blocking during process shutdown.

PiperOrigin-RevId: 932327625
Change-Id: If91777dba934814b3f9faa17f9055869920575d7

python/mujoco/experimental/studio/sample/async.py

@@ -14,13 +14,11 @@
 """This script runs a simulation and viewer in separate processes communicating asynchronously.
 
 In this example, we will run the viewer in an independent process communicating
-via multiprocessing queues. Controls are provided to simulate network transit
+via multiprocessing queues. A slider is provided to adjust the state send rate.
-latency and adjust the communication rates.
 
 You must provide a mjcf model file via the first command-line argument.
 """
 
-import dataclasses
 import multiprocessing
 import os
 import sys
@@ -30,79 +28,41 @@ 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
+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, viewer_protocol.GFX_MODES, 'Rendering graphics mode.'
+    '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')
 
 
-@dataclasses.dataclass
-class SimToView:
-  """A message sent from the simulation process to the viewer process."""
-
-  model: mujoco.MjModel | None = None
-  data: mujoco.MjData | None = None
-  state: np.ndarray | None = None
-  state_sig: int = 0
-  send_time: float = 0.0
-
-
-@dataclasses.dataclass
-class ViewToSim:
-  """A message sent from the viewer process to the simulation process."""
-
-  state: np.ndarray | None = None
-  state_sig: int = 0
-  reset: bool = False
-  send_rate: float = 60.0
-
-
-class Network:
-  """Simulated networking parameters."""
-
-  def __init__(self) -> None:
-    self.transit_buffer = []
-    self.send_rate = 60.0
-    self.network_delay = 0.2
-
-  def get_arrived(self, q: multiprocessing.Queue) -> SimToView | None:
-    now = time.time()
-    while not q.empty():
-      self.transit_buffer.append(q.get())
-    arrived = None
-    while (
-        self.transit_buffer
-        and now >= self.transit_buffer[0].send_time + self.network_delay
-    ):
-      arrived = self.transit_buffer.pop(0)
-    return arrived
-
-
 def view(
     sim_to_view: multiprocessing.Queue,
     view_to_sim: multiprocessing.Queue,
 ) -> None:
   """Entry-point for process that renders the simulation."""
   # Block until the first message (containing the model) arrives.
-  msg = sim_to_view.get()
+  first_msg = sim_to_view.get()
-  assert msg.model is not None, 'First message must contain the MuJoCo model.'
+  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(msg.model, xfrc_sig)
+  xfrc_size = mujoco.mj_stateSize(first_msg.model, xfrc_sig)
   xfrc_state = np.zeros(xfrc_size, np.float64)
 
-  app = studio_app.StudioApp(msg.model, msg.data)
+  data = mujoco.MjData(first_msg.model)
-  network = Network()
+  app = studio_app.StudioApp(first_msg.model, data)
+  send_rate = 60.0
   viewer = _viewer.NativeViewer(
       app.model,
       title=title,
@@ -113,26 +73,28 @@ def view(
 
   while viewer.is_running() and app.is_running():
 
-    # Determine which messages have arrived through the simulated network.
+    # Drain the queue, keeping only the latest message.
-    arrived = network.get_arrived(sim_to_view)
+    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 arrived is not None and arrived.state is not None:
+    if msg is not None and msg.state is not None:
-      mujoco.mj_setState(app.model, app.data, arrived.state, arrived.state_sig)
+      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 arrived is not None:
+    if msg is not None:
       mujoco.mj_getState(app.model, app.data, xfrc_state, xfrc_sig)
       view_to_sim.put(
-          ViewToSim(
+          vp.ViewToSim(
-              send_rate=network.send_rate, state=xfrc_state, state_sig=xfrc_sig
+              send_rate=send_rate, state=xfrc_state, state_sig=xfrc_sig
           )
       )
 
@@ -145,24 +107,25 @@ def view(
         | int(imgui.WindowFlags.NoCollapse),
     ):
       imgui.PushItemWidth(200.0)
-      _, network.network_delay = imgui.SliderFloat(
+      updated, send_rate = imgui.SliderFloat(
-          'Network Latency (s)', network.network_delay, 0.0, 2.0
+          'Send Rate (Hz)', send_rate, 1.0, 120.0
-      )
-      updated, network.send_rate = imgui.SliderFloat(
-          'Send Rate (Hz)', network.send_rate, 1.0, 120.0
       )
       if updated:
-        view_to_sim.put(ViewToSim(send_rate=network.send_rate))
+        view_to_sim.put(vp.ViewToSim(send_rate=send_rate))
 
       imgui.SetNextItemWidth(-1)
       if imgui.Button('Reset Simulation'):
-        view_to_sim.put(ViewToSim(reset=True, send_rate=network.send_rate))
+        view_to_sim.put(vp.ViewToSim(reset=True, send_rate=send_rate))
 
       imgui.PopItemWidth()
     imgui.End()
 
     viewer.sync(app.model, app.data)
 
+  # Prevent multiprocessing.Queue atexit handler from blocking on exit.
+  sim_to_view.cancel_join_thread()
+  view_to_sim.cancel_join_thread()
+
 
 def sim(
     data: mujoco.MjData,
@@ -173,14 +136,14 @@ def sim(
 ) -> None:
   """Entry-point for process that runs the simulation."""
 
-  sim_to_view.put(SimToView(model=model, data=data))
+  sim_to_view.put(vp.SimToView(model=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 = ViewToSim()
+  msg = vp.ViewToSim()
   last_send_time = time.time()
 
   while view_process.is_alive():
@@ -204,17 +167,21 @@ def sim(
     if now - last_send_time >= 1.0 / max(1.0, msg.send_rate):
       mujoco.mj_getState(model, data, integration_state, integration_sig)
       sim_to_view.put(
-          SimToView(
+          vp.SimToView(
               state=integration_state,
               state_sig=integration_sig,
-              send_time=now,
           )
       )
       last_send_time = now
 
 
 def main(argv: list[str]) -> None:
-  app = studio_app.StudioApp.from_argv(argv)
+  if len(argv) < 2:
+    print('Usage: async <model_path.xml>')
+    sys.exit(1)
+
+  data = parser.parse(argv[1])
+  model = data.model
 
   # Queues for communication between the simulation and viewer processes.
   sim_to_view = multiprocessing.Queue()
@@ -227,7 +194,13 @@ def main(argv: list[str]) -> None:
   view_process.start()
 
   # Start the simulation in the main process.
-  sim(app.data, app.model, sim_to_view, view_to_sim, view_process)
+  sim(data, model, sim_to_view, view_to_sim, view_process)
+
+  # Prevent multiprocessing.Queue atexit handler from blocking on exit.
+  sim_to_view.cancel_join_thread()
+  view_to_sim.cancel_join_thread()
+
+  view_process.join(timeout=5.0)
 
 
 if __name__ == '__main__':

python/mujoco/experimental/studio/viewer_protocol.py

@@ -16,10 +16,12 @@
 StudioApp uses the protocol for convenience methods that accept any viewer.
 """
 
+import dataclasses
+from typing import Any
 from typing import Protocol
-
 import mujoco
 from mujoco.experimental.studio import ux
+import numpy as np
 
 GFX_MODES = (
     'classic',
@@ -33,6 +35,27 @@ GFX_MODES = (
 )
 
 
+@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
+  user_data: dict[str, Any] = dataclasses.field(default_factory=dict)
+
+
+@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
+  send_rate: float = 60.0
+  user_data: dict[str, Any] = dataclasses.field(default_factory=dict)
+
+
 class Viewer(Protocol):
   """Structural interface for any viewer."""