diff --git a/python/mujoco/experimental/studio/sample/async.py b/python/mujoco/experimental/studio/sample/async.py deleted file mode 100644 index b3b60808..00000000 --- a/python/mujoco/experimental/studio/sample/async.py +++ /dev/null @@ -1,207 +0,0 @@ -# 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 processes communicating asynchronously. - -In this example, we will run the viewer in an independent process communicating -via multiprocessing queues. A slider is provided to adjust the state send rate. - -You must provide a mjcf model file via the first command-line argument. -""" - -import multiprocessing -import os -import sys -import time - -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: 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. - 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) - send_rate = 60.0 - config = vp.ViewerConfig( - title=title, - width=_WIDTH.value, - height=_HEIGHT.value, - gfx=_GFX.value, - ) - viewer = _viewer.NativeViewer(config) - - 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( - send_rate=send_rate, 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) - updated, send_rate = imgui.SliderFloat( - 'Send Rate (Hz)', send_rate, 1.0, 120.0 - ) - if updated: - view_to_sim.put(vp.ViewToSim(send_rate=send_rate)) - - imgui.SetNextItemWidth(-1) - if imgui.Button('Reset Simulation'): - 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, - model: mujoco.MjModel, - sim_to_view: multiprocessing.Queue, - view_to_sim: multiprocessing.Queue, - view_process: multiprocessing.Process, -) -> None: - """Entry-point for process that runs the simulation.""" - - 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 = vp.ViewToSim() - last_send_time = time.time() - - while view_process.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 process. - 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 paced by the requested send_rate. - now = time.time() - 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( - vp.SimToView( - state=integration_state, - state_sig=integration_sig, - ) - ) - last_send_time = now - - -def main(argv: list[str]) -> None: - if len(argv) < 2: - print('Usage: async ') - 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() - view_to_sim = multiprocessing.Queue() - - # Start the viewer process. - view_process = multiprocessing.Process( - target=view, args=(sim_to_view, view_to_sim) - ) - view_process.start() - - # Start the simulation in the main 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__': - absl_app.run(main) diff --git a/python/mujoco/experimental/studio/sample/async_thread.py b/python/mujoco/experimental/studio/sample/async_thread.py deleted file mode 100644 index afc4d72e..00000000 --- a/python/mujoco/experimental/studio/sample/async_thread.py +++ /dev/null @@ -1,185 +0,0 @@ -# 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) - config = vp.ViewerConfig( - title=title, - width=_WIDTH.value, - height=_HEIGHT.value, - gfx=_GFX.value, - ) - viewer = _viewer.NativeViewer(config) - - 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) diff --git a/python/mujoco/experimental/studio/sample/implot.py b/python/mujoco/experimental/studio/sample/implot.py index 8188ce00..6f001b6b 100644 --- a/python/mujoco/experimental/studio/sample/implot.py +++ b/python/mujoco/experimental/studio/sample/implot.py @@ -24,22 +24,26 @@ import math import os import sys -from absl import app as absl_app -from absl import flags as absl_flags +from absl import app as _app +from absl import flags as _flags import mujoco -from mujoco.experimental.studio import native_viewer as _viewer -from mujoco.experimental.studio import studio_app -from mujoco.experimental.studio import viewer_protocol +from mujoco.experimental.studio import launch_passive +from mujoco.experimental.studio import messages as msg +from mujoco.experimental.studio import parser +from mujoco.experimental.studio import sim +from mujoco.experimental.studio import viewer_app as va +from mujoco.experimental.studio import viewer_protocol as vp import numpy as np from mujoco.experimental.dear_imgui import dear_imgui as imgui from mujoco.experimental.implot import implot -_GFX = absl_flags.DEFINE_enum( - 'gfx', None, viewer_protocol.GFX_MODES, 'Rendering graphics mode.' +_GFX = _flags.DEFINE_enum('gfx', None, vp.GFX_MODES, 'Graphics mode.') +_WIDTH = _flags.DEFINE_integer('width', 1200, 'Width of the output image.') +_HEIGHT = _flags.DEFINE_integer('height', 800, 'Height of the output image') +_VIEWER = _flags.DEFINE_enum_class( + 'viewer', vp.ViewerMode.NATIVE, vp.ViewerMode, 'Viewer 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') _N_HISTORY = 100 @@ -102,47 +106,39 @@ def _setup_angle_axis(plot_size: imgui.Vec2) -> None: ) -def main(argv: list[str]) -> None: - app = studio_app.StudioApp.from_argv(argv) - title = os.path.basename(sys.argv[0]) +class PlottingUi(va.ViewerGuiHook): + """Custom GUI component maintaining history buffers for ImPlot charts.""" - config = viewer_protocol.ViewerConfig( - title=title, - width=_WIDTH.value, - height=_HEIGHT.value, - gfx=_GFX.value, - ) - viewer = _viewer.NativeViewer(config) - - # Variables for the custom UI. - centroid = [np.zeros(3) for _ in range(_N_HISTORY)] - euler = [np.zeros(3) for _ in range(_N_HISTORY)] - body_id = -1 - - # Main viewer loop. - while viewer.is_running(): - if not app.update(viewer.camera, viewer.vis_options, viewer.perturb): - break - - # Build standard Studio UI. - app.build_gui(viewer.camera, viewer.vis_options, viewer.render_flags) + def __init__(self): + self.centroid = [np.zeros(3) for _ in range(_N_HISTORY)] + self.euler = [np.zeros(3) for _ in range(_N_HISTORY)] + self.body_id = -1 + def build_gui(self, app: va.ViewerApp) -> None: # Inspect the perturb.select body - if viewer.perturb.select > 0: - body_id = viewer.perturb.select + if app.viewer.perturb.select > 0: + self.body_id = app.viewer.perturb.select # Display selected body information. - if body_id > 0: + if self.body_id > 0: body_name = mujoco.mj_id2name( - app.model, int(mujoco.mjtObj.mjOBJ_BODY), body_id + app.model, int(mujoco.mjtObj.mjOBJ_BODY), self.body_id ) + io = imgui.GetIO() + imgui.SetNextWindowPos( + imgui.Vec2(io.DisplaySize.x * 0.5, io.DisplaySize.y * 0.5), + imgui.Cond.FirstUseEver, + imgui.Vec2(0.5, 0.5), + ) imgui.SetNextWindowSize(imgui.Vec2(1200, 600), imgui.Cond.FirstUseEver) # Note: The window title uses the special "###" markup to ensure the imgui # ID for the window is constant for all body names. This is needed for # the window to retain its state for all bodies. - window_title = f'Inspect Body {body_name or "(???)"!r} ({body_id})###Plot' + window_title = ( + f'Inspect Body {body_name or "(???)"!r} ({self.body_id})###Plot' + ) if imgui.Begin(window_title): avail = imgui.GetContentRegionAvail() wide = avail.x > avail.y @@ -156,10 +152,10 @@ def main(argv: list[str]) -> None: plot_flags = _setup_plot_flags(plot_size) if implot.BeginPlot('Centroid vs Time', plot_size, flags=plot_flags): _setup_time_axis(plot_size) - _setup_xpos_axis(centroid, plot_size) - implot.PlotLine('x', range(_N_HISTORY), [c[0] for c in centroid]) - implot.PlotLine('y', range(_N_HISTORY), [c[1] for c in centroid]) - implot.PlotLine('z', range(_N_HISTORY), [c[2] for c in centroid]) + _setup_xpos_axis(self.centroid, plot_size) + implot.PlotLine('x', range(_N_HISTORY), [c[0] for c in self.centroid]) + implot.PlotLine('y', range(_N_HISTORY), [c[1] for c in self.centroid]) + implot.PlotLine('z', range(_N_HISTORY), [c[2] for c in self.centroid]) implot.EndPlot() if wide: @@ -168,34 +164,63 @@ def main(argv: list[str]) -> None: if implot.BeginPlot('Euler Angle vs Time', plot_size, flags=plot_flags): _setup_time_axis(plot_size) _setup_angle_axis(plot_size) - implot.PlotLine('roll', range(_N_HISTORY), [e[0] for e in euler]) - implot.PlotLine('pitch', range(_N_HISTORY), [e[1] for e in euler]) - implot.PlotLine('yaw', range(_N_HISTORY), [e[2] for e in euler]) + implot.PlotLine('roll', range(_N_HISTORY), [e[0] for e in self.euler]) + implot.PlotLine( + 'pitch', range(_N_HISTORY), [e[1] for e in self.euler] + ) + implot.PlotLine('yaw', range(_N_HISTORY), [e[2] for e in self.euler]) implot.EndPlot() imgui.End() # Update plot data - centroid.pop(0) - euler.pop(0) - if body_id > 0: - centroid.append(app.data.xpos[body_id].copy()) + self.centroid.pop(0) + self.euler.pop(0) + if self.body_id > 0 and self.body_id < app.model.nbody: + self.centroid.append(app.data.xpos[self.body_id].copy()) # Convert quaternion to Euler angles via rotation matrix. - quat = app.data.xquat[body_id] + quat = app.data.xquat[self.body_id] mat = np.zeros(9) mujoco.mju_quat2Mat(mat, quat) # mat is row-major 3x3: R[i,j] = mat[3*i + j]. roll = math.atan2(mat[7], mat[8]) pitch = math.atan2(-mat[6], math.sqrt(mat[7] ** 2 + mat[8] ** 2)) yaw = math.atan2(mat[3], mat[0]) - euler.append(np.degrees(np.array([roll, pitch, yaw]))) + self.euler.append(np.degrees(np.array([roll, pitch, yaw]))) else: - centroid.append(np.zeros(3)) - euler.append(np.zeros(3)) + self.centroid.append(np.zeros(3)) + self.euler.append(np.zeros(3)) - viewer.sync(app.model, app.data) - viewer.stop() +def main(argv: list[str]) -> None: + if len(argv) < 2: + print('Usage: implot ') + sys.exit(1) + + data = parser.parse(argv[1]) + if data is None: + print(f'Error loading model from {argv[1]!r}') + sys.exit(1) + model = data.model + + title = os.path.basename(sys.argv[0]) + plot_ui = PlottingUi() + + config = vp.ViewerConfig( + title=title, + width=_WIDTH.value, + height=_HEIGHT.value, + gfx=_GFX.value or '', + viewer_mode=_VIEWER.value, + ) + + with launch_passive.launch_passive(config, viewer_gui_hook=plot_ui) as handle: + handle.send_to_viewer(msg.ModelEvent(model=model)) + + step_control = sim.StepControl() + while handle.is_running(): + step_control.advance(model, data) + model, data, step_control = handle.sync(model, data, step_control) if __name__ == '__main__': - absl_app.run(main) + _app.run(main)