diff --git a/python/mujoco/experimental/studio/sample/async_thread.py b/python/mujoco/experimental/studio/sample/async_thread.py new file mode 100644 index 00000000..5f4cbfd2 --- /dev/null +++ b/python/mujoco/experimental/studio/sample/async_thread.py @@ -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 ') + 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)