diff --git a/.gitignore b/.gitignore index 4606f98..b120dbb 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,164 @@ dist/ *.egg-info mujoco_viewer/__pycache__ mujoco_viewer/MUJOCO_LOG.TXT + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/#use-with-ide +.pdm.toml + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ \ No newline at end of file diff --git a/README.md b/README.md index 0d47d24..6f55544 100644 --- a/README.md +++ b/README.md @@ -72,4 +72,4 @@ img = viewer.read_pixels(camid=2) - `title`: set the title of the window, for example: `viewer = mujoco_viewer.MujocoViewer(model, data, title='My Demo')` (defaults to `mujoco-python-viewer`). - `width`: set the window width, for example: `viewer = mujoco_viewer.MujocoViewer(model, data, width=300)` (defaults to full screen's width). - `height`: set the window height, for example: `viewer = mujoco_viewer.MujocoViewer(model, data, height=300)` (defaults to full screen's height). -- `hide_menus`: set whether the overlay menus should be hidden or not (defaults to `False`). +- `hide_menus`: set whether the overlay menus and graph should be hidden or not (defaults to `False`). diff --git a/examples/actuator_force_rendering.py b/examples/actuator_force_rendering.py new file mode 100644 index 0000000..d33dd42 --- /dev/null +++ b/examples/actuator_force_rendering.py @@ -0,0 +1,76 @@ +import mujoco +import mujoco_viewer +import math + +MODEL_XML = """ + + + +""" + +model = mujoco.MjModel.from_xml_string(MODEL_XML) +data = mujoco.MjData(model) + +# create the viewer object +viewer = mujoco_viewer.MujocoViewer(model, data, hide_menus=True) + +f_render_list = [ + ["hinge_1", "pos_servo_1", "force_at_hinge_1"], + ["hinge_2", "pos_servo_2", "force_at_hinge_2"], + ["hinge_3", "pos_servo_3", "force_at_hinge_3"], +] + +for _ in range(10000): + # Render forces + viewer.show_actuator_forces( + f_render_list=f_render_list, + rgba_list=[1, 0.5, 1, 0.5], + force_scale=0.05, + arrow_radius=0.05, + show_force_labels=True, + ) + + # step and render + mujoco.mj_step(model, data) + viewer.render() + if not viewer.is_alive: + break + +# close +viewer.close() diff --git a/examples/graphs.py b/examples/graphs.py new file mode 100644 index 0000000..81264cc --- /dev/null +++ b/examples/graphs.py @@ -0,0 +1,90 @@ +import mujoco +import mujoco_viewer +import math + +MODEL_XML = """ + + + +""" + +model = mujoco.MjModel.from_xml_string(MODEL_XML) +data = mujoco.MjData(model) + +# create the viewer object +viewer = mujoco_viewer.MujocoViewer(model, data) + + +# Initialize the graph lines +viewer.add_graph_line(line_name="line_1") +viewer.add_graph_line(line_name="downscaled_force_sensor") +viewer.add_graph_line(line_name="position_sensor") + +# Initialize the Legend +viewer.show_graph_legend(show_legend=True) + +# For a time-based graph, +# x_axis_time is the total time that you want your viewing window to see +# It needs to be set to grater than [model.opt.timestep*50] +# If you want to over ride that, change the "override" parameter to True +viewer.set_grid_divisions(x_div=5, + y_div=4, + x_axis_time=model.opt.timestep * 1000, + override=True) + +for i in range(10000): + viewer.update_graph_line(line_name="line_1", line_data=math.sin(i / 10.0)) + viewer.update_graph_line(line_name="downscaled_force_sensor", line_data=data.actuator_force[0]/100.0) + viewer.update_graph_line(line_name="position_sensor", line_data=data.sensordata[0]) + + # step and render + mujoco.mj_step(model, data) + viewer.render() + if not viewer.is_alive: + break + +# close +viewer.close() diff --git a/mujoco_viewer/callbacks.py b/mujoco_viewer/callbacks.py index 92c63b8..28c3062 100644 --- a/mujoco_viewer/callbacks.py +++ b/mujoco_viewer/callbacks.py @@ -18,6 +18,7 @@ def __init__(self, hide_menus): self._last_mouse_x = 0 self._last_mouse_y = 0 self._paused = False + self._hide_graph = False self._transparent = False self._contacts = False self._joints = False @@ -97,6 +98,9 @@ def _key_callback(self, window, key, scancode, action, mods): self.model.geom_rgba[:, 3] /= 5.0 else: self.model.geom_rgba[:, 3] *= 5.0 + # Toggle Graph overlay + elif key == glfw.KEY_G: + self._hide_graph = not self._hide_graph # Display inertia elif key == glfw.KEY_I: self._inertias = not self._inertias @@ -115,6 +119,10 @@ def _key_callback(self, window, key, scancode, action, mods): self.vopt.flags[ mujoco.mjtVisFlag.mjVIS_CONVEXHULL ] = self._convex_hull_rendering + # Reload Simulation + elif key == glfw.KEY_BACKSPACE: + mujoco.mj_resetData(self.model, self.data) + mujoco.mj_forward(self.model, self.data) # Wireframe Rendering elif key == glfw.KEY_W: self._wire_frame = not self._wire_frame diff --git a/mujoco_viewer/mujoco_viewer.py b/mujoco_viewer/mujoco_viewer.py index 75c7355..41a607d 100644 --- a/mujoco_viewer/mujoco_viewer.py +++ b/mujoco_viewer/mujoco_viewer.py @@ -6,7 +6,6 @@ import yaml from .callbacks import Callbacks - class MujocoViewer(Callbacks): def __init__( self, @@ -16,8 +15,12 @@ def __init__( title="mujoco-python-viewer", width=None, height=None, + window_start_x_pixel_offset=6, + window_start_y_pixel_offset=30, hide_menus=False): super().__init__(hide_menus) + if hide_menus is True: + self._hide_graph = True self.model = model self.data = data @@ -40,14 +43,23 @@ def __init__( if not height: _, height = glfw.get_video_mode(glfw.get_primary_monitor()).size - - if self.render_mode == 'offscreen': - glfw.window_hint(glfw.VISIBLE, 0) - + + # disable rendering, by default + glfw.window_hint(glfw.VISIBLE, 0) self.window = glfw.create_window( width, height, title, None, None) glfw.make_context_current(self.window) + glfw.set_window_pos(self.window, + window_start_x_pixel_offset, + window_start_y_pixel_offset) + + # Select if offscreen or on window mode + if self.render_mode == 'offscreen': + glfw.window_hint(glfw.VISIBLE, 0) + else: + glfw.show_window(self.window) glfw.swap_interval(1) + framebuffer_width, framebuffer_height = glfw.get_framebuffer_size( self.window) @@ -70,8 +82,37 @@ def __init__( self.cam = mujoco.MjvCamera() self.scn = mujoco.MjvScene(self.model, maxgeom=10000) self.pert = mujoco.MjvPerturb() + self.fig = mujoco.MjvFigure() + mujoco.mjv_defaultFigure(self.fig) + + # Points for sampling of sensors... dictates smoothness of graph + # Do not exceed: self._num_pnts = 1000 + self._num_pnts = 1000 + self._data_graph_line_names = [] + self._line_datas = [] + + for n in range(0, mujoco.mjMAXLINE): + for i in range(0, self._num_pnts): + self.fig.linedata[n][2 * i] = float(-i) + self.ctx = mujoco.MjrContext( self.model, mujoco.mjtFontScale.mjFONTSCALE_150.value) + + # Adjust placement and size of graph + width, height = glfw.get_framebuffer_size(self.window) + width_adjustment = width % 4 + self.graph_viewport = mujoco.MjrRect( + int(3 * width / 4) + width_adjustment, + 0, + int(width / 4), + int(height / 4), + ) + mujoco.mjr_figure(self.graph_viewport, self.fig, self.ctx) + self.fig.flg_extend = 1 + self.fig.flg_symmetric = 0 + + # Makes the graph to be in autorange + self.axis_autorange() # load camera from configuration (if available) pathlib.Path( @@ -118,6 +159,155 @@ def __init__( self._overlay = {} self._markers = [] + def set_grid_divisions(self, x_div: int, y_div: int, x_axis_time: float = 0.0, override=False): + if override is False: + assert x_axis_time >= self.model.opt.timestep * 50, "Set [x_axis_time] >= [self.model.opt.timestep * 50], inorder to get a suitable sampling rate" + self.fig.gridsize[0] = x_div + 1 + self.fig.gridsize[1] = y_div + 1 + if x_axis_time != 0.0: + self._num_pnts = x_axis_time / self.model.opt.timestep + print("self._num_pnts: ", self._num_pnts) + if self._num_pnts > 1000: + self._num_pnts = 1000 + new_x_axis_time = self.model.opt.timestep * self._num_pnts + print( + f"Minimum x_axis_time is: {new_x_axis_time}" + + " reduce the x_axis_time" + f" OR Maximum time_step is: " + + f"{self.model.opt.timestep*self._num_pnts}" + + " increase the timestep" + ) + # assert x_axis_time == + assert 1 <= self._num_pnts <= 1000, ( + "num_pnts should be [1000], it is currently:", + f"{self._num_pnts}", + ) + # self._num_pnts = num_pnts + self._time_per_div = (self.model.opt.timestep * self._num_pnts) / ( + x_div + ) + self.set_x_label( + xname=f"time/div: {self._time_per_div}s" + + f" total: {self.model.opt.timestep * self._num_pnts}" + ) + + def axis_autorange(self): + """ + Call this function to auto-range the graph + """ + self.fig.range[0][0] = 1.0 + self.fig.range[0][1] = -1.0 + self.fig.range[1][0] = 1.0 + self.fig.range[1][1] = -1.0 + + def set_graph_name(self, name: str): + assert type(name) == str, "name is not a string" + self.fig.title = name + + def show_graph_legend(self, show_legend: bool = True): + if show_legend is True: + for i in range(0, len(self._data_graph_line_names)): + self.fig.linename[i] = self._data_graph_line_names[i] + self.fig.flg_legend = True + + def set_x_label(self, xname: str): + assert type(xname) == str, "xname is not a string" + self.fig.xlabel = xname + + def add_graph_line(self, line_name, line_data=0.0): + assert ( + type(line_name) == str + ), f"Line_name is not a string: {type(line_name)}" + if line_name in self._data_graph_line_names: + print("line name already exists") + else: + self._data_graph_line_names.append(line_name) + self._line_datas.append(line_data) + + def update_graph_line(self, line_name, line_data): + if line_name in self._data_graph_line_names: + idx = self._data_graph_line_names.index(line_name) + self._line_datas[idx] = line_data + else: + raise NameError( + "line name is not valid, add it to list before calling update" + ) + + def sensorupdate(self): + pnt = int(mujoco.mju_min(self._num_pnts, self.fig.linepnt[0] + 1)) + + for n in range(0, len(self._line_datas)): + for i in range(pnt - 1, 0, -1): + self.fig.linedata[n][2 * i + 1] = self.fig.linedata[n][ + 2 * i - 1 + ] + self.fig.linepnt[n] = pnt + self.fig.linedata[n][1] = self._line_datas[n] + + def update_graph_size(self, size_div_x=None, size_div_y=None): + if size_div_x is None and size_div_y is None: + width, height = glfw.get_framebuffer_size(self.window) + width_adjustment = width % 3 + self.graph_viewport.left = int(2 * width / 3) + width_adjustment + self.graph_viewport.width = int(width / 3) + self.graph_viewport.height = int(height / 3) + + else: + assert size_div_x is not None and size_div_y is None, "" + width, height = glfw.get_framebuffer_size(self.window) + width_adjustment = width % size_div_x + self.graph_viewport.left = ( + int((size_div_x - 1) * width / size_div_x) + width_adjustment + ) + self.graph_viewport.width = int(width / size_div_x) + self.graph_viewport.height = int(height / size_div_x) + + def show_actuator_forces( + self, + f_render_list, + rgba_list=[1, 0, 1, 1], + force_scale=0.05, + arrow_radius=0.03, + show_force_labels=False, + ) -> None: + """f_render_list: [ ["jnt_name1","act_name_1","lable1"] , + ["jnt_name2","act_name_2","lable2"] ] + """ + if show_force_labels is False: + for i in range(0, len(f_render_list)): + self.add_marker( + pos=self.data.joint(f_render_list[i][0]).xanchor, + mat=self.rotation_matrix_from_vectors( + vec1=[0.0, 0.0, 1.0], + vec2=self.data.joint(f_render_list[i][0]).xaxis), + size=[ + arrow_radius, + arrow_radius, + self.data.actuator(f_render_list[i][1]).force* force_scale, + ], + rgba=rgba_list, + type=mujoco.mjtGeom.mjGEOM_ARROW, + label="", + ) + else: + for i in range(0, len(f_render_list)): + self.add_marker( + pos=self.data.joint(f_render_list[i][0]).xanchor, + mat=self.rotation_matrix_from_vectors( + vec1=[0.0, 0.0, 1.0], + vec2=self.data.joint(f_render_list[i][0]).xaxis), + size=[ + arrow_radius, + arrow_radius, + self.data.actuator(f_render_list[i][1]).force* force_scale, + ], + rgba=rgba_list, + type=mujoco.mjtGeom.mjGEOM_ARROW, + label=f_render_list[i][2] + + ":" + + str(self.data.actuator(f_render_list[i][1]).force[0]), + ) + def add_marker(self, **marker_params): self._markers.append(marker_params) @@ -205,6 +395,10 @@ def add_overlay(gridpos, text1, text2): topleft, "[J]oints", "On" if self._joints else "Off") + add_overlay( + topleft, + "[G]raph Viewer", + "Off" if self._hide_graph else "On") add_overlay( topleft, "[I]nertia", @@ -250,6 +444,8 @@ def add_overlay(gridpos, text1, text2): add_overlay(topleft, "Cap[t]ure frame", "Saved as %s" % fname) else: add_overlay(topleft, "Cap[t]ure frame", "") + add_overlay(topleft, "[ESC] to Quit Application", "") + add_overlay(topleft, "[BACKSPACE] to Reload Sim", "") add_overlay( bottomleft, "FPS", "%d%s" % @@ -355,6 +551,27 @@ def update(): t1, t2, self.ctx) + + # Handle graph and pausing interactions + if ( + not self._paused + and not self._hide_graph + ): + self.sensorupdate() + self.update_graph_size() + mujoco.mjr_figure( + self.graph_viewport, self.fig, self.ctx + ) + elif self._hide_graph and self._paused: + self.update_graph_size() + elif not self._hide_graph and self._paused: + mujoco.mjr_figure( + self.graph_viewport, self.fig, self.ctx + ) + elif self._hide_graph and not self._paused: + self.sensorupdate() + self.update_graph_size() + glfw.swap_buffers(self.window) glfw.poll_events() self._time_per_render = 0.9 * self._time_per_render + \ @@ -391,3 +608,17 @@ def close(self): self.is_alive = False glfw.terminate() self.ctx.free() + + def rotation_matrix_from_vectors(self, vec1, vec2): + """ Find the rotation matrix that aligns vec1 to vec2 + :param vec1: A 3d "source" vector + :param vec2: A 3d "destination" vector + :return mat: A transform matrix (3x3) which when applied to vec1, aligns it with vec2. + """ + a, b = (vec1 / np.linalg.norm(vec1)).reshape(3), (vec2 / np.linalg.norm(vec2)).reshape(3) + v = np.cross(a, b) + c = np.dot(a, b) + s = np.linalg.norm(v) + kmat = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]]) + rotation_matrix = np.eye(3) + kmat + kmat.dot(kmat) * ((1 - c) / (s ** 2)) + return rotation_matrix