Configure Copybara export for Dear ImGui and ImPlot Python bindings
Following the export declarations in Dear ImGui and ImPlot METADATA, this change updates MuJoCo's Copybara configuration (copy.bara.sky) to export and transform the Python bindings. `//third_party/dear_imgui/google/py` exports to `python/mujoco/experimental/dear_imgui` and `//third_party/implot/google/py` exports to `python/mujoco/experimental/implot`. PiperOrigin-RevId: 925293624 Change-Id: Ie6e32d247a6f7fc24bb36ae7060f2075d8efeb26
This commit is contained in:
Matija Kecman
committed by
Copybara-Service
Copybara-Service
parent
062b0f1ea6
commit
4cf4a5665d
@@ -0,0 +1,231 @@
|
||||
# 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. Controls are provided to simulate network transit
|
||||
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
|
||||
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 sim as _sim
|
||||
from mujoco.experimental.studio import studio_app
|
||||
from mujoco.experimental.studio import ux
|
||||
import numpy as np
|
||||
|
||||
from mujoco.experimental.dear_imgui import dear_imgui as imgui
|
||||
|
||||
_GFX = absl_flags.DEFINE_string('gfx', '', '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()
|
||||
assert 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_state = np.zeros(xfrc_size, np.float64)
|
||||
|
||||
app = studio_app.StudioApp(msg.model, msg.data)
|
||||
network = Network()
|
||||
viewer = _viewer.NativeViewer(
|
||||
app.model,
|
||||
title=title,
|
||||
width=_WIDTH.value,
|
||||
height=_HEIGHT.value,
|
||||
gfx=_GFX.value,
|
||||
)
|
||||
|
||||
while viewer.is_running() and app.is_running():
|
||||
|
||||
# Determine which messages have arrived through the simulated network.
|
||||
arrived = network.get_arrived(sim_to_view)
|
||||
|
||||
# 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:
|
||||
mujoco.mj_setState(app.model, app.data, arrived.state, arrived.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:
|
||||
mujoco.mj_getState(app.model, app.data, xfrc_state, xfrc_sig)
|
||||
view_to_sim.put(
|
||||
ViewToSim(
|
||||
send_rate=network.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)
|
||||
_, network.network_delay = imgui.SliderFloat(
|
||||
'Network Latency (s)', network.network_delay, 0.0, 2.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))
|
||||
|
||||
imgui.SetNextItemWidth(-1)
|
||||
if imgui.Button('Reset Simulation'):
|
||||
view_to_sim.put(ViewToSim(reset=True, send_rate=network.send_rate))
|
||||
|
||||
imgui.PopItemWidth()
|
||||
imgui.End()
|
||||
|
||||
viewer.sync(app.model, app.data)
|
||||
|
||||
|
||||
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(SimToView(model=model, data=data))
|
||||
|
||||
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()
|
||||
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(
|
||||
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)
|
||||
|
||||
# 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(app.data, app.model, sim_to_view, view_to_sim, view_process)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
absl_app.run(main)
|
||||
@@ -0,0 +1,199 @@
|
||||
# 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.
|
||||
"""Example to run studio in the native viewer with responsive ImPlot UI.
|
||||
|
||||
This script runs a Studio viewer in-process and adds an 'Inspect Body' window
|
||||
using ImGui and ImPlot bindings to visualize selected body data. The example
|
||||
demonstrates how responsive UI layout rules are easily implemented.
|
||||
|
||||
Provide an MJCF model file via the first command-line argument to launch.
|
||||
"""
|
||||
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
|
||||
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 studio_app
|
||||
import numpy as np
|
||||
|
||||
from mujoco.experimental.dear_imgui import dear_imgui as imgui
|
||||
from mujoco.experimental.implot import implot
|
||||
|
||||
_GFX = absl_flags.DEFINE_string('gfx', '', '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')
|
||||
|
||||
|
||||
_N_HISTORY = 100
|
||||
|
||||
_PLOT_FLAGS = (
|
||||
implot.Flags.NoInputs.value # Disable pan/zoom mouse interaction.
|
||||
| implot.Flags.NoMenus.value # Disable right-click context menu.
|
||||
| implot.Flags.NoBoxSelect.value # Disable drag-to-select regions.
|
||||
)
|
||||
|
||||
_AXIS_FLAGS = (
|
||||
implot.AxisFlags.NoGridLines.value # Hide background grid lines.
|
||||
| implot.AxisFlags.NoTickMarks.value # Hide small tick marks on the axis.
|
||||
)
|
||||
|
||||
|
||||
def _setup_plot_flags(plot_size: imgui.Vec2) -> int:
|
||||
flags = _PLOT_FLAGS
|
||||
if min(plot_size.x, plot_size.y) < 300:
|
||||
flags |= implot.Flags.NoTitle.value
|
||||
if min(plot_size.x, plot_size.y) < 200:
|
||||
flags |= implot.Flags.NoLegend.value
|
||||
return flags
|
||||
|
||||
|
||||
def _setup_time_axis(plot_size: imgui.Vec2) -> None:
|
||||
flags = _AXIS_FLAGS
|
||||
if plot_size.x < 300:
|
||||
flags |= implot.AxisFlags.NoTickLabels.value
|
||||
implot.SetupAxis(implot.Axis.X1, '', flags)
|
||||
implot.SetupAxisLimits(implot.Axis.X1, 0, _N_HISTORY)
|
||||
|
||||
|
||||
def _setup_xpos_axis(centroid: list[np.ndarray], plot_size: imgui.Vec2) -> None:
|
||||
flags = _AXIS_FLAGS
|
||||
if plot_size.y < 300:
|
||||
flags |= implot.AxisFlags.NoTickLabels.value
|
||||
implot.SetupAxis(implot.Axis.Y1, '', flags)
|
||||
min_y = min(c[1] for c in centroid)
|
||||
max_y = max(c[1] for c in centroid)
|
||||
margin = max((max_y - min_y) * 0.1, 0.05)
|
||||
implot.SetupAxisLimits(
|
||||
implot.Axis.Y1,
|
||||
min_y - margin,
|
||||
max_y + margin,
|
||||
cond=implot.Cond.Always,
|
||||
)
|
||||
|
||||
|
||||
def _setup_angle_axis(plot_size: imgui.Vec2) -> None:
|
||||
flags = _AXIS_FLAGS
|
||||
if plot_size.y < 300:
|
||||
flags |= implot.AxisFlags.NoTickLabels.value
|
||||
implot.SetupAxis(implot.Axis.Y1, '', flags)
|
||||
implot.SetupAxisLimits(implot.Axis.Y1, -185.0, 185.0)
|
||||
implot.SetupAxisTicks(
|
||||
implot.Axis.Y1,
|
||||
[-180.0, -90.0, 0.0, 90.0, 180.0],
|
||||
['-180', '-90', '0', '90', '180'],
|
||||
)
|
||||
|
||||
|
||||
def main(argv: list[str]) -> None:
|
||||
app = studio_app.StudioApp.from_argv(argv)
|
||||
title = os.path.basename(sys.argv[0])
|
||||
|
||||
# Initialize the viewer.
|
||||
viewer = _viewer.NativeViewer(
|
||||
app.model,
|
||||
title=title,
|
||||
width=_WIDTH.value,
|
||||
height=_HEIGHT.value,
|
||||
gfx=_GFX.value,
|
||||
)
|
||||
|
||||
# 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)
|
||||
|
||||
# Inspect the perturb.select body
|
||||
if viewer.perturb.select > 0:
|
||||
body_id = viewer.perturb.select
|
||||
|
||||
# Display selected body information.
|
||||
if body_id > 0:
|
||||
body_name = mujoco.mj_id2name(
|
||||
app.model, int(mujoco.mjtObj.mjOBJ_BODY), body_id
|
||||
)
|
||||
|
||||
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'
|
||||
if imgui.Begin(window_title):
|
||||
avail = imgui.GetContentRegionAvail()
|
||||
wide = avail.x > avail.y
|
||||
|
||||
# Add a small padding factor to prevent scrollbars.
|
||||
plot_size = imgui.Vec2(
|
||||
avail.x * 0.5 - 4 if wide else avail.x,
|
||||
avail.y if wide else avail.y * 0.5 - 4,
|
||||
)
|
||||
|
||||
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])
|
||||
implot.EndPlot()
|
||||
|
||||
if wide:
|
||||
imgui.SameLine()
|
||||
|
||||
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.EndPlot()
|
||||
imgui.End()
|
||||
|
||||
# Update plot data
|
||||
centroid.pop(0)
|
||||
euler.pop(0)
|
||||
if body_id > 0:
|
||||
centroid.append(app.data.xpos[body_id].copy())
|
||||
# Convert quaternion to Euler angles via rotation matrix.
|
||||
quat = app.data.xquat[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])))
|
||||
else:
|
||||
centroid.append(np.zeros(3))
|
||||
euler.append(np.zeros(3))
|
||||
|
||||
viewer.sync(app.model, app.data)
|
||||
|
||||
viewer.stop()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
absl_app.run(main)
|
||||
@@ -0,0 +1,73 @@
|
||||
# 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.
|
||||
"""Render a MuJoCo model to an image."""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
from absl import app
|
||||
from absl import flags
|
||||
import mujoco
|
||||
from mujoco.experimental.studio import parser
|
||||
from mujoco.experimental.studio import renderer
|
||||
from PIL import Image
|
||||
|
||||
_MODEL = flags.DEFINE_string('model', '', 'Model file to load.')
|
||||
_OUTPUT = flags.DEFINE_string('output', '', 'Output file to save.')
|
||||
_GFX = flags.DEFINE_string('gfx', '', 'Renderer to use.')
|
||||
_WIDTH = flags.DEFINE_integer('width', 320, 'Width of the output image.')
|
||||
_HEIGHT = flags.DEFINE_integer('height', 240, 'Height of the output image.')
|
||||
_STEPS = flags.DEFINE_integer('steps', 1, 'Number of steps before render.')
|
||||
|
||||
|
||||
def main(argv):
|
||||
if len(argv) > 1:
|
||||
raise app.UsageError('Too many command-line arguments.')
|
||||
if not _MODEL.value:
|
||||
raise ValueError('`model` flag is required.')
|
||||
if not _OUTPUT.value:
|
||||
raise ValueError('`output flag is required.')
|
||||
|
||||
try:
|
||||
data = parser.parse(_MODEL.value)
|
||||
model = data.model
|
||||
except Exception as ex: # pylint: disable=broad-except
|
||||
print(f'Error loading model from `{_MODEL.value}`: {ex}')
|
||||
sys.exit(-1)
|
||||
|
||||
for _ in range(_STEPS.value):
|
||||
mujoco.mj_step(model, data)
|
||||
|
||||
try:
|
||||
r = renderer.Renderer(_GFX.value)
|
||||
r.Init(model)
|
||||
pixels = r.Render(
|
||||
model, data, None, None, None, _WIDTH.value, _HEIGHT.value
|
||||
)
|
||||
except Exception as ex: # pylint: disable=broad-except
|
||||
print(f'Error rendering model: {ex}')
|
||||
sys.exit(-2)
|
||||
|
||||
try:
|
||||
img = Image.frombytes('RGB', (_WIDTH.value, _HEIGHT.value), pixels)
|
||||
img.save(_OUTPUT.value, format=os.path.splitext(_OUTPUT.value)[1][1:])
|
||||
except Exception as ex: # pylint: disable=broad-except
|
||||
print(f'Error saving image to `{_OUTPUT.value}`: {ex}')
|
||||
sys.exit(-3)
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
app.run(main)
|
||||
Reference in New Issue
Block a user