diff --git a/.github/workflows/autopep8.yml b/.github/workflows/autopep8.yml deleted file mode 100644 index 42d63bb..0000000 --- a/.github/workflows/autopep8.yml +++ /dev/null @@ -1,21 +0,0 @@ -name: Format python code -on: push -jobs: - autopep8: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v2 - - name: autopep8 - uses: peter-evans/autopep8@v1 - with: - args: $(git ls-files '*.py') --recursive --in-place --aggressive --aggressive . - - name: Create Pull Request - uses: peter-evans/create-pull-request@v3 - with: - commit-message: autopep8 action fixes - title: Fixes by autopep8 action - body: This is an auto-generated PR with fixes by autopep8. - labels: autopep8, automated pr - reviewers: rohanpsingh - branch: autopep8-patches - base: ${{ github.head_ref }} diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml deleted file mode 100644 index cc88396..0000000 --- a/.github/workflows/python-publish.yml +++ /dev/null @@ -1,37 +0,0 @@ -# This workflow will upload a Python Package using Twine when a release is created -# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries - -# This workflow uses actions that are not certified by GitHub. -# They are provided by a third-party and are governed by -# separate terms of service, privacy policy, and support -# documentation. - -name: Upload Python Package - -on: - release: - types: [published] - -permissions: - contents: read - -jobs: - deploy: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v3 - - name: Set up Python - uses: actions/setup-python@v3 - with: - python-version: '3.x' - - name: Install dependencies - run: | - python -m pip install --upgrade pip - pip install build - - name: Build package - run: python -m build - - name: Publish package - uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29 - with: - user: __token__ - password: ${{ secrets.PYPI_API_TOKEN }} diff --git a/LICENSE b/LICENSE index 74df77e..7fb68a3 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,9 @@ BSD 2-Clause License -Copyright (c) 2022, Rohan P. Singh +mujoco_py +├── Copyright (c) 2022, Rohan P. Singh +│   └── Copyright (c) 2022, Jack Xu + All rights reserved. Redistribution and use in source and binary forms, with or without diff --git a/README.md b/README.md index d40dfdb..7767bff 100644 --- a/README.md +++ b/README.md @@ -1,75 +1,27 @@ # Viewer for MuJoCo in Python -Interactive renderer to use with the official Python bindings for MuJoCo. - -Starting with version 2.1.2, MuJoCo comes with native Python bindings officially supported by the MuJoCo devs. - -If you have been a user of `mujoco-py`, you might be looking to migrate. -Some pointers on migration are available [here](https://mujoco.readthedocs.io/en/latest/python.html#migration-notes-for-mujoco-py). - -# Install +Interactive renderer to use with the official Python bindings for MuJoCo 2.2.x + ++ Develop along with [jx-mujoco-python-engine](https://github.com/jaku-jaku/jx-mujoco-python-engine) ++ A more efficient, robust and thread safe viewer with native MuJoCo 2.2.x with glfw ++ Works with Mac M1 with MuJoCo 2.2.x +## Points to ponder: +> This repo is/will-be a re-implementation of https://github.com/rohanpsingh/mujoco-python-viewer with custom features +> Additional credits: I do see a lot of similar structure as used in `mjviewer` from mujoco-py, so I will also give credits to https://github.com/openai/mujoco-py/blob/master/mujoco_py/mjviewer.py +> Starting with version 2.1.2, MuJoCo comes with native Python bindings officially supported by the MuJoCo devs. +> If you have been a user of `mujoco-py`, you might be looking to migrate. +> Some pointers on migration are available [here](https://mujoco.readthedocs.io/en/latest/python.html#migration-notes-for-mujoco-py). + +## Tested OS Platform: +- Mac M1 +- Ubuntu 18.04 + +## Install ```sh -$ git clone https://github.com/rohanpsingh/mujoco-python-viewer -$ cd mujoco-python-viewer +$ git clone https://github.com/jaku-jaku/jx-mujoco-python-viewer +$ cd jx-mujoco-python-viewer $ pip install -e . ``` -Or, install via Pip. -```sh -$ pip install mujoco-python-viewer -``` - -# Usage -#### How to render in a window? -```py -import mujoco -import mujoco_viewer - -model = mujoco.MjModel.from_xml_path('humanoid.xml') -data = mujoco.MjData(model) - -# create the viewer object -viewer = mujoco_viewer.MujocoViewer(model, data) - -# simulate and render -for _ in range(10000): - if viewer.is_alive: - mujoco.mj_step(model, data) - viewer.render() - else: - break - -# close -viewer.close() -``` - -The render should pop up and the simulation should be running. -Double-click on a geom and hold `Ctrl` to apply forces (right) and torques (left). - -![ezgif-2-6758c40cdf](https://user-images.githubusercontent.com/16384313/161459985-a47e74dc-92c9-4a0b-99fc-92d1b5b04163.gif) - - -Press `ESC` to quit. -Other key bindings are shown in the overlay menu (Press `H` or hold `Alt`). - - - -#### How to render offscreen? -```py -import mujoco -import mujoco_viewer - -model = mujoco.MjModel.from_xml_path('humanoid.xml') -data = mujoco.MjData(model) - -viewer = mujoco_viewer.MujocoViewer(model, data, 'offscreen') -mujoco.mj_forward(model, data) -img = viewer.read_pixels(camid=2) -## do something cool with img -``` - -# Optional Parameters -- `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 `True`). +## Demo (WIP): +- https://github.com/UW-Advanced-Robotics-Lab/simulation-mujoco-summit-wam \ No newline at end of file diff --git a/examples/markers_demo.py b/examples/markers_demo.py index 1a40a86..fe374af 100644 --- a/examples/markers_demo.py +++ b/examples/markers_demo.py @@ -28,7 +28,7 @@ # create the viewer object viewer = mujoco_viewer.MujocoViewer(model, data) -while viewer.is_alive: +while True: # sim step mujoco.mj_step(model, data) @@ -36,29 +36,31 @@ x_dir = [[0, 0, 1], [0, 1, 0], [-1, 0, 0]] y_dir = [[1, 0, 0], [0, 0, 1], [0, -1, 0]] z_dir = [[1, 0, 0], [0, 1, 0], [0, 0, 1]] - viewer.add_marker( + viewer.add_marker_safe( pos=[ 0, 0, 0], size=[ 0.05, 0.05, 0.05], rgba=[ 1, 1, 1, 1], type=mujoco.mjtGeom.mjGEOM_SPHERE, label="origin") - viewer.add_marker( + viewer.add_marker_safe( pos=[ 0, 0, 0], mat=x_dir, size=[ 0.01, 0.01, 2], rgba=[ 1, 0, 0, 0.2], type=mujoco.mjtGeom.mjGEOM_ARROW, label="") - viewer.add_marker( + viewer.add_marker_safe( pos=[ 0, 0, 0], mat=y_dir, size=[ 0.01, 0.01, 2], rgba=[ 0, 1, 0, 0.2], type=mujoco.mjtGeom.mjGEOM_ARROW, label="") - viewer.add_marker( + viewer.add_marker_safe( pos=[ 0, 0, 0], mat=z_dir, size=[ 0.01, 0.01, 2], rgba=[ 0, 0, 1, 0.2], type=mujoco.mjtGeom.mjGEOM_ARROW, label="") # render - viewer.render() - + viewer.process_safe() + viewer.update_safe() + viewer.render_safe() + # close -viewer.close() +viewer.signal_termination_safe(if_immediate=True) diff --git a/examples/offscreen_demo.py b/examples/offscreen_demo.py index 0ba2a2c..549a53b 100644 --- a/examples/offscreen_demo.py +++ b/examples/offscreen_demo.py @@ -5,8 +5,12 @@ model = mujoco.MjModel.from_xml_path('humanoid.xml') data = mujoco.MjData(model) -viewer = mujoco_viewer.MujocoViewer(model, data, 'offscreen') +viewer = mujoco_viewer.MujocoViewer(model, data, if_on_scrn=False) mujoco.mj_step(model, data) -img = viewer.read_pixels(camid=2) -img = PIL.Image.fromarray(img) -img.show() +viewer.process_safe() +viewer.update_safe() +viewer.render_safe() + +# img = viewer.read_pixels(camid=2) +# img = PIL.Image.fromarray(img) +# img.show() diff --git a/examples/sample.py b/examples/sample.py index dd7ad03..08d4631 100644 --- a/examples/sample.py +++ b/examples/sample.py @@ -10,9 +10,9 @@ # simulate and render for _ in range(100000): mujoco.mj_step(model, data) - viewer.render() - if not viewer.is_alive: - break - + viewer.process_safe() + viewer.update_safe() + viewer.render_safe() + # close -viewer.close() +viewer.signal_termination_safe(if_immediate=True) diff --git a/examples/simple_linkage.py b/examples/simple_linkage.py index 7c25e27..f5b2deb 100644 --- a/examples/simple_linkage.py +++ b/examples/simple_linkage.py @@ -59,7 +59,9 @@ # simulate and render for _ in range(100000): mujoco.mj_step(model, data) - viewer.render() - + viewer.process_safe() + viewer.update_safe() + viewer.render_safe() + # close -viewer.close() +viewer.signal_termination_safe(if_immediate=True) diff --git a/mujoco_viewer/mujoco_viewer.py b/mujoco_viewer/mujoco_viewer.py index 5d63329..ba5c66b 100644 --- a/mujoco_viewer/mujoco_viewer.py +++ b/mujoco_viewer/mujoco_viewer.py @@ -1,362 +1,645 @@ -import mujoco -import glfw -from threading import Lock +#! /usr/bin/env python +""" `mujoco_viewer.py` + + @author: Jack (Jianxiang) Xu + Contacts : projectbyjx@gmail.com + Last edits : July 27, 2022 + + @description: + This library will provide rendering pipeline (thread-safe) +""" +#===================================# +# I M P O R T - L I B R A R I E S # +#===================================# + +# python libraries: +from ast import Lambda +from turtle import left import numpy as np import time import imageio +import copy +import sys +from typing import Optional, Dict, Union, Any, Tuple +from enum import IntFlag, auto +from dataclasses import dataclass -class MujocoViewer: - def __init__( - self, - model, - data, - mode='window', - title="mujoco-python-viewer", - width=None, - height=None, - hide_menus=True): - self.model = model - self.data = data - self.render_mode = mode - if self.render_mode not in ['offscreen', 'window']: - raise NotImplementedError( - "Invalid mode. Only 'offscreen' and 'window' are supported.") - - self.is_alive = True - - self._gui_lock = Lock() - self._button_left_pressed = False - self._button_right_pressed = False - self._left_double_click_pressed = False - self._right_double_click_pressed = False - self._last_left_click_time = None - self._last_right_click_time = None - self._last_mouse_x = 0 - self._last_mouse_y = 0 - self._paused = False - self._transparent = False - self._contacts = False - self._joints = False - self._shadows = True - self._wire_frame = False - self._convex_hull_rendering = False - self._inertias = False - self._com = False - self._render_every_frame = True - self._image_idx = 0 - self._image_path = "/tmp/frame_%07d.png" - self._time_per_render = 1 / 60.0 - self._run_speed = 1.0 - self._loop_count = 0 - self._advance_by_one_step = False - self._hide_menus = hide_menus - - # glfw init - glfw.init() - - if not width: - width, _ = glfw.get_video_mode(glfw.get_primary_monitor()).size - - if not height: - _, height = glfw.get_video_mode(glfw.get_primary_monitor()).size - - if self.render_mode == 'offscreen': - glfw.window_hint(glfw.VISIBLE, 0) - - self.window = glfw.create_window( - width, height, title, None, None) - glfw.make_context_current(self.window) - glfw.swap_interval(1) - - framebuffer_width, framebuffer_height = glfw.get_framebuffer_size( - self.window) - - # install callbacks only for 'window' mode - if self.render_mode == 'window': - window_width, _ = glfw.get_window_size(self.window) - self._scale = framebuffer_width * 1.0 / window_width - - # set callbacks - glfw.set_cursor_pos_callback( - self.window, self._cursor_pos_callback) - glfw.set_mouse_button_callback( - self.window, self._mouse_button_callback) - glfw.set_scroll_callback(self.window, self._scroll_callback) - glfw.set_key_callback(self.window, self._key_callback) - - # create options, camera, scene, context - self.vopt = mujoco.MjvOption() - self.cam = mujoco.MjvCamera() - self.scn = mujoco.MjvScene(self.model, maxgeom=10000) - self.pert = mujoco.MjvPerturb() - self.ctx = mujoco.MjrContext( - self.model, mujoco.mjtFontScale.mjFONTSCALE_150.value) - - # get viewport - self.viewport = mujoco.MjrRect( - 0, 0, framebuffer_width, framebuffer_height) +# python 3rd party libraries: +import mujoco +import glfw +from threading import Lock - # overlay, markers - self._overlay = {} - self._markers = [] +from icecream import ic + + +# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # +# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # +# # ___ ___ ___ ___ ___ ___ # # +# # /__/\ /__/\ / /\ / /\ / /\ / /\ # # +# # | |::\ \ \:\ / /:/ / /::\ / /:/ / /::\ # # +# # | |:|:\ \ \:\ /__/::\ / /:/\:\ / /:/ / /:/\:\ # # +# # __|__|:|\:\ ___ \ \:\ \__\/\:\ / /:/ \:\ / /:/ ___ / /:/ \:\ # # +# # /__/::::| \:\ /__/\ \__\:\ \ \:\ /__/:/ \__\:\ /__/:/ / /\ /__/:/ \__\:\ # # +# # \ \:\~~\__\/ \ \:\ / /:/ \__\:\ \ \:\ / /:/ \ \:\ / /:/ \ \:\ / /:/ # # +# # \ \:\ \ \:\ /:/ / /:/ \ \:\ /:/ \ \:\ /:/ \ \:\ /:/ # # +# # \ \:\ \ \:\/:/ /__/:/ \ \:\/:/ \ \:\/:/ \ \:\/:/ # # +# # \ \:\ \ \::/ \__\/ \ \::/ \ \::/ \ \::/ # # +# # \__\/ \__\/ \__\/ \__\/ \__\/ # # +# # # # +# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # +# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # + +#=======================# +# D E F I N I T I O N # +#=======================# +class _VIEWER_STATUS_MAP(IntFlag): + # Status: + IS_ALIVE = (1<<0) + # Visual: + COM = (1<<2) + JOINTS = (1<<3) + SHADOWS = (1<<4) + CONTACTS = (1<<5) + INERTIAS = (1<<6) + WIRE_FRAME = (1<<7) + TRANSPARENT = (1<<8) + RENDER_ON_SCREEN = (1<<9) + TOGGLE_COORD_FRAME = (1<<10) + CONVEX_HULL_RENDERING = (1<<11) + HIDE_MENUS = (1<<12) + GROUP_0 = (1<<13) + GROUP_1 = (1<<14) + GROUP_2 = (1<<15) + GROUP_3 = (1<<16) + # MOD KEY: + LEFT_SHIFT_PRESSED = (1<<17) + LEFT_CONTROL_PRESSED = (1<<18) + # Functional: + TERMINATE = (1<<1) + CAPTURE_SCRNSHOT = (1<<20) + SAVE_SCRNSHOT_ON_FLY = (1<<19) + # TODO: to-be-implemented + PAUSED = (1<<21) # TODO: to implement on Engine side (partially implemented here) + STEP_BY_STEP = (1<<22) # TODO: to implement on Engine side (unset bit unless engine resets it) + RENDER_EVERY_FRAME = (1<<23) # TODO: to be implemented + # mux configs: + DEFAULT = GROUP_0 | GROUP_1 | GROUP_2 | GROUP_3 | SHADOWS + +@dataclass +class _CAMERA_DATA: + frame_buffer = None + depth_buffer = None + mj_cam = None + mj_viewport = None + frame_stamp = None +class _MJ_CONFIG: + vopt = None + pert = None + scn = None + ctx = None - def _key_callback(self, window, key, scancode, action, mods): - if action != glfw.RELEASE: - if key == glfw.KEY_LEFT_ALT: - self._hide_menus = False - return - # Switch cameras - elif key == glfw.KEY_TAB: - self.cam.fixedcamid += 1 - self.cam.type = mujoco.mjtCamera.mjCAMERA_FIXED - if self.cam.fixedcamid >= self.model.ncam: - self.cam.fixedcamid = -1 - self.cam.type = mujoco.mjtCamera.mjCAMERA_FREE - # Pause simulation - elif key == glfw.KEY_SPACE and self._paused is not None: - self._paused = not self._paused - # Advances simulation by one step. - elif key == glfw.KEY_RIGHT and self._paused is not None: - self._advance_by_one_step = True - self._paused = True - # Slows down simulation - elif key == glfw.KEY_S: - self._run_speed /= 2.0 - # Speeds up simulation - elif key == glfw.KEY_F: - self._run_speed *= 2.0 - # Turn off / turn on rendering every frame. - elif key == glfw.KEY_D: - self._render_every_frame = not self._render_every_frame - # Capture screenshot - elif key == glfw.KEY_T: - img = np.zeros( - (glfw.get_framebuffer_size( - self.window)[1], glfw.get_framebuffer_size( - self.window)[0], 3), dtype=np.uint8) - mujoco.mjr_readPixels(img, None, self.viewport, self.ctx) - imageio.imwrite(self._image_path % self._image_idx, np.flipud(img)) - self._image_idx += 1 - # Display contact forces - elif key == glfw.KEY_C: - self._contacts = not self._contacts - self.vopt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = self._contacts - self.vopt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTFORCE] = self._contacts - elif key == glfw.KEY_J: - self._joints = not self._joints - self.vopt.flags[mujoco.mjtVisFlag.mjVIS_JOINT] = self._joints - # Display coordinate frames - elif key == glfw.KEY_E: - self.vopt.frame = 1 - self.vopt.frame - # Hide overlay menu - elif key == glfw.KEY_LEFT_ALT: - self._hide_menus = True - elif key == glfw.KEY_H: - self._hide_menus = not self._hide_menus - # Make transparent - elif key == glfw.KEY_R: - self._transparent = not self._transparent - if self._transparent: - self.model.geom_rgba[:, 3] /= 5.0 +class MujocoViewer: + #=========================# + # P L A C E H O L D E R # + #=========================# + @dataclass + class _MOUSE_DATA: + button_left_pressed = False + button_right_pressed = False + left_double_click_pressed = False + right_double_click_pressed = False + last_left_click_time = None + last_right_click_time = None + last_mouse_click_x = 0 + last_mouse_click_y = 0 + dx = 0 + dy = 0 + last_x = 0 + last_y = 0 + y_offset_transient = 0 + button_released = False + @dataclass + class _VIEWER_CONFIG: + title: str = "DEFAULT" + window_size: Optional[Tuple[int, int]] = None + image_output_format: str = "output/frame_%07d.png" + # FPS: float = 60.0 [NOT IMPLEMENTED] + simulation_run_speed: float = 1.0 # [UNUSED] + sensor_config = dict() + cam_id: int = -1 + cam_type = mujoco.mjtCamera.mjCAMERA_FREE + N_CAM: int = 0 + _viewer_status: _VIEWER_STATUS_MAP = _VIEWER_STATUS_MAP.DEFAULT + + def index_camera_unsafe(self): + self.cam_id += 1 + if self.cam_id >= self.N_CAM: + self.cam_id = -1 + self.cam_type = mujoco.mjtCamera.mjCAMERA_FREE else: - self.model.geom_rgba[:, 3] *= 5.0 - # Display inertia - elif key == glfw.KEY_I: - self._inertias = not self._inertias - self.vopt.flags[mujoco.mjtVisFlag.mjVIS_INERTIA] = self._inertias - # Display center of mass - elif key == glfw.KEY_M: - self._com = not self._com - self.vopt.flags[mujoco.mjtVisFlag.mjVIS_COM] = self._com - # Shadow Rendering - elif key == glfw.KEY_O: - self._shadows = not self._shadows - self.scn.flags[mujoco.mjtRndFlag.mjRND_SHADOW] = self._shadows - # Convex-Hull rendering - elif key == glfw.KEY_V: - self._convex_hull_rendering = not self._convex_hull_rendering - self.vopt.flags[ - mujoco.mjtVisFlag.mjVIS_CONVEXHULL - ] = self._convex_hull_rendering - # Wireframe Rendering - elif key == glfw.KEY_W: - self._wire_frame = not self._wire_frame - self.scn.flags[mujoco.mjtRndFlag.mjRND_WIREFRAME] = self._wire_frame - # Geom group visibility - elif key in (glfw.KEY_0, glfw.KEY_1, glfw.KEY_2, glfw.KEY_3, glfw.KEY_4): - self.vopt.geomgroup[key - glfw.KEY_0] ^= 1 - # Quit - if key == glfw.KEY_ESCAPE: - print("Pressed ESC") - print("Quitting.") - glfw.set_window_should_close(self.window, True) - return - - def _cursor_pos_callback(self, window, xpos, ypos): - if not (self._button_left_pressed or self._button_right_pressed): - return - - mod_shift = ( - glfw.get_key(window, glfw.KEY_LEFT_SHIFT) == glfw.PRESS or - glfw.get_key(window, glfw.KEY_RIGHT_SHIFT) == glfw.PRESS) - if self._button_right_pressed: - action = mujoco.mjtMouse.mjMOUSE_MOVE_H if mod_shift else mujoco.mjtMouse.mjMOUSE_MOVE_V - elif self._button_left_pressed: - action = mujoco.mjtMouse.mjMOUSE_ROTATE_H if mod_shift else mujoco.mjtMouse.mjMOUSE_ROTATE_V - else: - action = mujoco.mjtMouse.mjMOUSE_ZOOM - - dx = int(self._scale * xpos) - self._last_mouse_x - dy = int(self._scale * ypos) - self._last_mouse_y - width, height = glfw.get_framebuffer_size(window) + self.cam_type = mujoco.mjtCamera.mjCAMERA_FIXED + + def scale_simulation_run_speed_unsafe(self, factor): + self.simulation_run_speed *= factor + + def update_viewer_status_unsafe(self, status, if_enable): + if if_enable: + self._viewer_status |= status + else: + self._viewer_status &= ~status + + def toggle_viewer_status_unsafe(self, status): + self._viewer_status ^= status + + def set_viewer_status_unsafe(self, status): + self._viewer_status |= status + + def clear_viewer_status_unsafe(self, status): + self._viewer_status &= ~status + + def get_viewer_status_unsafe(self, status): + return bool(self._viewer_status & status) + @dataclass + class _MJ_DATA: + # cache: + model = None + data = None + # mujoco placeholder: + config:_MJ_CONFIG = _MJ_CONFIG() + camera_data: _CAMERA_DATA = _CAMERA_DATA() + + @dataclass + class _GUI_DATA: + # glfw placeholder: + glfw_window = None + frame_buffer_size = None + frame_window_size = None + frame_scale = None + viewer_overlay = None + viewer_markers = None + frame_rendering_time = 1 + frame_stamp = 0 + + #===============================# + # I N I T I A L I Z A T I O N # + #===============================# + # internal: + _mouse_data_lock = Lock() + _mouse_data: _MOUSE_DATA = _MOUSE_DATA() + + _viewer_config_lock = Lock() + _viewer_config: _VIEWER_CONFIG = _VIEWER_CONFIG() + + _mj_lock = Lock() + _mj: _MJ_DATA = _MJ_DATA() + + _gui_lock = Lock() + _gui_data: _GUI_DATA = _GUI_DATA() + + _camera_data_lock = Lock() + _camera_data: Dict[str, _CAMERA_DATA] = {} + + _camera_mj_config: _MJ_CONFIG = _MJ_CONFIG() # config for camera rendering + + def __init__(self, + mj_model, + mj_data, + sensor_config = None, + title: str = "unnamed", + image_output_format: str = "output/frame_%07d.png", + window_size: Optional[Tuple[int, int]] = None, + if_hide_menus: bool = True, + if_render_on_screen: bool = True, + if_save_scrnshot_on_fly: bool = True, + ): + # update config: + if sensor_config: + self._viewer_config.sensor_config = sensor_config + self._viewer_config.title = title + self._viewer_config.window_size = window_size + self._viewer_config.update_viewer_status_unsafe(_VIEWER_STATUS_MAP.HIDE_MENUS, if_hide_menus) + self._viewer_config.update_viewer_status_unsafe(_VIEWER_STATUS_MAP.RENDER_ON_SCREEN, if_render_on_screen) + self._viewer_config.N_CAM = mj_model.ncam + self._viewer_config.update_viewer_status_unsafe(_VIEWER_STATUS_MAP.SAVE_SCRNSHOT_ON_FLY, if_save_scrnshot_on_fly) + self._viewer_config.image_output_format = image_output_format + + # cache mj objects: + self._mj.model = mj_model + self._mj.data = mj_data + + ## module initialization sequence: + self._init_glfw_safe() # depends on config + self._init_Mujoco_safe() # depends on glfw + + def _init_Mujoco_safe(self): + # get viewport + with self._viewer_config_lock: + ww, wh = self._viewer_config.window_size + + # create options, camera, scene, context + with self._mj_lock: + self._mj.config.scn = mujoco.MjvScene(self._mj.model, maxgeom=10000) + self._mj.config.ctx = mujoco.MjrContext(self._mj.model, mujoco.mjtFontScale.mjFONTSCALE_150.value) + self._mj.config.vopt = mujoco.MjvOption() + self._mj.config.pert = mujoco.MjvPerturb() + # camera buffer: + self._mj.camera_data.mj_cam = mujoco.MjvCamera() + self._mj.camera_data.mj_viewport = mujoco.MjrRect(0, 0, ww, wh) + self._mj.camera_data.depth_buffer = None # disabled + self._mj.camera_data.frame_buffer = np.zeros((wh, ww, 3), dtype=np.uint8) + self._mj.camera_data.frame_stamp = self._gui_data.frame_stamp + + # initialize overlay, markers + with self._gui_lock: + self._gui_data.viewer_overlay = {} + self._gui_data.viewer_markers = [] + + # create camera sensor buffers: + with self._camera_data_lock: + self._camera_mj_config.vopt = mujoco.MjvOption() + self._camera_mj_config.pert = mujoco.MjvPerturb() + self._camera_mj_config.scn = mujoco.MjvScene(self._mj.model, maxgeom=10000) + self._camera_mj_config.ctx = mujoco.MjrContext(self._mj.model, mujoco.mjtFontScale.mjFONTSCALE_150.value) + # - default config: + # --> (disable group0: collision meshes) + self._camera_mj_config.vopt.geomgroup[0] = False + self._camera_mj_config.vopt.geomgroup[1] = True + # --> (disable ) + + for camera_node, camera in self._viewer_config.sensor_config.items(): + w = camera["width"] + h = camera["height"] + cam_id = camera["id"] + # Create Buffer: + with self._camera_data_lock: + self._camera_data[camera_node] = _CAMERA_DATA() + self._camera_data[camera_node].mj_viewport = mujoco.MjrRect(0, 0, w, h) + self._camera_data[camera_node].frame_buffer = np.zeros((h, w, 3), dtype=np.uint8) + self._camera_data[camera_node].depth_buffer = np.zeros((h, w), dtype=np.float32) + self._camera_data[camera_node].frame_stamp = self._gui_data.frame_stamp + self._camera_data[camera_node].mj_cam = mujoco.MjvCamera() + self._camera_data[camera_node].mj_cam.type = mujoco.mjtCamera.mjCAMERA_FIXED + self._camera_data[camera_node].mj_cam.fixedcamid = cam_id + + def _init_glfw_safe(self): + # - fetch config: + with self._viewer_config_lock: + window_size = self._viewer_config.window_size + title = self._viewer_config.title + if_on_scrn = self._viewer_config.get_viewer_status_unsafe(_VIEWER_STATUS_MAP.RENDER_ON_SCREEN) + + # - init: + glfw.init() + + # - Create Window: + if not window_size: # auto-window + window_size = glfw.get_video_mode(glfw.get_primary_monitor()).size + with self._viewer_config_lock: + self._viewer_config.window_size = window_size + ww, wh = window_size + + glfw_window = glfw.create_window(ww, wh, "[{}] Main Window".format(title), None, None) + + glfw.window_hint(glfw.VISIBLE, if_on_scrn) + glfw.make_context_current(glfw_window) + glfw.swap_interval(1) + frame_buffer_size = glfw.get_framebuffer_size(glfw_window) + frame_window_size = glfw.get_window_size(glfw_window) + frame_scale = frame_buffer_size[0] * 1.0 / frame_window_size[0] + + # - cache: + with self._gui_lock: + self._gui_data.glfw_window = glfw_window + self._gui_data.frame_buffer_size = frame_buffer_size + self._gui_data.frame_window_size = frame_window_size + self._gui_data.frame_scale = frame_scale + + # - initialize interactive on-screen modules: + if if_on_scrn: + # - set callbacks + glfw.set_cursor_pos_callback (glfw_window, self._cursor_pos_callback ) + glfw.set_mouse_button_callback (glfw_window, self._mouse_button_callback ) + glfw.set_scroll_callback (glfw_window, self._scroll_callback ) + glfw.set_key_callback (glfw_window, self._key_callback ) + glfw.set_framebuffer_size_callback (glfw_window, self._framebuffer_size_callback ) + + with self._viewer_config_lock: + self._viewer_config.set_viewer_status_unsafe(_VIEWER_STATUS_MAP.IS_ALIVE) + + #===================# + # C A L L B A C K # + #===================# + def _framebuffer_size_callback(self, window, width, height): + frame_buffer_size = glfw.get_framebuffer_size(window) + frame_window_size = glfw.get_window_size(window) + frame_scale = frame_buffer_size[0] * 1.0 / frame_window_size[0] + # - cache: with self._gui_lock: - if self.pert.active: - mujoco.mjv_movePerturb( - self.model, - self.data, - action, - dx / height, - dy / height, - self.scn, - self.pert) + self._gui_data.frame_buffer_size = frame_buffer_size + self._gui_data.frame_window_size = frame_window_size + self._gui_data.frame_scale = frame_scale + + # === CALLBACK MAP === # + _key_stroke_callback_map = { + glfw.KEY_TAB : lambda config: config.index_camera_unsafe(), + glfw.KEY_SPACE : lambda config: config.toggle_viewer_status_unsafe (_VIEWER_STATUS_MAP.PAUSED), + glfw.KEY_RIGHT : lambda config: config.set_viewer_status_unsafe (_VIEWER_STATUS_MAP.PAUSED | _VIEWER_STATUS_MAP.STEP_BY_STEP), + glfw.KEY_S : lambda config: config.scale_simulation_run_speed_unsafe(0.5), # slow down + glfw.KEY_F : lambda config: config.scale_simulation_run_speed_unsafe(2.0), # faster + glfw.KEY_D : lambda config: config.toggle_viewer_status_unsafe (_VIEWER_STATUS_MAP.RENDER_EVERY_FRAME), + glfw.KEY_P : lambda config: config.set_viewer_status_unsafe (_VIEWER_STATUS_MAP.CAPTURE_SCRNSHOT), + glfw.KEY_C : lambda config: config.toggle_viewer_status_unsafe (_VIEWER_STATUS_MAP.CONTACTS), + glfw.KEY_J : lambda config: config.toggle_viewer_status_unsafe (_VIEWER_STATUS_MAP.JOINTS), + glfw.KEY_R : lambda config: config.toggle_viewer_status_unsafe (_VIEWER_STATUS_MAP.TOGGLE_COORD_FRAME), + glfw.KEY_H : lambda config: config.toggle_viewer_status_unsafe (_VIEWER_STATUS_MAP.HIDE_MENUS), + glfw.KEY_T : lambda config: config.toggle_viewer_status_unsafe (_VIEWER_STATUS_MAP.TRANSPARENT), + glfw.KEY_I : lambda config: config.toggle_viewer_status_unsafe (_VIEWER_STATUS_MAP.INERTIAS), + glfw.KEY_M : lambda config: config.toggle_viewer_status_unsafe (_VIEWER_STATUS_MAP.COM), + glfw.KEY_O : lambda config: config.toggle_viewer_status_unsafe (_VIEWER_STATUS_MAP.SHADOWS), + glfw.KEY_V : lambda config: config.toggle_viewer_status_unsafe (_VIEWER_STATUS_MAP.CONVEX_HULL_RENDERING), + glfw.KEY_W : lambda config: config.toggle_viewer_status_unsafe (_VIEWER_STATUS_MAP.WIRE_FRAME), + glfw.KEY_ESCAPE : lambda config: config.toggle_viewer_status_unsafe (_VIEWER_STATUS_MAP.TERMINATE), + glfw.KEY_LEFT_SHIFT : lambda config: config.clear_viewer_status_unsafe (_VIEWER_STATUS_MAP.LEFT_SHIFT_PRESSED), # release + glfw.KEY_LEFT_CONTROL: lambda config: config.clear_viewer_status_unsafe (_VIEWER_STATUS_MAP.LEFT_CONTROL_PRESSED), # release + # Number keys + glfw.KEY_0 : lambda config: config.toggle_viewer_status_unsafe (_VIEWER_STATUS_MAP.GROUP_0), + glfw.KEY_1 : lambda config: config.toggle_viewer_status_unsafe (_VIEWER_STATUS_MAP.GROUP_1), + glfw.KEY_2 : lambda config: config.toggle_viewer_status_unsafe (_VIEWER_STATUS_MAP.GROUP_2), + glfw.KEY_3 : lambda config: config.toggle_viewer_status_unsafe (_VIEWER_STATUS_MAP.GROUP_3), + } + def _key_callback(self, window, key, scancode, action, mods): + + if action == glfw.RELEASE: + # register key upon release + if key in self._key_stroke_callback_map: + with self._viewer_config_lock: + self._key_stroke_callback_map[key](self._viewer_config) else: - mujoco.mjv_moveCamera( - self.model, - action, - dx / height, - dy / height, - self.scn, - self.cam) + pass # unregistered key press + else: + if key == glfw.KEY_LEFT_SHIFT: + with self._viewer_config_lock: + self._viewer_config.set_viewer_status_unsafe(_VIEWER_STATUS_MAP.LEFT_SHIFT_PRESSED) + elif key == glfw.KEY_LEFT_CONTROL: + with self._viewer_config_lock: + self._viewer_config.set_viewer_status_unsafe(_VIEWER_STATUS_MAP.LEFT_CONTROL_PRESSED) + return - self._last_mouse_x = int(self._scale * xpos) - self._last_mouse_y = int(self._scale * ypos) + def _cursor_pos_callback(self, window, mouse_x, mouse_y): + with self._mouse_data_lock: + self._mouse_data.dx = mouse_x - self._mouse_data.last_x + self._mouse_data.dy = mouse_y - self._mouse_data.last_y + self._mouse_data.last_x = mouse_x + self._mouse_data.last_y = mouse_y def _mouse_button_callback(self, window, button, act, mods): - self._button_left_pressed = button == glfw.MOUSE_BUTTON_LEFT and act == glfw.PRESS - self._button_right_pressed = button == glfw.MOUSE_BUTTON_RIGHT and act == glfw.PRESS - - x, y = glfw.get_cursor_pos(window) - self._last_mouse_x = int(self._scale * x) - self._last_mouse_y = int(self._scale * y) - - # detect a left- or right- doubleclick - self._left_double_click_pressed = False - self._right_double_click_pressed = False + # Local: + left_double_click_pressed = False + right_double_click_pressed = False + last_left_click_time = None + last_right_click_time = None + + # read cache: + with self._mouse_data_lock: + last_left_click_time = self._mouse_data.last_left_click_time + last_right_click_time = self._mouse_data.last_right_click_time + + # process: time_now = glfw.get_time() - - if self._button_left_pressed: - if self._last_left_click_time is None: - self._last_left_click_time = glfw.get_time() - - time_diff = (time_now - self._last_left_click_time) - if time_diff > 0.01 and time_diff < 0.3: - self._left_double_click_pressed = True - self._last_left_click_time = time_now - - if self._button_right_pressed: - if self._last_right_click_time is None: - self._last_right_click_time = glfw.get_time() - - time_diff = (time_now - self._last_right_click_time) - if time_diff > 0.01 and time_diff < 0.2: - self._right_double_click_pressed = True - self._last_right_click_time = time_now - - # set perturbation - key = mods == glfw.MOD_CONTROL - newperturb = 0 - if key and self.pert.select > 0: + button_left_pressed = button == glfw.MOUSE_BUTTON_LEFT and act == glfw.PRESS + button_right_pressed = button == glfw.MOUSE_BUTTON_RIGHT and act == glfw.PRESS + + if button_left_pressed: + if last_left_click_time: + delta_t = time_now - last_left_click_time + if 0.01 < delta_t < 0.3: + left_double_click_pressed = True + last_left_click_time = time_now + if button_right_pressed: + if last_right_click_time: + delta_t = time_now - last_right_click_time + if 0.01 < delta_t < 0.3: + right_double_click_pressed = True + last_right_click_time = time_now + + x, y = glfw.get_cursor_pos(window) + # cache: + with self._mouse_data_lock: + self._mouse_data.last_mouse_click_x = x + self._mouse_data.last_mouse_click_y = y + self._mouse_data.button_left_pressed = button_left_pressed + self._mouse_data.button_right_pressed = button_right_pressed + self._mouse_data.left_double_click_pressed = left_double_click_pressed + self._mouse_data.right_double_click_pressed = right_double_click_pressed + self._mouse_data.last_left_click_time = last_left_click_time + self._mouse_data.last_right_click_time = last_right_click_time + self._mouse_data.button_released = (act == glfw.RELEASE) + + def _scroll_callback(self, window, x_offset, y_offset_transient): + with self._mouse_data_lock: + self._mouse_data.y_offset_transient = y_offset_transient + + #====================================# + # P R I V A T E F U N C T I O N # + #====================================# + def _on_terminate_safe(self): + print("Terminating mujoco_viewer!") + glfw.set_window_should_close(self._gui_data.glfw_window, True) + glfw.terminate() + with self._mj_lock: + self._mj.config.ctx.free() + with self._viewer_config_lock: + self._viewer_config.clear_viewer_status_unsafe(_VIEWER_STATUS_MAP.IS_ALIVE) + print("mujoco_viewer has been Terminated!") + sys.exit() + + def _process_viewer_config(self): + # - Fetch Inputs: + with self._viewer_config_lock: + current_viewer_status = self._viewer_config._viewer_status + current_cam_id = self._viewer_config.cam_id + current_cam_type = self._viewer_config.cam_type + with self._gui_lock: + frame_buffer_size = self._gui_data.frame_buffer_size + + # - Update MJ: + with self._mj_lock: + ### change cam id view: + self._mj.camera_data.mj_cam.fixedcamid = current_cam_id + self._mj.camera_data.mj_cam.type = current_cam_type + ### coord frame toggle: + if bool(current_viewer_status & _VIEWER_STATUS_MAP.TOGGLE_COORD_FRAME): + self._mj.config.vopt.frame = 1 - self._mj.config.vopt.frame + current_viewer_status ^= ~(_VIEWER_STATUS_MAP.TOGGLE_COORD_FRAME) # clear status + ### transparency: + if bool(current_viewer_status & _VIEWER_STATUS_MAP.TRANSPARENT): + self._mj.model.geom_rgba[:, 3] = 0.5 + else: + self._mj.model.geom_rgba[:, 3] = 1.0 + ### visual features: + self._mj.config.vopt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = bool(current_viewer_status & _VIEWER_STATUS_MAP.CONTACTS) + self._mj.config.vopt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTFORCE] = bool(current_viewer_status & _VIEWER_STATUS_MAP.CONTACTS) + self._mj.config.vopt.flags[mujoco.mjtVisFlag.mjVIS_JOINT] = bool(current_viewer_status & _VIEWER_STATUS_MAP.JOINTS) + self._mj.config.vopt.flags[mujoco.mjtVisFlag.mjVIS_COM] = bool(current_viewer_status & _VIEWER_STATUS_MAP.INERTIAS) + self._mj.config.vopt.flags[mujoco.mjtVisFlag.mjVIS_INERTIA] = bool(current_viewer_status & _VIEWER_STATUS_MAP.COM) + self._mj.config.vopt.flags[mujoco.mjtVisFlag.mjVIS_CONVEXHULL] = bool(current_viewer_status & _VIEWER_STATUS_MAP.CONVEX_HULL_RENDERING) + ### OpenGL rendering effects. + self._mj.config.scn.flags[mujoco.mjtRndFlag.mjRND_WIREFRAME] = bool(current_viewer_status & _VIEWER_STATUS_MAP.WIRE_FRAME) + self._mj.config.scn.flags[mujoco.mjtRndFlag.mjRND_SHADOW] = bool(current_viewer_status & _VIEWER_STATUS_MAP.SHADOWS) + ### Geom Group: + self._mj.config.vopt.geomgroup[0] = bool(current_viewer_status & _VIEWER_STATUS_MAP.GROUP_0) + self._mj.config.vopt.geomgroup[1] = bool(current_viewer_status & _VIEWER_STATUS_MAP.GROUP_1) + self._mj.config.vopt.geomgroup[2] = bool(current_viewer_status & _VIEWER_STATUS_MAP.GROUP_2) + self._mj.config.vopt.geomgroup[3] = bool(current_viewer_status & _VIEWER_STATUS_MAP.GROUP_3) + ### viewport dimension: + self._mj.camera_data.mj_viewport.width, self._mj.camera_data.mj_viewport.height = frame_buffer_size + + def _process_mouse_interactions(self): + """ Process mouse interactions: + - selection + - dragging interaction + - camera angles and zoom + + TODO: double check the perturbation + """ + # - local cache safely: + with self._mouse_data_lock: + current_mouse_data = copy.deepcopy(self._mouse_data) + with self._viewer_config_lock: + current_viewer_status = copy.deepcopy(self._viewer_config._viewer_status) + with self._gui_lock: + width, height = self._gui_data.frame_buffer_size + scale = self._gui_data.frame_scale + + ### Selection ### + # - mouse interaction: + if_mod_control = bool(current_viewer_status & _VIEWER_STATUS_MAP.LEFT_CONTROL_PRESSED) + + # - perturbation: + new_perturb = 0 + if if_mod_control and self._mj.config.pert.select > 0: # right: translate, left: rotate - if self._button_right_pressed: - newperturb = mujoco.mjtPertBit.mjPERT_TRANSLATE - if self._button_left_pressed: - newperturb = mujoco.mjtPertBit.mjPERT_ROTATE - - # perturbation onste: reset reference - if newperturb and not self.pert.active: - mujoco.mjv_initPerturb( - self.model, self.data, self.scn, self.pert) - self.pert.active = newperturb - - # handle doubleclick - if self._left_double_click_pressed or self._right_double_click_pressed: - # determine selection mode - selmode = 0 - if self._left_double_click_pressed: - selmode = 1 - if self._right_double_click_pressed: - selmode = 2 - if self._right_double_click_pressed and key: - selmode = 3 - + if current_mouse_data.button_right_pressed: + new_perturb = mujoco.mjtPertBit.mjPERT_TRANSLATE + if current_mouse_data.button_left_pressed: + new_perturb = mujoco.mjtPertBit.mjPERT_ROTATE + # perturbation: reset reference + if new_perturb and not self._mj.config.pert.active: + mujoco.mjv_initPerturb(self._mj.model, self._mj.data, self._mj.config.scn, self._mj.config.pert) + self._mj.config.pert.active = new_perturb + + # - interaction: + selection_mode = 0 + if current_mouse_data.left_double_click_pressed: + selection_mode = 1 + elif current_mouse_data.right_double_click_pressed: + selection_mode = 3 if if_mod_control else 2 + + if selection_mode: + with self._mj_lock: + vp_w, vp_h = self._mj.camera_data.mj_viewport.width, self._mj.camera_data.mj_viewport.height # find geom and 3D click point, get corresponding body - width, height = self.viewport.width, self.viewport.height - aspectratio = width / height - relx = x / width - rely = (self.viewport.height - y) / height - selpnt = np.zeros((3, 1), dtype=np.float64) - selgeom = np.zeros((1, 1), dtype=np.int32) - selskin = np.zeros((1, 1), dtype=np.int32) - selbody = mujoco.mjv_select( - self.model, - self.data, - self.vopt, - aspectratio, - relx, - rely, - self.scn, - selpnt, - selgeom, - selskin) - - # set lookat point, start tracking is requested - if selmode == 2 or selmode == 3: - # set cam lookat - if selbody >= 0: - self.cam.lookat = selpnt.flatten() - # switch to tracking camera if dynamic body clicked - if selmode == 3 and selbody > 0: - self.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING - self.cam.trackbodyid = selbody - self.cam.fixedcamid = -1 - # set body selection - else: - if selbody >= 0: - # record selection - self.pert.select = selbody - self.pert.skinselect = selskin - # compute localpos - vec = selpnt.flatten() - self.data.xpos[selbody] - mat = self.data.xmat[selbody].reshape(3, 3) - self.pert.localpos = self.data.xmat[selbody].reshape( - 3, 3).dot(vec) + aspect_ratio = vp_w / vp_h + rel_x = int(scale * (current_mouse_data.last_mouse_click_x)) / vp_w + rel_y = (vp_h - int(scale * (current_mouse_data.last_mouse_click_y))) / vp_h + selected_pnt = np.zeros((3, 1), dtype=np.float64) + selected_geom = np.zeros((1, 1), dtype=np.int32) + selected_skin = np.zeros((1, 1), dtype=np.int32) + # apply selection: + with self._mj_lock: + selected_body = mujoco.mjv_select( + self._mj.model, + self._mj.data, + self._mj.config.vopt, + aspect_ratio, + rel_x, + rel_y, + self._mj.config.scn, + selected_pnt, + selected_geom, + selected_skin + ) + + # set lookat point, start tracking is requested + if current_mouse_data.right_double_click_pressed: + # set cam lookat + if selected_body >= 0: + self._mj.camera_data.mj_cam.lookat = selected_pnt.flatten() + # switch to tracking camera if dynamic body clicked + # if if_mod_control and selected_body > 0: + # self._viewer_config.cam_id = -1 + # self._viewer_config.cam_type = mujoco.mjtCamera.mjCAMERA_TRACKING + # self._mj.camera_data.mj_cam.trackbodyid = selected_body + # set body selection else: - self.pert.select = 0 - self.pert.skinselect = -1 - # stop perturbation on select - self.pert.active = 0 - - # 3D release - if act == glfw.RELEASE: - self.pert.active = 0 - - def _scroll_callback(self, window, x_offset, y_offset): - with self._gui_lock: + if selected_body >= 0: + # record selection + self._mj.config.pert.select = selected_body + self._mj.config.pert.skinselect = selected_skin + # compute localpos + vec = selected_pnt.flatten() - self._mj.data.xpos[selected_body] + mat = self._mj.data.xmat[selected_body].reshape(3, 3) + self._mj.config.pert.localpos = mat.dot(vec) + else: + self._mj.config.pert.select = 0 + self._mj.config.pert.skinselect = -1 + # stop perturbation on select + self._mj.config.pert.active = 0 + + if current_mouse_data.button_released: + # 3D release + with self._mj_lock: + self._mj.config.pert.active = 0 + + ### Camera ### + if current_mouse_data.button_left_pressed or current_mouse_data.button_right_pressed: + # - process mouse interaction: + if_mod = bool(current_viewer_status & _VIEWER_STATUS_MAP.LEFT_SHIFT_PRESSED) + if current_mouse_data.button_left_pressed: + mouse_action = mujoco.mjtMouse.mjMOUSE_ROTATE_H if if_mod else mujoco.mjtMouse.mjMOUSE_ROTATE_V + elif current_mouse_data.button_right_pressed: + mouse_action = mujoco.mjtMouse.mjMOUSE_MOVE_H if if_mod else mujoco.mjtMouse.mjMOUSE_MOVE_V + + dx = current_mouse_data.dx + dy = current_mouse_data.dy + dx_pix = int(scale * dx) + dy_pix = int(scale * dy) + + # - update: + with self._mj_lock: + if self._mj.config.pert.active: + mujoco.mjv_movePerturb( + self._mj.model, self._mj.data, + mouse_action, + dx_pix / height, + dy_pix / height, + self._mj.config.scn, self._mj.config.pert) + else: + mujoco.mjv_moveCamera( + self._mj.model, + mouse_action, + dx_pix / height, + dy_pix / height, + self._mj.config.scn, self._mj.camera_data.mj_cam) + # - zoom: + with self._mj_lock: mujoco.mjv_moveCamera( - self.model, mujoco.mjtMouse.mjMOUSE_ZOOM, 0, -0.05 * y_offset, self.scn, self.cam) - - def add_marker(self, **marker_params): - self._markers.append(marker_params) - - def _add_marker_to_scene(self, marker): - if self.scn.ngeom >= self.scn.maxgeom: - raise RuntimeError( - 'Ran out of geoms. maxgeom: %d' % - self.scn.maxgeom) - - g = self.scn.geoms[self.scn.ngeom] + self._mj.model, + mujoco.mjtMouse.mjMOUSE_ZOOM, + 0, -0.05 * current_mouse_data.y_offset_transient, + self._mj.config.scn, self._mj.camera_data.mj_cam) + + with self._mouse_data_lock: + self._mouse_data.y_offset_transient = 0 + + @staticmethod + def _add_marker_to_scene(marker, scn): + if scn.ngeom >= scn.maxgeom: + raise RuntimeError('Ran out of geoms. maxgeom: %d' % scn.maxgeom) + + g = scn.geoms[scn.ngeom] # default values. g.dataid = -1 g.objtype = mujoco.mjtObj.mjOBJ_UNKNOWN @@ -393,223 +676,312 @@ def _add_marker_to_scene(self, marker): key, type(value))) else: raise ValueError("mjtGeom doesn't have field %s" % key) - - self.scn.ngeom += 1 - + scn.ngeom += 1 return def _create_overlay(self): - topleft = mujoco.mjtGridPos.mjGRID_TOPLEFT - topright = mujoco.mjtGridPos.mjGRID_TOPRIGHT - bottomleft = mujoco.mjtGridPos.mjGRID_BOTTOMLEFT - bottomright = mujoco.mjtGridPos.mjGRID_BOTTOMRIGHT - - def add_overlay(gridpos, text1, text2): - if gridpos not in self._overlay: - self._overlay[gridpos] = ["", ""] - self._overlay[gridpos][0] += text1 + "\n" - self._overlay[gridpos][1] += text2 + "\n" - - if self._render_every_frame: - add_overlay(topleft, "", "") - else: - add_overlay( - topleft, - "Run speed = %.3f x real time" % - self._run_speed, - "[S]lower, [F]aster") - add_overlay( - topleft, - "Ren[d]er every frame", - "On" if self._render_every_frame else "Off") - add_overlay( - topleft, "Switch camera (#cams = %d)" % - (self.model.ncam + 1), "[Tab] (camera ID = %d)" % - self.cam.fixedcamid) - add_overlay( - topleft, - "[C]ontact forces", - "On" if self._contacts else "Off") - add_overlay( - topleft, - "[J]oints", - "On" if self._joints else "Off") - add_overlay( - topleft, - "[I]nertia", - "On" if self._inertias else "Off") - add_overlay( - topleft, - "Center of [M]ass", - "On" if self._com else "Off") - add_overlay( - topleft, "Shad[O]ws", "On" if self._shadows else "Off" - ) - add_overlay( - topleft, - "T[r]ansparent", - "On" if self._transparent else "Off") - add_overlay( - topleft, - "[W]ireframe", - "On" if self._wire_frame else "Off") - add_overlay( - topleft, - "Con[V]ex Hull Rendering", - "On" if self._convex_hull_rendering else "Off", + with self._viewer_config_lock: + v_status = copy.deepcopy(self._viewer_config._viewer_status) + run_spd = self._viewer_config.simulation_run_speed + N_cam = self._viewer_config.N_CAM + image_output_format = self._viewer_config.image_output_format + with self._mj_lock: + camid = self._mj.camera_data.mj_cam.fixedcamid + mj_timestep = self._mj.model.opt.timestep + mj_data_time = self._mj.data.time + mj_data_solviter = self._mj.data.solver_iter + with self._gui_lock: + frame_stamp = self._gui_data.frame_stamp + frame_rendering_time = self._gui_data.frame_rendering_time + + if_render_every_frame = bool(v_status & _VIEWER_STATUS_MAP.RENDER_EVERY_FRAME) + if_contacts = bool(v_status & _VIEWER_STATUS_MAP.CONTACTS) + if_joints = bool(v_status & _VIEWER_STATUS_MAP.JOINTS) + if_inertias = bool(v_status & _VIEWER_STATUS_MAP.INERTIAS) + if_com = bool(v_status & _VIEWER_STATUS_MAP.COM) + if_shadows = bool(v_status & _VIEWER_STATUS_MAP.SHADOWS) + if_transparent = bool(v_status & _VIEWER_STATUS_MAP.TRANSPARENT) + if_wire_frame = bool(v_status & _VIEWER_STATUS_MAP.WIRE_FRAME) + if_convex_hull_rendering = bool(v_status & _VIEWER_STATUS_MAP.CONVEX_HULL_RENDERING) + if_paused = bool(v_status & _VIEWER_STATUS_MAP.PAUSED) + if_toggle_coord = bool(v_status & _VIEWER_STATUS_MAP.TOGGLE_COORD_FRAME) + + _TOP_LEFT = mujoco.mjtGridPos.mjGRID_TOPLEFT + _BTM_LEFT = mujoco.mjtGridPos.mjGRID_BOTTOMLEFT + + group_avail = "{}{}{}{}".format( + "[0]" if (v_status & _VIEWER_STATUS_MAP.GROUP_0) else "", + "[1]" if (v_status & _VIEWER_STATUS_MAP.GROUP_1) else "", + "[2]" if (v_status & _VIEWER_STATUS_MAP.GROUP_2) else "", + "[3]" if (v_status & _VIEWER_STATUS_MAP.GROUP_3) else "", ) - if self._paused is not None: - if not self._paused: - add_overlay(topleft, "Stop", "[Space]") - else: - add_overlay(topleft, "Start", "[Space]") - add_overlay( - topleft, - "Advance simulation by one step", - "[right arrow]") - add_overlay( - topleft, - "Referenc[e] frames", - "On" if self.vopt.frame == 1 else "Off") - add_overlay(topleft, "[H]ide Menus", "") - if self._image_idx > 0: - fname = self._image_path % (self._image_idx - 1) - add_overlay(topleft, "Cap[t]ure frame", "Saved as %s" % fname) - else: - add_overlay(topleft, "Cap[t]ure frame", "") - add_overlay(topleft, "Toggle geomgroup visibility", "0-4") - - add_overlay( - bottomleft, "FPS", "%d%s" % - (1 / self._time_per_render, "")) - add_overlay( - bottomleft, "Solver iterations", str( - self.data.solver_iter + 1)) - add_overlay( - bottomleft, "Step", str( - round( - self.data.time / self.model.opt.timestep))) - add_overlay(bottomleft, "timestep", "%.5f" % self.model.opt.timestep) - - def apply_perturbations(self): - self.data.xfrc_applied = np.zeros_like(self.data.xfrc_applied) - mujoco.mjv_applyPerturbPose(self.model, self.data, self.pert, 0) - mujoco.mjv_applyPerturbForce(self.model, self.data, self.pert) - - def read_pixels(self, camid=None): - if self.render_mode is 'window': - raise NotImplementedError( - "Use 'render()' in 'window' mode.") - - if camid is not None: - if camid == -1: - self.cam.type = mujoco.mjtCamera.mjCAMERA_FREE - else: - self.cam.type = mujoco.mjtCamera.mjCAMERA_FIXED - self.cam.fixedcamid = camid - - self.viewport.width, self.viewport.height = glfw.get_framebuffer_size( - self.window) - # update scene + overlay = { + _TOP_LEFT: ["", ""], + _BTM_LEFT: ["", ""], + } + + _overlay_layout = { + _TOP_LEFT: { + "[S]lower, [F]aster" : "Run speed = %.3f x real time" % run_spd, + "Ren[d]er every frame" : "On" if if_render_every_frame else "Off", + "Switch camera (#cams = %d)" %(N_cam + 1) : "[Tab] (camera ID = %d)" % camid, + "[C]ontact forces" : "On" if if_contacts else "Off", + "[J]oints" : "On" if if_joints else "Off", + "[I]nertia" : "On" if if_inertias else "Off", + "Center of [M]ass" : "On" if if_com else "Off", + "Shad[O]ws" : "On" if if_shadows else "Off", + "[T]ransparent" : "On" if if_transparent else "Off", + "[W]ireframe" : "On" if if_wire_frame else "Off", + "Con[V]ex Hull Rendering" : "On" if if_convex_hull_rendering else "Off", + "Stop" if if_paused else "Start" : "[Space]", + "Step Through" : "->[right arrow]" if if_paused else "[Space] to stop first", + "[R]eference frames" : "On" if if_toggle_coord else "Off", + "[H]ide Menus" : "", + "Ca[p]ture frame" : ("path: " + image_output_format) % frame_stamp, + "Toggle geomgroup visibility" : group_avail, + }, + _BTM_LEFT: { + "FPS" : "%d" % (1/frame_rendering_time), + "Solver iterations" : str(mj_data_solviter + 1), + "Step" : str(round(mj_data_time / mj_timestep)), + "timestep" : "%.5f" % mj_timestep, + } + } + for grid_loc in _overlay_layout: + for feature, description in _overlay_layout[grid_loc].items(): + overlay[grid_loc][0] += feature + "\n" + overlay[grid_loc][1] += description + "\n" + + # override overlays: + with self._gui_lock: + self._gui_data.viewer_overlay = overlay + + #==================================# + # P U B L I C F U N C T I O N # + #==================================# + def signal_termination_safe(self, if_immediate:bool=False): + """ Signal the viewer to self-destroy upon next update/rendering (Or immediately if True) + """ + with self._viewer_config_lock: + self._viewer_config.set_viewer_status_unsafe(_VIEWER_STATUS_MAP.TERMINATE) + if if_immediate: + self._on_terminate_safe() + + def add_marker_safe(self, **marker_params): + with self._gui_lock: + self._gui_data.viewer_markers.append(marker_params) + + def apply_perturbations_safe(self): + """ Apply perturbation to mujoco: pose + force + """ + with self._mj_lock: + self._mj.data.xfrc_applied = np.zeros_like(self._mj.data.xfrc_applied) + mujoco.mjv_applyPerturbPose(self._mj.model, self._mj.data, self._mj.config.pert, 0) + mujoco.mjv_applyPerturbForce(self._mj.model, self._mj.data, self._mj.config.pert) + + def process_safe(self): + """ Process: safely process configs from the physical / software modifications + """ + # - Process user interactions: + self._process_viewer_config() + self._process_mouse_interactions() + return + + def update_safe(self): + # - generate overlays: + self._create_overlay() + + # - Update contexts: + ## Main Camera contents: + with self._viewer_config_lock: + is_on_scrn = self._viewer_config.get_viewer_status_unsafe(_VIEWER_STATUS_MAP.RENDER_ON_SCREEN) + + with self._mj_lock: + mj_timestep = self._mj.model.opt.timestep + mj_data_time = self._mj.data.time + mujoco.mjv_updateScene( - self.model, - self.data, - self.vopt, - self.pert, - self.cam, + self._mj.model, + self._mj.data, + self._mj.config.vopt, + self._mj.config.pert, + self._mj.camera_data.mj_cam, mujoco.mjtCatBit.mjCAT_ALL.value, - self.scn) - # render - mujoco.mjr_render(self.viewport, self.scn, self.ctx) - - img = np.zeros( - (glfw.get_framebuffer_size( - self.window)[1], glfw.get_framebuffer_size( - self.window)[0], 3), dtype=np.uint8) - mujoco.mjr_readPixels(img, None, self.viewport, self.ctx) - return np.flipud(img) - - def render(self): - if self.render_mode is 'offscreen': - raise NotImplementedError( - "Use 'read_pixels()' for 'offscreen' mode.") - if not self.is_alive: - raise Exception( - "GLFW window does not exist but you tried to render.") - if glfw.window_should_close(self.window): - self.close() + self._mj.config.scn + ) + + if is_on_scrn: + # marker items + with self._gui_lock: + for marker in self._gui_data.viewer_markers: + self._add_marker_to_scene(marker, self._mj.config.scn) + + + def render_safe(self): + """ Rendering: assume it has been udpated with `update_safe` + """ + render_start = time.time() + + with self._viewer_config_lock: + is_on_scrn = self._viewer_config.get_viewer_status_unsafe(_VIEWER_STATUS_MAP.RENDER_ON_SCREEN) + is_alive = self._viewer_config.get_viewer_status_unsafe(_VIEWER_STATUS_MAP.IS_ALIVE) + is_paused = self._viewer_config.get_viewer_status_unsafe(_VIEWER_STATUS_MAP.PAUSED) + if_capture_req = self._viewer_config.get_viewer_status_unsafe(_VIEWER_STATUS_MAP.CAPTURE_SCRNSHOT) + if_hide_menu = self._viewer_config.get_viewer_status_unsafe(_VIEWER_STATUS_MAP.HIDE_MENUS) + if_terminate = self._viewer_config.get_viewer_status_unsafe(_VIEWER_STATUS_MAP.TERMINATE) + + # signal exit on window close: + if_window_closed = False + with self._gui_lock: + if_window_closed = (glfw.window_should_close(self._gui_data.glfw_window)) + current_frame_stamp = self._gui_data.frame_stamp + + # process to determine if we shall exit(): + if if_terminate or if_window_closed: + self.signal_termination_safe(if_immediate=True) return - # mjv_updateScene, mjr_render, mjr_overlay - def update(): - # fill overlay items - self._create_overlay() - - render_start = time.time() - self.viewport.width, self.viewport.height = glfw.get_framebuffer_size( - self.window) - with self._gui_lock: - # update scene - mujoco.mjv_updateScene( - self.model, - self.data, - self.vopt, - self.pert, - self.cam, - mujoco.mjtCatBit.mjCAT_ALL.value, - self.scn) - # marker items - for marker in self._markers: - self._add_marker_to_scene(marker) - # render - mujoco.mjr_render(self.viewport, self.scn, self.ctx) + if not is_alive: + self.signal_termination_safe(if_immediate=True) + raise Exception("The GLFW window does not exist! Program is dead :( ---> Terminating ...") + + with self._gui_lock: + width, height = self._gui_data.frame_buffer_size + scale = self._gui_data.frame_scale + + # - init: + img_buffer = np.zeros((height, width, 3), dtype=np.uint8) + + # - update scene + render_start = time.time() + with self._mj_lock: + mj_timestep = self._mj.model.opt.timestep + mj_data_time = self._mj.data.time + + # render + mujoco.mjr_render(self._mj.camera_data.mj_viewport, self._mj.config.scn, self._mj.config.ctx) + + # read images: + if if_capture_req: + mujoco.mjr_readPixels(img_buffer, None, self._mj.camera_data.mj_viewport, self._mj.config.ctx) + + # on-scrn render: + if is_on_scrn: # overlay items - for gridpos, [t1, t2] in self._overlay.items(): - menu_positions = [mujoco.mjtGridPos.mjGRID_TOPLEFT, - mujoco.mjtGridPos.mjGRID_BOTTOMLEFT] - if gridpos in menu_positions and self._hide_menus: - continue - - mujoco.mjr_overlay( - mujoco.mjtFontScale.mjFONTSCALE_150, - gridpos, - self.viewport, - t1, - t2, - self.ctx) - glfw.swap_buffers(self.window) - glfw.poll_events() - self._time_per_render = 0.9 * self._time_per_render + \ - 0.1 * (time.time() - render_start) - - # clear overlay - self._overlay.clear() - - if self._paused: - while self._paused: - update() - if glfw.window_should_close(self.window): - self.close() - break - if self._advance_by_one_step: - self._advance_by_one_step = False - break - else: - self._loop_count += self.model.opt.timestep / \ - (self._time_per_render * self._run_speed) - if self._render_every_frame: - self._loop_count = 1 - while self._loop_count > 0: - update() - self._loop_count -= 1 - - # clear markers - self._markers[:] = [] + with self._gui_lock: + for gridpos, [t1, t2] in self._gui_data.viewer_overlay.items(): + # skip rendering for these locations if hidden: + if gridpos in [mujoco.mjtGridPos.mjGRID_TOPLEFT, mujoco.mjtGridPos.mjGRID_BOTTOMLEFT] \ + and if_hide_menu: + continue + # render overlay + mujoco.mjr_overlay( + mujoco.mjtFontScale.mjFONTSCALE_150, + gridpos, self._mj.camera_data.mj_viewport, + t1, t2, self._mj.config.ctx + ) + # render on screen: + with self._gui_lock: + glfw.swap_buffers(self._gui_data.glfw_window) + glfw.poll_events() + else: + print("TODO: Have not tested off creen mode") + + # clear gui data + with self._gui_lock: + self._gui_data.viewer_overlay.clear() + self._gui_data.viewer_markers.clear() + self._gui_data.frame_rendering_time = 0.9 * self._gui_data.frame_rendering_time + 0.1 * (time.time() - render_start) + self._gui_data.frame_stamp = int(round(mj_data_time / mj_timestep)) # apply perturbation (should this come before mj_step?) - self.apply_perturbations() - - def close(self): - self.is_alive = False - glfw.terminate() - self.ctx.free() + self.apply_perturbations_safe() + + if if_capture_req: + # cache buffer: + with self._mj_lock: + # img_buffer_flipped = np.flipud(img_buffer) + self._mj.camera_data.frame_buffer = img_buffer + self._mj.camera_data.frame_stamp = current_frame_stamp + # reset capture flag: + with self._viewer_config_lock: + if_capture_req = self._viewer_config.clear_viewer_status_unsafe(_VIEWER_STATUS_MAP.CAPTURE_SCRNSHOT) + # output: + self.save_last_available_captured_scrnshot_safe() + + return + + def render_sensor_cameras_safe(self, if_on_scrn_as_overlay=False, if_segmentation=False): + for camera_name, camera in self._camera_data.items(): + self._camera_mj_config.scn.flags[mujoco.mjtRndFlag.mjRND_WIREFRAME] = False + self._camera_mj_config.scn.flags[mujoco.mjtRndFlag.mjRND_SHADOW] = True + self._camera_mj_config.scn.flags[mujoco.mjtRndFlag.mjRND_SEGMENT] = if_segmentation + self._camera_mj_config.scn.flags[mujoco.mjtRndFlag.mjRND_IDCOLOR] = if_segmentation + # update scene + with self._mj_lock: + # change cam to the camera sensors + mujoco.mjv_updateScene( + self._mj.model, + self._mj.data, + self._camera_mj_config.vopt, + self._camera_mj_config.pert, + camera.mj_cam, + mujoco.mjtCatBit.mjCAT_ALL.value, + self._camera_mj_config.scn) + + # render off-screen + with self._camera_data_lock: + mujoco.mjr_render(camera.mj_viewport, self._camera_mj_config.scn, self._camera_mj_config.ctx) + + # cache: + mujoco.mjr_readPixels(camera.frame_buffer, camera.depth_buffer, camera.mj_viewport, self._camera_mj_config.ctx) + camera.frame_stamp = self._gui_data.frame_stamp + + + def acquire_sensor_camera_frames_safe(self): + camera_buffers = {"frame_buffer":{}, "depth_buffer":{}, "frame_stamp":{}} + with self._camera_data_lock: + for camera_name, camera in self._camera_data.items(): + camera_buffers["depth_buffer"][camera_name] = copy.deepcopy(camera.depth_buffer) + camera_buffers["frame_buffer"][camera_name] = copy.deepcopy(camera.frame_buffer) + camera_buffers["frame_stamp"][camera_name] = copy.deepcopy(camera.frame_stamp) + # if write_to: [NOT-USED: Implementation to save frames directly] + # imageio.imwrite( + # "{}/{}.png".format(write_to, camera_name.replace("\\", "_")), + # camera_buffers["frame_buffer"][camera_name] + # ) + # imageio.imwrite( + # "{}/{}_gray.png".format(write_to, camera_name.replace("\\", "_")), + # camera_buffers["depth_buffer"][camera_name] + # ) + return camera_buffers + + def save_last_available_captured_scrnshot_safe(self): + with self._viewer_config_lock: + if_new_capture_in_pending = self._viewer_config.get_viewer_status_unsafe(_VIEWER_STATUS_MAP.CAPTURE_SCRNSHOT) + image_output_format = self._viewer_config.image_output_format + + if not if_new_capture_in_pending: + with self._mj_lock: + img_buffer = self._mj.camera_data.frame_buffer + frame_stamp = self._mj.camera_data.frame_stamp + if img_buffer is not None and frame_stamp > 0: + img_buffer_corrected = np.flipud(img_buffer) + # save image: + imageio.imwrite(image_output_format % frame_stamp, img_buffer_corrected) + self._mj.camera_data.frame_stamp = -1 # prevent to capture again + + def is_key_registered_to_pause_program_safe(self): + with self._viewer_config_lock: + status = self._viewer_config.get_viewer_status_unsafe(_VIEWER_STATUS_MAP.PAUSED) + return status + + def is_key_registered_to_step_to_next_safe(self): + with self._viewer_config_lock: + status = self._viewer_config.get_viewer_status_unsafe(_VIEWER_STATUS_MAP.STEP_BY_STEP) + return status + + def reset_key_registered_to_step_to_next_safe(self): + with self._viewer_config_lock: + self._viewer_config.clear_viewer_status_unsafe(_VIEWER_STATUS_MAP.STEP_BY_STEP) + \ No newline at end of file diff --git a/setup.py b/setup.py index 0d52292..5551e58 100755 --- a/setup.py +++ b/setup.py @@ -2,22 +2,24 @@ import os from setuptools import find_packages, setup -VERSION = '0.1.0' +VERSION = '1.0.1' INSTALL_REQUIRES = ( - ['mujoco >= 2.1.5', - 'glfw >= 2.5.0', - 'imageio'] + [ + 'mujoco >= 2.2.0', + 'glfw >= 2.5.0', + 'imageio' + ] ) setup( name='mujoco-python-viewer', version=VERSION, - author='Rohan P. Singh', - author_email='rohan565singh@gmail.com', - url='https://github.com/rohanpsingh/mujoco-python-viewer', + author='Jack Xu', + author_email='projectbyjx@gmail.com', + url='https://github.com/jaku-jaku/jx-mujoco-python-viewer', description='Interactive renderer for MuJoCo Python', - long_description='Interactive renderer for MuJoCo Python', + long_description='Interactive renderer for MuJoCo Python, inspired by mujoco-python-viewer and mujoco-py', install_requires=INSTALL_REQUIRES, packages=find_packages(), python_requires='>=3.6',