feat(sim): save air-insert and rollout validation updates
This commit is contained in:
Logic
200
render.py
Normal file
200
render.py
Normal file
@@ -0,0 +1,200 @@
|
||||
import mujoco
|
||||
from mujoco import viewer
|
||||
import sys
|
||||
import numpy as np
|
||||
import time
|
||||
import threading
|
||||
|
||||
|
||||
class MjBasicRenderer:
|
||||
def __new__(cls, *args, **kwargs):
|
||||
return super().__new__(cls)
|
||||
|
||||
def __init__(self, mj_model=None, mj_data=None):
|
||||
# keyboard flag
|
||||
self.render_paused = True
|
||||
self.exit_flag = False
|
||||
# init param
|
||||
self.mj_model = mj_model
|
||||
self.mj_data = mj_data
|
||||
self.renderer = "viewer" # default
|
||||
self.viewer = None
|
||||
self._image = None
|
||||
|
||||
# Set up mujoco viewer
|
||||
self.image_renderer = mujoco.Renderer(self.mj_model)
|
||||
|
||||
def __del__(self):
|
||||
pass
|
||||
|
||||
def _init_renderer(self):
|
||||
"""Initialize renderer, choose official renderer with "viewer"(joined from version 2.3.3),
|
||||
another renderer with "mujoco_viewer"
|
||||
"""
|
||||
|
||||
def key_callback(keycode):
|
||||
if keycode == 32: # space
|
||||
self.render_paused = not self.render_paused
|
||||
elif keycode == 256: # escape
|
||||
self.exit_flag = not self.exit_flag
|
||||
|
||||
if self.renderer == "viewer":
|
||||
# This function does not block, allowing user code to continue execution.
|
||||
self.viewer = viewer.launch_passive(
|
||||
self.mj_model,
|
||||
self.mj_data,
|
||||
key_callback=key_callback,
|
||||
show_left_ui=False,
|
||||
show_right_ui=False,
|
||||
)
|
||||
self.set_renderer_config()
|
||||
else:
|
||||
raise ValueError("Invalid renderer for some reason.")
|
||||
|
||||
def render(self):
|
||||
"""mujoco render"""
|
||||
if self.viewer is not None and self.render_paused is True:
|
||||
if self.viewer.is_running() and self.exit_flag is False:
|
||||
self.viewer: viewer.Handle
|
||||
self.viewer.sync()
|
||||
else:
|
||||
self.viewer.close()
|
||||
|
||||
def set_renderer_config(self):
|
||||
"""Setup mujoco global config while using viewer as renderer.
|
||||
It should be noted that the render thread need locked.
|
||||
"""
|
||||
self.viewer.cam.lookat = np.array([0.4, 0, 0.5])
|
||||
self.viewer.cam.azimuth -= 0.005
|
||||
with self.viewer.lock():
|
||||
self.viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = int(
|
||||
self.mj_data.time % 2
|
||||
)
|
||||
|
||||
|
||||
try:
|
||||
import cv2
|
||||
except ImportError:
|
||||
print("Could not import cv2, please install it to enable camera viewer.")
|
||||
|
||||
|
||||
class MjMultiRenderer(MjBasicRenderer):
|
||||
|
||||
# __slots__=('mj_model','mj_data','renderer','enable_camera_viewer')
|
||||
def __new__(cls, *args, **kwargs):
|
||||
return super().__new__(cls)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mj_model=None,
|
||||
mj_data=None,
|
||||
renderer=None,
|
||||
enable_camera_viewer=False,
|
||||
enable_depth=False,
|
||||
):
|
||||
super().__init__(mj_model, mj_data)
|
||||
self._depth = None
|
||||
self.renderer = renderer
|
||||
self._init_renderer()
|
||||
self.enable_camera_viewer = enable_camera_viewer
|
||||
if self.enable_camera_viewer:
|
||||
self.enable_depth = enable_depth
|
||||
self._init_window()
|
||||
else:
|
||||
self.enable_depth = False
|
||||
print("No Camera View")
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
|
||||
def _init_renderer(self):
|
||||
"""
|
||||
Initialize renderer, choose official renderer with "viewer"(joined from version 2.3.3)
|
||||
"""
|
||||
if self.renderer == "unity":
|
||||
# TODO: Support unity renderer.
|
||||
raise ValueError("Unity renderer init failed for no supporting reason")
|
||||
elif self.renderer == "viewer":
|
||||
super()._init_renderer()
|
||||
print("mujoco viewer init !")
|
||||
else:
|
||||
raise ValueError("renderer init failed for some reason.")
|
||||
|
||||
def _init_window(self, name="Camera view"):
|
||||
if not self.enable_depth:
|
||||
cv2.namedWindow(name, cv2.WINDOW_NORMAL)
|
||||
else:
|
||||
cv2.namedWindow(name, cv2.WINDOW_NORMAL)
|
||||
cv2.namedWindow("Camera depth view", cv2.WINDOW_NORMAL)
|
||||
|
||||
def render(self):
|
||||
"""render mujoco"""
|
||||
if self.renderer == "viewer":
|
||||
super().render()
|
||||
elif self.renderer == "unity":
|
||||
# TODO: Support unity renderer.
|
||||
raise ValueError("Unity renderer not supported now.")
|
||||
else:
|
||||
raise ValueError("Invalid renderer for some reason.")
|
||||
|
||||
def render(self):
|
||||
"""mujoco render"""
|
||||
if self.viewer is not None and self.render_paused is True:
|
||||
if self.viewer.is_running() and self.exit_flag is False:
|
||||
self.viewer: viewer.Handle
|
||||
self.viewer.sync()
|
||||
else:
|
||||
self.viewer.close()
|
||||
|
||||
def camera_render(self, cam=None):
|
||||
if self.enable_camera_viewer:
|
||||
if not self.enable_depth:
|
||||
rgb, depth = self.render_from_camera(cam)
|
||||
rgb = cv2.resize(rgb, (1920, 1600))
|
||||
cv2.imshow("Camera view", rgb)
|
||||
cv2.waitKey(1)
|
||||
|
||||
else:
|
||||
rgb, depth = self.render_from_camera(cam)
|
||||
cv2.imshow("Camera view", rgb)
|
||||
cv2.imshow("Camera depth view", depth)
|
||||
cv2.waitKey(1)
|
||||
|
||||
else:
|
||||
print("camera info disable")
|
||||
return
|
||||
|
||||
def render_from_camera(self, cam=None):
|
||||
self.image_renderer.update_scene(self.mj_data, camera=cam)
|
||||
if self.enable_depth is True:
|
||||
self.image_renderer.enable_depth_rendering()
|
||||
org = self.image_renderer.render()
|
||||
depth = org[:, :]
|
||||
self.image_renderer.disable_depth_rendering()
|
||||
org = self.image_renderer.render()
|
||||
image = org[:, :, ::-1]
|
||||
else:
|
||||
org = self.image_renderer.render()
|
||||
image = org[:, :, ::-1]
|
||||
depth = np.zeros([240, 320])
|
||||
return image, depth
|
||||
|
||||
def close(self):
|
||||
"""close the environment."""
|
||||
if self.enable_camera_viewer and self.viewer.is_running() == False:
|
||||
cv2.destroyAllWindows()
|
||||
self.viewer.close()
|
||||
# sys.exit(0)
|
||||
|
||||
# def get_cam_intrinsic(self, fovy=45.0, width=320, height=240):
|
||||
# aspect = width * 1.0 / height
|
||||
# fovx = np.degrees(2 * np.arctan(aspect * np.tan(np.radians(fovy / 2))))
|
||||
|
||||
# cx = 0.5 * width
|
||||
# cy = 0.5 * height
|
||||
# fx = cx / np.tan(fovx * np.pi / 180 * 0.5)
|
||||
# fy = cy / np.tan(fovy * np.pi / 180 * 0.5)
|
||||
|
||||
# K = np.array([[fx, 0, cx],
|
||||
# [0, fy, cy],
|
||||
# [0, 0, 1]], dtype=np.float32)
|
||||
Reference in New Issue
Block a user