chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor Studio implot sample to use a new message passing and launch APIs

Browse Source 
Deleted async examples since they add noise and are not particularly useful

PiperOrigin-RevId: 941117641
Change-Id: I4f621801f652e0b410a7acf905cda58143dc3fb0
This commit is contained in:
Matija Kecman
2026-07-01 08:41:48 -07:00
committed by Copybara-Service
parent a07ae6f849
commit 6e8a79c657
3 changed files with 80 additions and 447 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/experimental/studio/sample/async.py
-207
View File
@@ -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 <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()
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)
python/mujoco/experimental/studio/sample/async_thread.py
-185
View File
@@ -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 <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)
python/mujoco/experimental/studio/sample/implot.py
+80 -55
View File
@@ -24,22 +24,26 @@ import math
import os import os
import sys import sys
from absl import app as absl_app from absl import app as _app
from absl import flags as absl_flags from absl import flags as _flags
import mujoco import mujoco
from mujoco.experimental.studio import native_viewer as _viewer from mujoco.experimental.studio import launch_passive
from mujoco.experimental.studio import studio_app from mujoco.experimental.studio import messages as msg
from mujoco.experimental.studio import viewer_protocol 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 import numpy as np
from mujoco.experimental.dear_imgui import dear_imgui as imgui from mujoco.experimental.dear_imgui import dear_imgui as imgui
from mujoco.experimental.implot import implot from mujoco.experimental.implot import implot
_GFX = absl_flags.DEFINE_enum( _GFX = _flags.DEFINE_enum('gfx', None, vp.GFX_MODES, 'Graphics mode.')
'gfx', None, viewer_protocol.GFX_MODES, 'Rendering 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 _N_HISTORY = 100
@@ -102,47 +106,39 @@ def _setup_angle_axis(plot_size: imgui.Vec2) -> None:
) )
def main(argv: list[str]) -> None: class PlottingUi(va.ViewerGuiHook):
app = studio_app.StudioApp.from_argv(argv) """Custom GUI component maintaining history buffers for ImPlot charts."""
title = os.path.basename(sys.argv[0])
config = viewer_protocol.ViewerConfig( def __init__(self):
title=title, self.centroid = [np.zeros(3) for _ in range(_N_HISTORY)]
width=_WIDTH.value, self.euler = [np.zeros(3) for _ in range(_N_HISTORY)]
height=_HEIGHT.value, self.body_id = -1
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 build_gui(self, app: va.ViewerApp) -> None:
# Inspect the perturb.select body # Inspect the perturb.select body
if viewer.perturb.select > 0: if app.viewer.perturb.select > 0:
body_id = viewer.perturb.select self.body_id = app.viewer.perturb.select
# Display selected body information. # Display selected body information.
if body_id > 0: if self.body_id > 0:
body_name = mujoco.mj_id2name( 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) imgui.SetNextWindowSize(imgui.Vec2(1200, 600), imgui.Cond.FirstUseEver)
# Note: The window title uses the special "###" markup to ensure the imgui # 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 # ID for the window is constant for all body names. This is needed for
# the window to retain its state for all bodies. # 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): if imgui.Begin(window_title):
avail = imgui.GetContentRegionAvail() avail = imgui.GetContentRegionAvail()
wide = avail.x > avail.y wide = avail.x > avail.y
@@ -156,10 +152,10 @@ def main(argv: list[str]) -> None:
plot_flags = _setup_plot_flags(plot_size) plot_flags = _setup_plot_flags(plot_size)
if implot.BeginPlot('Centroid vs Time', plot_size, flags=plot_flags): if implot.BeginPlot('Centroid vs Time', plot_size, flags=plot_flags):
_setup_time_axis(plot_size) _setup_time_axis(plot_size)
_setup_xpos_axis(centroid, plot_size) _setup_xpos_axis(self.centroid, plot_size)
implot.PlotLine('x', range(_N_HISTORY), [c[0] for c in centroid]) implot.PlotLine('x', range(_N_HISTORY), [c[0] for c in self.centroid])
implot.PlotLine('y', range(_N_HISTORY), [c[1] for c in 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 centroid]) implot.PlotLine('z', range(_N_HISTORY), [c[2] for c in self.centroid])
implot.EndPlot() implot.EndPlot()
if wide: if wide:
@@ -168,34 +164,63 @@ def main(argv: list[str]) -> None:
if implot.BeginPlot('Euler Angle vs Time', plot_size, flags=plot_flags): if implot.BeginPlot('Euler Angle vs Time', plot_size, flags=plot_flags):
_setup_time_axis(plot_size) _setup_time_axis(plot_size)
_setup_angle_axis(plot_size) _setup_angle_axis(plot_size)
implot.PlotLine('roll', range(_N_HISTORY), [e[0] 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 euler]) implot.PlotLine(
implot.PlotLine('yaw', range(_N_HISTORY), [e[2] for e in euler]) '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() implot.EndPlot()
imgui.End() imgui.End()
# Update plot data # Update plot data
centroid.pop(0) self.centroid.pop(0)
euler.pop(0) self.euler.pop(0)
if body_id > 0: if self.body_id > 0 and self.body_id < app.model.nbody:
centroid.append(app.data.xpos[body_id].copy()) self.centroid.append(app.data.xpos[self.body_id].copy())
# Convert quaternion to Euler angles via rotation matrix. # 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) mat = np.zeros(9)
mujoco.mju_quat2Mat(mat, quat) mujoco.mju_quat2Mat(mat, quat)
# mat is row-major 3x3: R[i,j] = mat[3*i + j]. # mat is row-major 3x3: R[i,j] = mat[3*i + j].
roll = math.atan2(mat[7], mat[8]) roll = math.atan2(mat[7], mat[8])
pitch = math.atan2(-mat[6], math.sqrt(mat[7] ** 2 + mat[8] ** 2)) pitch = math.atan2(-mat[6], math.sqrt(mat[7] ** 2 + mat[8] ** 2))
yaw = math.atan2(mat[3], mat[0]) 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: else:
centroid.append(np.zeros(3)) self.centroid.append(np.zeros(3))
euler.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 <model_path.xml>')
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__': if __name__ == '__main__':
absl_app.run(main) _app.run(main)