andrew
622 lines
18 KiB
Python
622 lines
18 KiB
Python
# Copyright 2022 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
|
||
#
|
||
# http://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.
|
||
# ==============================================================================
|
||
"""Interactive GUI viewer for MuJoCo."""
|
||
|
||
import abc
|
||
import atexit
|
||
import contextlib
|
||
import math
|
||
import os
|
||
import queue
|
||
import sys
|
||
import threading
|
||
import time
|
||
from typing import Callable, List, Optional, Tuple, Union
|
||
import weakref
|
||
|
||
import glfw
|
||
import numpy as np
|
||
|
||
import mujoco
|
||
from mujoco import _simulate
|
||
|
||
if not glfw._glfw: # pylint: disable=protected-access
|
||
raise RuntimeError('GLFW dynamic library handle is not available')
|
||
else:
|
||
_simulate.set_glfw_dlhandle(glfw._glfw._handle) # pylint: disable=protected-access
|
||
|
||
# Logarithmically spaced realtime slow-down coefficients (percent).
|
||
PERCENT_REALTIME = (
|
||
100, 80, 66, 50, 40, 33, 25, 20, 16, 13,
|
||
10, 8, 6.6, 5, 4, 3.3, 2.5, 2, 1.6, 1.3,
|
||
1, 0.8, 0.66, 0.5, 0.4, 0.33, 0.25, 0.2, 0.16, 0.13,
|
||
0.1
|
||
) # fmt: skip
|
||
|
||
# Maximum time mis-alignment before re-sync.
|
||
MAX_SYNC_MISALIGN = 0.1
|
||
|
||
# Fraction of refresh available for simulation.
|
||
SIM_REFRESH_FRACTION = 0.7
|
||
|
||
CallbackType = Callable[[mujoco.MjModel, mujoco.MjData], None]
|
||
LoaderType = Callable[[], Tuple[mujoco.MjModel, mujoco.MjData]]
|
||
KeyCallbackType = Callable[[int], None]
|
||
|
||
# Loader function that also returns a file path for the GUI to display.
|
||
_LoaderWithPathType = Callable[[], Tuple[mujoco.MjModel, mujoco.MjData, str]]
|
||
_InternalLoaderType = Union[LoaderType, _LoaderWithPathType]
|
||
|
||
_Simulate = _simulate.Simulate
|
||
|
||
|
||
class Handle:
|
||
"""A handle for interacting with a MuJoCo viewer."""
|
||
|
||
def __init__(
|
||
self,
|
||
sim: _Simulate,
|
||
cam: mujoco.MjvCamera,
|
||
opt: mujoco.MjvOption,
|
||
pert: mujoco.MjvPerturb,
|
||
user_scn: Optional[mujoco.MjvScene],
|
||
):
|
||
self._sim = weakref.ref(sim)
|
||
self._cam = cam
|
||
self._opt = opt
|
||
self._pert = pert
|
||
self._user_scn = user_scn
|
||
|
||
@property
|
||
def cam(self):
|
||
return self._cam
|
||
|
||
@property
|
||
def opt(self):
|
||
return self._opt
|
||
|
||
@property
|
||
def perturb(self):
|
||
return self._pert
|
||
|
||
@property
|
||
def user_scn(self):
|
||
return self._user_scn
|
||
|
||
@property
|
||
def m(self):
|
||
sim = self._sim()
|
||
if sim is not None:
|
||
return sim.m
|
||
return None
|
||
|
||
@property
|
||
def d(self):
|
||
sim = self._sim()
|
||
if sim is not None:
|
||
return sim.d
|
||
return None
|
||
|
||
@property
|
||
def viewport(self):
|
||
sim = self._sim()
|
||
if sim is not None:
|
||
return sim.viewport
|
||
return None
|
||
|
||
def set_figures(
|
||
self, viewports_figures: Union[Tuple[mujoco.MjrRect, mujoco.MjvFigure],
|
||
List[Tuple[mujoco.MjrRect, mujoco.MjvFigure]]]
|
||
):
|
||
"""Overlay figures on the viewer.
|
||
|
||
Args:
|
||
viewports_figures: Single tuple or list of tuples of (viewport, figure)
|
||
viewport: Rectangle defining position and size of the figure
|
||
figure: MjvFigure object containing the figure data to display
|
||
"""
|
||
sim = self._sim()
|
||
if sim is not None:
|
||
# Convert single tuple to list if needed
|
||
if isinstance(viewports_figures, tuple):
|
||
viewports_figures = [viewports_figures]
|
||
sim.set_figures(viewports_figures)
|
||
|
||
def clear_figures(self):
|
||
sim = self._sim()
|
||
if sim is not None:
|
||
sim.clear_figures()
|
||
|
||
def set_text(self, overlay_texts: Union[Tuple[Optional[int], Optional[int], Optional[str], Optional[str]],
|
||
List[Tuple[Optional[int], Optional[int], Optional[str], Optional[str]]]]):
|
||
"""Overlay text on the viewer.
|
||
|
||
Args:
|
||
overlay_texts: Single tuple or list of tuples of (font, gridpos, text1, text2)
|
||
font: Font style from mujoco.mjtFontScale
|
||
gridpos: Position of text box from mujoco.mjtGridPos
|
||
text1: Left text column, defaults to empty string if None
|
||
text2: Right text column, defaults to empty string if None
|
||
"""
|
||
sim = self._sim()
|
||
if sim is not None:
|
||
# Convert single tuple to list if needed
|
||
if isinstance(overlay_texts, tuple):
|
||
overlay_texts = [overlay_texts]
|
||
|
||
# Convert None values to empty strings
|
||
default_font = mujoco.mjtFontScale.mjFONTSCALE_150
|
||
default_gridpos = mujoco.mjtGridPos.mjGRID_TOPLEFT
|
||
processed_texts = [(
|
||
default_font if font is None else font,
|
||
default_gridpos if gridpos is None else gridpos,
|
||
"" if text1 is None else text1,
|
||
"" if text2 is None else text2)
|
||
for font, gridpos, text1, text2 in overlay_texts]
|
||
|
||
sim.set_text(processed_texts)
|
||
|
||
def clear_text(self):
|
||
sim = self._sim()
|
||
if sim is not None:
|
||
sim.clear_text()
|
||
|
||
def set_images(
|
||
self, viewports_images: Union[Tuple[mujoco.MjrRect, np.ndarray],
|
||
List[Tuple[mujoco.MjrRect, np.ndarray]]]
|
||
):
|
||
"""Overlay images on the viewer.
|
||
|
||
Args:
|
||
viewports_images: Single tuple or list of tuples of (viewport, image)
|
||
viewport: Rectangle defining position and size of the image
|
||
image: RGB image with shape (height, width, 3)
|
||
"""
|
||
sim = self._sim()
|
||
if sim is not None:
|
||
# Convert single tuple to list if needed
|
||
if isinstance(viewports_images, tuple):
|
||
viewports_images = [viewports_images]
|
||
|
||
processed_images = []
|
||
for viewport, image in viewports_images:
|
||
targ_shape = (viewport.height, viewport.width)
|
||
# Check if image is already the correct shape
|
||
if image.shape[:2] != targ_shape:
|
||
raise ValueError(f"Image shape {image.shape[:2]} does not match target shape {targ_shape}")
|
||
flipped = np.flip(image, axis=0)
|
||
contiguous = np.ascontiguousarray(flipped)
|
||
processed_images.append((viewport, contiguous))
|
||
sim.set_images(processed_images)
|
||
|
||
def clear_images(self):
|
||
sim = self._sim()
|
||
if sim is not None:
|
||
sim.clear_images()
|
||
|
||
def close(self):
|
||
sim = self._sim()
|
||
if sim is not None:
|
||
sim.exit()
|
||
|
||
def _get_sim(self) -> Optional[_Simulate]:
|
||
sim = self._sim()
|
||
if sim is not None:
|
||
try:
|
||
return sim if sim.exitrequest == 0 else None
|
||
except mujoco.UnexpectedError:
|
||
# UnexpectedError is raised when accessing `exitrequest` after the
|
||
# underlying simulate instance has been deleted in C++.
|
||
return None
|
||
return None
|
||
|
||
def is_running(self) -> bool:
|
||
return self._get_sim() is not None
|
||
|
||
def lock(self):
|
||
sim = self._get_sim()
|
||
if sim is not None:
|
||
return sim.lock()
|
||
return contextlib.nullcontext()
|
||
|
||
def sync(self):
|
||
sim = self._get_sim()
|
||
if sim is not None:
|
||
sim.sync() # locks internally
|
||
|
||
def update_hfield(self, hfieldid: int):
|
||
sim = self._get_sim()
|
||
if sim is not None:
|
||
sim.update_hfield(hfieldid) # locks internally and blocks until done
|
||
|
||
def update_mesh(self, meshid: int):
|
||
sim = self._get_sim()
|
||
if sim is not None:
|
||
sim.update_mesh(meshid) # locks internally and blocks until done
|
||
|
||
def update_texture(self, texid: int):
|
||
sim = self._get_sim()
|
||
if sim is not None:
|
||
sim.update_texture(texid) # locks internally and blocks until done
|
||
|
||
def __enter__(self):
|
||
return self
|
||
|
||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||
self.close()
|
||
|
||
|
||
# Abstract base dispatcher class for systems that require UI calls to be made
|
||
# on a specific thread (e.g. macOS). This is subclassed by system-specific
|
||
# Python launcher (mjpython) to implement the required dispatching mechanism.
|
||
class _MjPythonBase(metaclass=abc.ABCMeta):
|
||
|
||
def launch_on_ui_thread(
|
||
self,
|
||
model: mujoco.MjModel,
|
||
data: mujoco.MjData,
|
||
handle_return: Optional['queue.Queue[Handle]'],
|
||
key_callback: Optional[KeyCallbackType],
|
||
):
|
||
pass
|
||
|
||
|
||
# When running under mjpython, the launcher initializes this object.
|
||
_MJPYTHON: Optional[_MjPythonBase] = None
|
||
|
||
|
||
def _file_loader(path: str) -> _LoaderWithPathType:
|
||
"""Loads an MJCF model from file path."""
|
||
|
||
def load(path=path) -> Tuple[mujoco.MjModel, mujoco.MjData, str]:
|
||
if len(path) >= 4 and path[-4:] == '.mjb':
|
||
m = mujoco.MjModel.from_binary_path(path)
|
||
else:
|
||
m = mujoco.MjModel.from_xml_path(path)
|
||
d = mujoco.MjData(m)
|
||
return m, d, path
|
||
|
||
return load
|
||
|
||
|
||
def _reload(
|
||
simulate: _Simulate,
|
||
loader: _InternalLoaderType,
|
||
notify_loaded: Optional[Callable[[], None]] = None,
|
||
) -> Optional[Tuple[mujoco.MjModel, mujoco.MjData]]:
|
||
"""Internal function for reloading a model in the viewer."""
|
||
try:
|
||
simulate.load_message('') # path is unknown at this point
|
||
load_tuple = loader()
|
||
except Exception as e: # pylint: disable=broad-except
|
||
simulate.load_error = str(e)
|
||
simulate.load_message_clear()
|
||
else:
|
||
m, d = load_tuple[:2]
|
||
|
||
# If the loader does not raise an exception then we assume that it
|
||
# successfully created mjModel and mjData. This is specified in the type
|
||
# annotation, but we perform a runtime assertion here as well to prevent
|
||
# possible segmentation faults.
|
||
assert m is not None and d is not None
|
||
|
||
path = load_tuple[2] if len(load_tuple) == 3 else ''
|
||
simulate.load(m, d, path)
|
||
|
||
# Make sure any load_error message is cleared
|
||
simulate.load_error = ''
|
||
|
||
if notify_loaded:
|
||
notify_loaded()
|
||
|
||
return m, d
|
||
|
||
|
||
def _physics_loop(simulate: _Simulate, loader: Optional[_InternalLoaderType]):
|
||
"""Physics loop for the GUI, to be run in a separate thread."""
|
||
m: mujoco.MjModel = None
|
||
d: mujoco.MjData = None
|
||
ctrl_noise = np.array([])
|
||
reload = True
|
||
|
||
# CPU-sim synchronization point.
|
||
synccpu = 0.0
|
||
syncsim = 0.0
|
||
|
||
# Run until asked to exit.
|
||
while not simulate.exitrequest:
|
||
if simulate.droploadrequest:
|
||
simulate.droploadrequest = 0
|
||
loader = _file_loader(simulate.dropfilename)
|
||
reload = True
|
||
|
||
if simulate.uiloadrequest:
|
||
simulate.uiloadrequest_decrement()
|
||
reload = True
|
||
|
||
if reload and loader is not None:
|
||
result = _reload(simulate, loader)
|
||
if result is not None:
|
||
m, d = result
|
||
ctrl_noise = np.zeros((m.nu,))
|
||
|
||
reload = False
|
||
|
||
# Sleep for 1 ms or yield, to let main thread run.
|
||
if simulate.run != 0 and simulate.busywait != 0:
|
||
time.sleep(0)
|
||
else:
|
||
time.sleep(0.001)
|
||
|
||
with simulate.lock():
|
||
if m is not None:
|
||
assert d is not None
|
||
if simulate.run:
|
||
stepped = False
|
||
# Record CPU time at start of iteration.
|
||
startcpu = time.time()
|
||
|
||
elapsedcpu = startcpu - synccpu
|
||
elapsedsim = d.time - syncsim
|
||
|
||
# Inject noise.
|
||
if simulate.ctrl_noise_std != 0.0:
|
||
# Convert rate and scale to discrete time (Ornstein–Uhlenbeck).
|
||
rate = math.exp(
|
||
-m.opt.timestep / max(simulate.ctrl_noise_rate, mujoco.mjMINVAL)
|
||
)
|
||
scale = simulate.ctrl_noise_std * math.sqrt(1 - rate * rate)
|
||
|
||
for i in range(m.nu):
|
||
# Update noise.
|
||
ctrl_noise[i] = rate * ctrl_noise[
|
||
i
|
||
] + scale * mujoco.mju_standardNormal(None)
|
||
|
||
# Apply noise.
|
||
d.ctrl[i] = ctrl_noise[i]
|
||
|
||
# Requested slow-down factor.
|
||
slowdown = 100 / PERCENT_REALTIME[simulate.real_time_index]
|
||
|
||
# Misalignment: distance from target sim time > MAX_SYNC_MISALIGN.
|
||
misaligned = (
|
||
abs(elapsedcpu / slowdown - elapsedsim) > MAX_SYNC_MISALIGN
|
||
)
|
||
|
||
# Out-of-sync (for any reason): reset sync times, step.
|
||
if (
|
||
elapsedsim < 0
|
||
or elapsedcpu < 0
|
||
or synccpu == 0
|
||
or misaligned
|
||
or simulate.speed_changed
|
||
):
|
||
# Re-sync.
|
||
synccpu = startcpu
|
||
syncsim = d.time
|
||
simulate.speed_changed = False
|
||
|
||
# Run single step, let next iteration deal with timing.
|
||
mujoco.mj_step(m, d)
|
||
stepped = True
|
||
|
||
# In-sync: step until ahead of cpu.
|
||
else:
|
||
measured = False
|
||
prevsim = d.time
|
||
refreshtime = SIM_REFRESH_FRACTION / simulate.refresh_rate
|
||
# Step while sim lags behind CPU and within refreshtime.
|
||
while (
|
||
(d.time - syncsim) * slowdown < (time.time() - synccpu)
|
||
) and ((time.time() - startcpu) < refreshtime):
|
||
# Measure slowdown before first step.
|
||
if not measured and elapsedsim:
|
||
simulate.measured_slowdown = elapsedcpu / elapsedsim
|
||
measured = True
|
||
|
||
# Call mj_step.
|
||
mujoco.mj_step(m, d)
|
||
stepped = True
|
||
|
||
# Break if reset.
|
||
if d.time < prevsim:
|
||
break
|
||
|
||
# save current state to history buffer
|
||
if stepped:
|
||
simulate.add_to_history()
|
||
|
||
else: # simulate.run is False: GUI is paused.
|
||
|
||
# Run mj_forward, to update rendering and joint sliders.
|
||
mujoco.mj_forward(m, d)
|
||
simulate.speed_changed = True
|
||
|
||
|
||
def _launch_internal(
|
||
model: Optional[mujoco.MjModel] = None,
|
||
data: Optional[mujoco.MjData] = None,
|
||
*,
|
||
run_physics_thread: bool,
|
||
loader: Optional[_InternalLoaderType] = None,
|
||
handle_return: Optional['queue.Queue[Handle]'] = None,
|
||
key_callback: Optional[KeyCallbackType] = None,
|
||
show_left_ui: bool = True,
|
||
show_right_ui: bool = True,
|
||
) -> None:
|
||
"""Internal API, so that the public API has more readable type annotations."""
|
||
if model is None and data is not None:
|
||
raise ValueError('mjData is specified but mjModel is not')
|
||
elif callable(model) and data is not None:
|
||
raise ValueError(
|
||
'mjData should not be specified when an mjModel loader is used'
|
||
)
|
||
elif loader is not None and model is not None:
|
||
raise ValueError('model and loader are both specified')
|
||
elif run_physics_thread and handle_return is not None:
|
||
raise ValueError('run_physics_thread and handle_return are both specified')
|
||
|
||
if loader is None and model is not None:
|
||
|
||
def _loader(m=model, d=data) -> Tuple[mujoco.MjModel, mujoco.MjData]:
|
||
if d is None:
|
||
d = mujoco.MjData(m)
|
||
return m, d
|
||
|
||
loader = _loader
|
||
|
||
cam = mujoco.MjvCamera()
|
||
opt = mujoco.MjvOption()
|
||
pert = mujoco.MjvPerturb()
|
||
if model and not run_physics_thread:
|
||
user_scn = mujoco.MjvScene(model, _Simulate.MAX_GEOM)
|
||
else:
|
||
user_scn = None
|
||
simulate = _Simulate(
|
||
cam, opt, pert, user_scn, run_physics_thread, key_callback
|
||
)
|
||
|
||
simulate.ui0_enable = show_left_ui
|
||
simulate.ui1_enable = show_right_ui
|
||
|
||
# Initialize GLFW if not using mjpython.
|
||
if _MJPYTHON is None:
|
||
if not glfw.init():
|
||
raise mujoco.FatalError('could not initialize GLFW')
|
||
atexit.register(glfw.terminate)
|
||
|
||
notify_loaded = None
|
||
if handle_return:
|
||
notify_loaded = lambda: handle_return.put_nowait(
|
||
Handle(simulate, cam, opt, pert, user_scn)
|
||
)
|
||
|
||
if run_physics_thread:
|
||
side_thread = threading.Thread(
|
||
target=_physics_loop, args=(simulate, loader)
|
||
)
|
||
else:
|
||
side_thread = threading.Thread(
|
||
target=_reload, args=(simulate, loader, notify_loaded)
|
||
)
|
||
|
||
def make_exit(simulate):
|
||
def exit_simulate():
|
||
simulate.exit()
|
||
|
||
return exit_simulate
|
||
|
||
exit_simulate = make_exit(simulate)
|
||
atexit.register(exit_simulate)
|
||
|
||
side_thread.start()
|
||
simulate.render_loop()
|
||
atexit.unregister(exit_simulate)
|
||
side_thread.join()
|
||
simulate.destroy()
|
||
|
||
|
||
def launch(
|
||
model: Optional[mujoco.MjModel] = None,
|
||
data: Optional[mujoco.MjData] = None,
|
||
*,
|
||
loader: Optional[LoaderType] = None,
|
||
show_left_ui: bool = True,
|
||
show_right_ui: bool = True,
|
||
) -> None:
|
||
"""Launches the Simulate GUI."""
|
||
_launch_internal(
|
||
model,
|
||
data,
|
||
run_physics_thread=True,
|
||
loader=loader,
|
||
show_left_ui=show_left_ui,
|
||
show_right_ui=show_right_ui,
|
||
)
|
||
|
||
|
||
def launch_from_path(path: str) -> None:
|
||
"""Launches the Simulate GUI from file path."""
|
||
_launch_internal(run_physics_thread=True, loader=_file_loader(path))
|
||
|
||
|
||
def launch_passive(
|
||
model: mujoco.MjModel,
|
||
data: mujoco.MjData,
|
||
*,
|
||
key_callback: Optional[KeyCallbackType] = None,
|
||
show_left_ui: bool = True,
|
||
show_right_ui: bool = True,
|
||
) -> Handle:
|
||
"""Launches a passive Simulate GUI without blocking the running thread."""
|
||
if not isinstance(model, mujoco.MjModel):
|
||
raise ValueError(f'`model` is not a mujoco.MjModel: got {model!r}')
|
||
if not isinstance(data, mujoco.MjData):
|
||
raise ValueError(f'`data` is not a mujoco.MjData: got {data!r}')
|
||
if key_callback is not None and not callable(key_callback):
|
||
raise ValueError(f'`key_callback` is not callable: got {key_callback!r}')
|
||
|
||
mujoco.mj_forward(model, data)
|
||
handle_return = queue.Queue(1)
|
||
|
||
if sys.platform != 'darwin':
|
||
thread = threading.Thread(
|
||
target=_launch_internal,
|
||
args=(model, data),
|
||
kwargs=dict(
|
||
run_physics_thread=False,
|
||
handle_return=handle_return,
|
||
key_callback=key_callback,
|
||
show_left_ui=show_left_ui,
|
||
show_right_ui=show_right_ui,
|
||
),
|
||
)
|
||
thread.daemon = True
|
||
thread.start()
|
||
else:
|
||
if not isinstance(_MJPYTHON, _MjPythonBase):
|
||
raise RuntimeError(
|
||
'`launch_passive` requires that the Python script be run under '
|
||
'`mjpython` on macOS'
|
||
)
|
||
_MJPYTHON.launch_on_ui_thread(
|
||
model,
|
||
data,
|
||
handle_return,
|
||
key_callback,
|
||
show_left_ui,
|
||
show_right_ui,
|
||
)
|
||
|
||
return handle_return.get()
|
||
|
||
|
||
if __name__ == '__main__':
|
||
# pylint: disable=g-bad-import-order
|
||
from absl import app # pylint: disable=g-import-not-at-top
|
||
from absl import flags # pylint: disable=g-import-not-at-top
|
||
|
||
_MJCF_PATH = flags.DEFINE_string('mjcf', None, 'Path to MJCF file.')
|
||
|
||
def main(argv) -> None:
|
||
del argv
|
||
if _MJCF_PATH.value is not None:
|
||
launch_from_path(os.path.expanduser(_MJCF_PATH.value))
|
||
else:
|
||
launch()
|
||
|
||
app.run(main)
|