diff --git a/examples/camera_feeds.py b/examples/camera_feeds.py index 56f92ca..2b94bfa 100644 --- a/examples/camera_feeds.py +++ b/examples/camera_feeds.py @@ -1,16 +1,15 @@ import random import threading import time + import cv2 + from stretch_mujoco.enums.actuators import Actuators from stretch_mujoco.enums.stretch_cameras import StretchCameras from stretch_mujoco.stretch_mujoco_simulator import StretchMujocoSimulator -def show_camera_feeds_sync( - sim: StretchMujocoSimulator, - print_fps: bool -): +def show_camera_feeds_sync(sim: StretchMujocoSimulator, print_fps: bool): """ Pull camera data from the simulator and display it using OpenCV. @@ -20,8 +19,10 @@ def show_camera_feeds_sync( camera_data = sim.pull_camera_data() if print_fps: - print(f"Physics fps: {sim.pull_status().fps}. Camera FPS: {camera_data.fps}. {sim.pull_status().sim_to_real_time_ratio_msg}") - + print( + f"Physics fps: {sim.pull_status().fps}. Camera FPS: {camera_data.fps}. {sim.pull_status().sim_to_real_time_ratio_msg}" + ) + for camera_name, pixels in camera_data.get_all(use_depth_color_map=True).items(): cv2.imshow(camera_name.name, pixels) @@ -33,6 +34,7 @@ def my_control_loop(): sim.move_to(Actuators.lift, random.random()) time.sleep(3) + if __name__ == "__main__": # You can use all the camera's, but it takes longer to render, and may affect the overall simulation FPS. # cameras_to_use = StretchCameras.all() @@ -49,4 +51,4 @@ def my_control_loop(): show_camera_feeds_sync(sim, True) except KeyboardInterrupt: - sim.stop() \ No newline at end of file + sim.stop() diff --git a/examples/draw_circles.py b/examples/draw_circles.py index 3739a66..2769d30 100644 --- a/examples/draw_circles.py +++ b/examples/draw_circles.py @@ -1,6 +1,6 @@ - import threading import time + import numpy as np from examples.camera_feeds import show_camera_feeds_sync @@ -9,11 +9,11 @@ from stretch_mujoco.stretch_mujoco_simulator import StretchMujocoSimulator -def draw_circle(n, diameter_m, arm_init, lift_init, sim:StretchMujocoSimulator): +def draw_circle(n, diameter_m, arm_init, lift_init, sim: StretchMujocoSimulator): """ From https://forum.hello-robot.com/t/creating-smooth-motion-using-trajectories/671 """ - t = np.linspace(0, 2*np.pi, n, endpoint=True) + t = np.linspace(0, 2 * np.pi, n, endpoint=True) x = (diameter_m / 2) * np.cos(t) + arm_init y = (diameter_m / 2) * np.sin(t) + lift_init circle_mat = np.c_[x, y] @@ -25,10 +25,11 @@ def draw_circle(n, diameter_m, arm_init, lift_init, sim:StretchMujocoSimulator): sim.wait_until_at_setpoint(Actuators.arm) sim.wait_until_at_setpoint(Actuators.lift) + def _run_draw_circle(): time.sleep(2) try: - while sim.is_running(): + while sim.is_running(): sim.move_to(Actuators.head_tilt, -1.5707) sim.move_to(Actuators.head_pan, -0.7853) @@ -39,13 +40,14 @@ def _run_draw_circle(): sim.move_to(Actuators.gripper, pos=-0.15) sim.wait_until_at_setpoint(Actuators.gripper) - + status = sim.pull_status() draw_circle(25, 0.2, status.arm.pos, status.lift.pos, sim) time.sleep(1) sim.home() time.sleep(2) - except ConnectionError: ... + except ConnectionError: + ... if __name__ == "__main__": @@ -65,4 +67,4 @@ def _run_draw_circle(): show_camera_feeds_sync(sim, True) except KeyboardInterrupt: - sim.stop() \ No newline at end of file + sim.stop() diff --git a/examples/gamepad_teleop.py b/examples/gamepad_teleop.py index c540c09..a4b5c24 100644 --- a/examples/gamepad_teleop.py +++ b/examples/gamepad_teleop.py @@ -3,15 +3,15 @@ import click import cv2 -from examples.camera_feeds import show_camera_feeds_sync from gamepad_controller import GamePadController +from examples.camera_feeds import show_camera_feeds_sync from stretch_mujoco import StretchMujocoSimulator -from stretch_mujoco.enums.stretch_cameras import StretchCameras from stretch_mujoco.enums.actuators import Actuators +from stretch_mujoco.enums.stretch_cameras import StretchCameras -sim:StretchMujocoSimulator -gamepad:GamePadController +sim: StretchMujocoSimulator +gamepad: GamePadController button_mapping = { "top_pad_pressed": ["wrist_pitch", 1, 0.05], diff --git a/examples/keyboard_teleop.py b/examples/keyboard_teleop.py index 5a21361..9c7d6d9 100644 --- a/examples/keyboard_teleop.py +++ b/examples/keyboard_teleop.py @@ -1,8 +1,8 @@ -from time import sleep -from pynput import keyboard from pprint import pprint +from time import sleep import click +from pynput import keyboard from examples.camera_feeds import show_camera_feeds_sync from examples.laser_scan import show_laser_scan @@ -95,7 +95,14 @@ def on_release(key, sim: StretchMujocoSimulator): @click.option("--imagery", is_flag=True, help="Show all the cameras' imagery") @click.option("--lidar", is_flag=True, help="Show the lidar scan in Matplotlib") @click.option("--print-ratio", is_flag=True, help="Print the sim-to-real time ratio to the cli.") -def main(scene_xml_path: str|None, robocasa_env: bool, imagery_nav: bool, imagery: bool, lidar:bool, print_ratio:bool): +def main( + scene_xml_path: str | None, + robocasa_env: bool, + imagery_nav: bool, + imagery: bool, + lidar: bool, + print_ratio: bool, +): cameras_to_use = StretchCameras.all() if imagery else [] if imagery_nav: cameras_to_use = [StretchCameras.cam_nav_rgb] @@ -110,9 +117,7 @@ def main(scene_xml_path: str|None, robocasa_env: bool, imagery_nav: bool, imager model, xml, objects_info = model_generation_wizard() sim = StretchMujocoSimulator( - model=model, - scene_xml_path=scene_xml_path, - cameras_to_use=cameras_to_use + model=model, scene_xml_path=scene_xml_path, cameras_to_use=cameras_to_use ) try: @@ -140,8 +145,8 @@ def main(scene_xml_path: str|None, robocasa_env: bool, imagery_nav: bool, imager try: show_laser_scan(scan_data=sensor_data.get_data(StretchSensors.base_lidar)) - except: ... - + except: + ... listener.stop() diff --git a/examples/laser_scan.py b/examples/laser_scan.py index 0487e8b..09e4328 100644 --- a/examples/laser_scan.py +++ b/examples/laser_scan.py @@ -1,7 +1,9 @@ import time + +import matplotlib import matplotlib.pyplot as plt import numpy as np -import matplotlib + from stretch_mujoco.enums.stretch_cameras import StretchCameras from stretch_mujoco.enums.stretch_sensors import StretchSensors from stretch_mujoco.stretch_mujoco_simulator import StretchMujocoSimulator @@ -9,7 +11,9 @@ try: # Some machines seem to need this for matplotlib to work. matplotlib.use("TkAgg") -except: ... +except: + ... + def show_laser_scan(scan_data: np.ndarray): @@ -23,20 +27,20 @@ def show_laser_scan(scan_data: np.ndarray): if len(filtered_distance) == 0: return time.sleep(1 / 15) - + degrees = np.array(range(len(scan_data))) degrees = np.radians(degrees) - + degrees_filtered = degrees[mask_lower & mask_upper] x = filtered_distance * np.cos(degrees_filtered) * -1 y = filtered_distance * np.sin(degrees_filtered) * -1 degrees_filtered = np.rad2deg(degrees_filtered) - front_idx = (degrees_filtered >= 150) & (degrees_filtered <= 210) # ~180 - back_idx = (degrees_filtered >= 330) | (degrees_filtered <= 30) # ~0 - right_idx= (degrees_filtered >= 60) & (degrees_filtered <= 120) # ~90 - left_idx = (degrees_filtered >= 240) & (degrees_filtered <= 300) # ~270 + front_idx = (degrees_filtered >= 150) & (degrees_filtered <= 210) # ~180 + back_idx = (degrees_filtered >= 330) | (degrees_filtered <= 30) # ~0 + right_idx = (degrees_filtered >= 60) & (degrees_filtered <= 120) # ~90 + left_idx = (degrees_filtered >= 240) & (degrees_filtered <= 300) # ~270 plt.scatter(x, y, color="r", s=5) plt.scatter(x[front_idx], y[front_idx], color="g", s=5) @@ -45,8 +49,8 @@ def show_laser_scan(scan_data: np.ndarray): plt.scatter(x[right_idx], y[right_idx], color="c", s=5) max_x = np.abs(x).max() max_y = np.abs(y).max() - plt.xlim([-max_x-1, max_x+1]) - plt.ylim([-max_y-1, max_y+1]) + plt.xlim([-max_x - 1, max_x + 1]) + plt.ylim([-max_y - 1, max_y + 1]) plt.legend(["All", "Front", "Back", "Left", "Right"]) plt.pause(1 / 15) @@ -70,7 +74,8 @@ def show_laser_scan(scan_data: np.ndarray): try: show_laser_scan(scan_data=sensor_data.get_data(StretchSensors.base_lidar)) - except: ... + except: + ... current_position = status.base.x diff --git a/examples/move_joints.py b/examples/move_joints.py index 145e9b5..e5bb0a5 100644 --- a/examples/move_joints.py +++ b/examples/move_joints.py @@ -1,9 +1,11 @@ import random + import numpy as np from stretch_mujoco.enums.actuators import Actuators from stretch_mujoco.stretch_mujoco_simulator import StretchMujocoSimulator + def lift_sequence(): LIFT_START_POS = 0.1 MOVE_LIFT_BY = 0.5 @@ -14,15 +16,19 @@ def lift_sequence(): start_lift_position = sim.pull_status().lift.pos if not np.isclose(start_lift_position, LIFT_START_POS, atol=0.05): - print(f"The lift did not move to the starting position. Should be at {LIFT_START_POS}, but is at {start_lift_position:.2f} instead.") + print( + f"The lift did not move to the starting position. Should be at {LIFT_START_POS}, but is at {start_lift_position:.2f} instead." + ) sim.move_by(Actuators.lift, MOVE_LIFT_BY) sim.wait_while_is_moving(Actuators.lift) current_lift_position = sim.pull_status().lift.pos - + if not np.isclose(start_lift_position, current_lift_position, atol=0.05): - print(f"The lift did not move by the specified amount. Asked to move from {start_lift_position:.4f} by {MOVE_LIFT_BY}, but ended up at {current_lift_position:.4f}. Should be {start_lift_position + MOVE_LIFT_BY :.4f}") + print( + f"The lift did not move by the specified amount. Asked to move from {start_lift_position:.4f} by {MOVE_LIFT_BY}, but ended up at {current_lift_position:.4f}. Should be {start_lift_position + MOVE_LIFT_BY :.4f}" + ) if __name__ == "__main__": @@ -51,8 +57,8 @@ def lift_sequence(): target *= -1 sim.set_base_velocity(v_linear=5.0, omega=30) - sim.move_to(Actuators.head_pan, random.random()-0.5) - sim.move_to(Actuators.head_tilt, random.random()-0.5) + sim.move_to(Actuators.head_pan, random.random() - 0.5) + sim.move_to(Actuators.head_tilt, random.random() - 0.5) sim.wait_until_at_setpoint(Actuators.head_pan) sim.wait_until_at_setpoint(Actuators.head_tilt) diff --git a/examples/robocasa_environment.py b/examples/robocasa_environment.py index f778331..516b6f6 100644 --- a/examples/robocasa_environment.py +++ b/examples/robocasa_environment.py @@ -17,7 +17,7 @@ def main(task: str, layout: int, style: int, write_to_file): # You can use all the camera's, but it takes longer to render, and may affect the overall simulation FPS. # cameras_to_use = StretchCameras.all() cameras_to_use = [StretchCameras.cam_d405_rgb] - + model, xml, objects_info = model_generation_wizard( task=task, layout=layout, diff --git a/examples/test_one_process.py b/examples/test_one_process.py index a51f39c..0f095cf 100644 --- a/examples/test_one_process.py +++ b/examples/test_one_process.py @@ -1,10 +1,8 @@ -from multiprocessing import Manager import signal import threading +from multiprocessing import Manager -from stretch_mujoco.mujoco_server import MujocoServer - -from stretch_mujoco.mujoco_server import MujocoServerProxies +from stretch_mujoco.mujoco_server import MujocoServer, MujocoServerProxies _manager = Manager() data_proxies = MujocoServerProxies.default(_manager) @@ -14,11 +12,11 @@ signal.signal(signal.SIGINT, lambda num, frame: event.set()) MujocoServer.launch_server( - scene_xml_path=None, - model=None, - camera_hz=30, + scene_xml_path=None, + model=None, + camera_hz=30, show_viewer_ui=True, - stop_mujoco_process_event=event, + stop_mujoco_process_event=event, data_proxies=data_proxies, - cameras_to_use=[] + cameras_to_use=[], ) diff --git a/examples/test_one_process_passive.py b/examples/test_one_process_passive.py index 99f1642..bffcac5 100644 --- a/examples/test_one_process_passive.py +++ b/examples/test_one_process_passive.py @@ -1,6 +1,6 @@ -from multiprocessing import Manager import signal import threading +from multiprocessing import Manager from stretch_mujoco.mujoco_server import MujocoServerProxies from stretch_mujoco.mujoco_server_passive import MujocoServerPassive @@ -13,11 +13,11 @@ signal.signal(signal.SIGINT, lambda num, frame: event.set()) MujocoServerPassive.launch_server( - scene_xml_path=None, - model=None, - camera_hz=30, + scene_xml_path=None, + model=None, + camera_hz=30, show_viewer_ui=True, - stop_mujoco_process_event=event, + stop_mujoco_process_event=event, data_proxies=data_proxies, - cameras_to_use=[] - ) + cameras_to_use=[], +) diff --git a/launch_sim.py b/launch_sim.py index f61a6f0..08dfd12 100644 --- a/launch_sim.py +++ b/launch_sim.py @@ -1,7 +1,7 @@ -import stretch_mujoco import click import cv2 +import stretch_mujoco from stretch_mujoco.enums.stretch_cameras import StretchCameras @@ -9,13 +9,9 @@ @click.option("--scene-xml-path", help="Path to a scene xml file") @click.option("--headless", is_flag=True, help="Run the simulation headless") @click.option("--imagery", is_flag=True, help="Show the cameras' imagery") -def main( - scene_xml_path: str, - headless: bool, - imagery: bool -) -> None: +def main(scene_xml_path: str, headless: bool, imagery: bool) -> None: cameras_to_use = StretchCameras.all() if imagery else [] - sim = stretch_mujoco.StretchMujocoSimulator(scene_xml_path,cameras_to_use=cameras_to_use) + sim = stretch_mujoco.StretchMujocoSimulator(scene_xml_path, cameras_to_use=cameras_to_use) sim.start(headless=headless) try: while sim.is_running(): diff --git a/stretch_mujoco/datamodels/status_command.py b/stretch_mujoco/datamodels/status_command.py index dfb354f..880a606 100644 --- a/stretch_mujoco/datamodels/status_command.py +++ b/stretch_mujoco/datamodels/status_command.py @@ -44,12 +44,12 @@ class StatusCommand: move_to: dict[str, CommandMove] = field(default_factory=dict) move_by: dict[str, CommandMove] = field(default_factory=dict) - base_velocity: CommandBaseVelocity = field(default_factory=lambda:CommandBaseVelocity(0, 0, False)) - keyframe: CommandKeyframe = field(default_factory=lambda:CommandKeyframe("", False)) + base_velocity: CommandBaseVelocity = field( + default_factory=lambda: CommandBaseVelocity(0, 0, False) + ) + keyframe: CommandKeyframe = field(default_factory=lambda: CommandKeyframe("", False)) coordinate_frame_arrows_viz: list[CommandCoordinateFrameArrowsViz] = field(default_factory=list) - - def set_move_to(self, command: CommandMove): """Sends a move_to command and removes the move_by command.""" self.move_to[command.actuator_name] = command diff --git a/stretch_mujoco/datamodels/status_stretch_camera.py b/stretch_mujoco/datamodels/status_stretch_camera.py index 0453b5c..0d01277 100644 --- a/stretch_mujoco/datamodels/status_stretch_camera.py +++ b/stretch_mujoco/datamodels/status_stretch_camera.py @@ -1,5 +1,6 @@ import copy from dataclasses import asdict, dataclass + import cv2 import numpy as np @@ -10,22 +11,25 @@ @dataclass class StatusStretchCameras: """ - A dataclass and helper methods to pack camera data. + A dataclass and helper methods to pack camera data. """ + time: float fps: float - cam_d405_rgb: np.ndarray|None = None - cam_d405_depth:np.ndarray|None = None - cam_d405_K: np.ndarray|None = None + cam_d405_rgb: np.ndarray | None = None + cam_d405_depth: np.ndarray | None = None + cam_d405_K: np.ndarray | None = None - cam_d435i_rgb:np.ndarray|None = None - cam_d435i_depth:np.ndarray|None = None - cam_d435i_K: np.ndarray|None = None + cam_d435i_rgb: np.ndarray | None = None + cam_d435i_depth: np.ndarray | None = None + cam_d435i_K: np.ndarray | None = None - cam_nav_rgb: np.ndarray|None = None + cam_nav_rgb: np.ndarray | None = None - def get_all(self, *, auto_rotate: bool = True, auto_correct_rgb=True, use_depth_color_map=False)-> dict[StretchCameras, np.ndarray]: + def get_all( + self, *, auto_rotate: bool = True, auto_correct_rgb=True, use_depth_color_map=False + ) -> dict[StretchCameras, np.ndarray]: """Returns the camera `{StretchCameras: pixels}` that are available (not None). `auto_rotate` will correct the rotation of the cam_d435i_rgb and cam_d435i_depth from their innately rotated optical frame. default: True @@ -33,21 +37,34 @@ def get_all(self, *, auto_rotate: bool = True, auto_correct_rgb=True, use_depth_ `use_depth_color_map` default: False Note: This get the values inside this dataclass; it does not poll from the simulator. - + Note: Alternatively, use `get_camera_data()` to get a specific camera's data. """ data: dict[StretchCameras, np.ndarray] = {} for camera in StretchCameras.all(): try: - data[camera] = self.get_camera_data(camera=camera, auto_rotate=auto_rotate, auto_correct_rgb=auto_correct_rgb,use_depth_color_map=use_depth_color_map) - except ValueError: ... # get_camera_data throws a ValueError when the value is None or doesn't exist. + data[camera] = self.get_camera_data( + camera=camera, + auto_rotate=auto_rotate, + auto_correct_rgb=auto_correct_rgb, + use_depth_color_map=use_depth_color_map, + ) + except ValueError: + ... # get_camera_data throws a ValueError when the value is None or doesn't exist. return data - - def get_camera_data(self, camera:StretchCameras, *, auto_rotate: bool = True, auto_correct_rgb=True, use_depth_color_map=False) -> np.ndarray: + + def get_camera_data( + self, + camera: StretchCameras, + *, + auto_rotate: bool = True, + auto_correct_rgb=True, + use_depth_color_map=False, + ) -> np.ndarray: """ Use this to get the camera data (pixels) using a StretchCameras instance. - + Throws a ValueError if the data is None. `auto_rotate` will correct the rotation of the cam_d435i_rgb and cam_d435i_depth from their innately rotated optical frame. default: True @@ -56,7 +73,7 @@ def get_camera_data(self, camera:StretchCameras, *, auto_rotate: bool = True, au Note: This get the values inside this dataclass; it does not poll from the simulator. """ - data:np.ndarray|None = None + data: np.ndarray | None = None if camera == StretchCameras.cam_d405_rgb and self.cam_d405_rgb is not None: data = self.cam_d405_rgb data = cv2.cvtColor(data, cv2.COLOR_RGB2BGR) if auto_correct_rgb else data @@ -78,10 +95,10 @@ def get_camera_data(self, camera:StretchCameras, *, auto_rotate: bool = True, au if data is None: raise ValueError(f"Tried to get {camera} data, but it is empty or not implemented.") - + return data - - def set_camera_data(self, camera:StretchCameras, data:np.ndarray): + + def set_camera_data(self, camera: StretchCameras, data: np.ndarray): """ Use this to match a StretchCameras enum instance with its property in StatusStretchCameras dataclass, to set the camera data property within this dataclass. @@ -102,7 +119,7 @@ def set_camera_data(self, camera:StretchCameras, data:np.ndarray): if camera == StretchCameras.cam_nav_rgb: self.cam_nav_rgb = data return - + raise NotImplementedError(f"Camera {camera} is not implemented.") @staticmethod @@ -111,15 +128,13 @@ def default(): Returns an empty instance with None or zeros for properties. """ return StatusStretchCameras(time=0, fps=0) - + def to_dict(self): return asdict(self) - + def copy(self): return StatusStretchCameras.from_dict(copy.copy(self.to_dict())) - + @staticmethod - def from_dict(dict_data:dict)-> "StatusStretchCameras": - return dataclass_from_dict(StatusStretchCameras, dict_data) #type: ignore - - + def from_dict(dict_data: dict) -> "StatusStretchCameras": + return dataclass_from_dict(StatusStretchCameras, dict_data) # type: ignore diff --git a/stretch_mujoco/datamodels/status_stretch_joints.py b/stretch_mujoco/datamodels/status_stretch_joints.py index 096d8be..2cefe00 100644 --- a/stretch_mujoco/datamodels/status_stretch_joints.py +++ b/stretch_mujoco/datamodels/status_stretch_joints.py @@ -1,7 +1,9 @@ import copy from dataclasses import asdict, dataclass + from stretch_mujoco.utils import dataclass_from_dict + @dataclass class PositionVelocity: pos: float @@ -11,24 +13,26 @@ class PositionVelocity: def default(): return PositionVelocity(0, 0) + @dataclass class BaseStatus: - x:float - y:float - theta:float - x_vel:float - theta_vel:float + x: float + y: float + theta: float + x_vel: float + theta_vel: float @staticmethod def default(): - return BaseStatus(0, 0, 0,0,0) + return BaseStatus(0, 0, 0, 0, 0) + @dataclass class StatusStretchJoints: time: float - fps:float + fps: float sim_to_real_time_ratio_msg: str - base:BaseStatus + base: BaseStatus lift: PositionVelocity arm: PositionVelocity head_pan: PositionVelocity @@ -38,20 +42,19 @@ class StatusStretchJoints: wrist_roll: PositionVelocity gripper: PositionVelocity - def __getitem__(self, name:str): + def __getitem__(self, name: str): """For backward compatibility: allows access with the square brackets []""" return getattr(self, name) - + def to_dict(self): return asdict(self) - + def copy(self): return StatusStretchJoints.from_dict(copy.copy(self.to_dict())) - - @staticmethod - def from_dict(dict_data:dict)-> "StatusStretchJoints": - return dataclass_from_dict(StatusStretchJoints, dict_data) #type: ignore + @staticmethod + def from_dict(dict_data: dict) -> "StatusStretchJoints": + return dataclass_from_dict(StatusStretchJoints, dict_data) # type: ignore @staticmethod def default(): @@ -70,5 +73,5 @@ def default(): PositionVelocity.default(), PositionVelocity.default(), PositionVelocity.default(), - PositionVelocity.default() + PositionVelocity.default(), ) diff --git a/stretch_mujoco/enums/actuators.py b/stretch_mujoco/enums/actuators.py index 0be1f48..ee92d49 100644 --- a/stretch_mujoco/enums/actuators.py +++ b/stretch_mujoco/enums/actuators.py @@ -24,7 +24,6 @@ class Actuators(Enum): gripper_left_finger = 12 gripper_right_finger = 13 - def get_joint_names_in_mjcf(self) -> list[str]: """ An actuator may have multiple joints in the MJCF. Return their names here. Useful for querying positions from Mujoco. @@ -57,7 +56,7 @@ def get_joint_names_in_mjcf(self) -> list[str]: return ["joint_head_tilt"] raise NotImplementedError(f"Joint names for {self} are not defined.") - + @staticmethod @cache def get_actuator_by_joint_names_in_mjcf(joint_name: str) -> "Actuators": @@ -91,11 +90,11 @@ def get_actuator_by_joint_names_in_mjcf(joint_name: str) -> "Actuators": return Actuators.left_wheel_vel if joint_name == "joint_right_wheel": return Actuators.right_wheel_vel - if joint_name == 'translate_mobile_base' or joint_name == 'position': + if joint_name == "translate_mobile_base" or joint_name == "position": return Actuators.base_translate - if joint_name == 'rotate_mobile_base': + if joint_name == "rotate_mobile_base": return Actuators.base_rotate - + if joint_name == "joint_lift": return Actuators.lift if "joint_arm" in joint_name: @@ -119,8 +118,6 @@ def get_actuator_by_joint_names_in_mjcf(joint_name: str) -> "Actuators": raise NotImplementedError(f"Actuator for {joint_name} is not defined.") - - def _get_status_attribute(self, is_position: bool, status: StatusStretchJoints) -> float: attribute_name = "pos" if is_position else "vel" if self == Actuators.arm: diff --git a/stretch_mujoco/enums/stretch_sensors.py b/stretch_mujoco/enums/stretch_sensors.py index f47ba29..e5ddda1 100644 --- a/stretch_mujoco/enums/stretch_sensors.py +++ b/stretch_mujoco/enums/stretch_sensors.py @@ -38,7 +38,7 @@ def lidar_names(resolution: int = 720): return [ f"{StretchSensors.base_lidar.name}{str(i).zfill(num_digits)}" for i in range(resolution) ] - + @staticmethod def from_mjmodel(mjmodel: mujoco._structs.MjModel) -> "list[StretchSensors]": """Get all the sensors in an mjmodel. We don't have the spec, only the compiled model. We're gonna try to find all the sensors.""" @@ -47,7 +47,7 @@ def from_mjmodel(mjmodel: mujoco._structs.MjModel) -> "list[StretchSensors]": try: index = 0 while True: - # We have no way of pulling the number of sensors via API. + # We have no way of pulling the number of sensors via API. # When we exceed the sensors in mjmodel.sensor, an IndexError will be thrown. name = mjmodel.sensor(index).name index += 1 @@ -60,7 +60,7 @@ def from_mjmodel(mjmodel: mujoco._structs.MjModel) -> "list[StretchSensors]": if len(remaining_sensors) == 0: break - except IndexError: ... + except IndexError: + ... return list(sensors) - diff --git a/stretch_mujoco/mujoco_server.py b/stretch_mujoco/mujoco_server.py index 9d4f882..7c9be2d 100644 --- a/stretch_mujoco/mujoco_server.py +++ b/stretch_mujoco/mujoco_server.py @@ -1,33 +1,32 @@ import contextlib -from dataclasses import dataclass -from multiprocessing.managers import DictProxy, SyncManager import signal import threading import time +from dataclasses import dataclass +from multiprocessing.managers import DictProxy, SyncManager from typing import Callable import click import mujoco -import mujoco._functions import mujoco._enums +import mujoco._functions import numpy as np from mujoco._structs import MjData, MjModel -import mujoco._enums +import stretch_mujoco.config as config +import stretch_mujoco.utils as utils +from stretch_mujoco.datamodels.status_command import CommandBaseVelocity, CommandMove, StatusCommand from stretch_mujoco.datamodels.status_stretch_camera import StatusStretchCameras from stretch_mujoco.datamodels.status_stretch_joints import StatusStretchJoints from stretch_mujoco.datamodels.status_stretch_sensors import StatusStretchSensors from stretch_mujoco.enums.actuators import Actuators from stretch_mujoco.enums.stretch_cameras import StretchCameras -import stretch_mujoco.config as config from stretch_mujoco.enums.stretch_sensors import StretchSensors from stretch_mujoco.mujoco_server_camera_manager import ( - MujocoServerCameraManagerThreaded, MujocoServerCameraManagerSync, + MujocoServerCameraManagerThreaded, ) -from stretch_mujoco.datamodels.status_command import CommandBaseVelocity, CommandMove, StatusCommand from stretch_mujoco.mujoco_server_sensor_manager import MujocoServerSensorManagerThreaded -import stretch_mujoco.utils as utils from stretch_mujoco.utils import FpsCounter @@ -90,7 +89,6 @@ def default(manager: SyncManager) -> "MujocoServerProxies": class BaseController: - def __init__(self, mujoco_server: "MujocoServer") -> None: self.mujoco_server = mujoco_server self.last_command: CommandMove | CommandBaseVelocity | None = None @@ -489,7 +487,7 @@ def _to_real_gripper_range(self, pos: float) -> float: config.robot_settings["gripper_min_max"], ) - def push_command(self, command_status:StatusCommand): + def push_command(self, command_status: StatusCommand): """ Handles setting mujoco ctrl properties to move joints. """ diff --git a/stretch_mujoco/mujoco_server_camera_manager.py b/stretch_mujoco/mujoco_server_camera_manager.py index 5720b17..cb81373 100644 --- a/stretch_mujoco/mujoco_server_camera_manager.py +++ b/stretch_mujoco/mujoco_server_camera_manager.py @@ -1,15 +1,16 @@ -from concurrent.futures import ThreadPoolExecutor, as_completed import platform import threading import time +from concurrent.futures import ThreadPoolExecutor, as_completed from typing import TYPE_CHECKING + import mujoco import mujoco._enums import numpy as np -from stretch_mujoco import config, utils -from stretch_mujoco.enums.stretch_cameras import StretchCameras +from stretch_mujoco import utils from stretch_mujoco.datamodels.status_stretch_camera import StatusStretchCameras +from stretch_mujoco.enums.stretch_cameras import StretchCameras from stretch_mujoco.utils import FpsCounter, switch_to_glfw_renderer if TYPE_CHECKING: @@ -96,7 +97,7 @@ def _pull_camera_data(self): def _create_camera_renderer(self, for_camera: StretchCameras): settings = for_camera.initial_camera_settings - # Update mujoco's offscreen gl buffer size to accomodate bigger resolutions: + # Update mujoco's offscreen gl buffer size to accommodate bigger resolutions: offscreen_buffer_width = self.mujoco_server.mjmodel.vis.global_.offwidth offscreen_buffer_height = self.mujoco_server.mjmodel.vis.global_.offheight @@ -109,7 +110,9 @@ def _create_camera_renderer(self, for_camera: StretchCameras): self.mujoco_server.mjmodel, width=settings.width, height=settings.height ) - renderer._scene_option.flags[mujoco._enums.mjtVisFlag.mjVIS_RANGEFINDER] = False # Disables the lidar yellow lines. + renderer._scene_option.flags[ + mujoco._enums.mjtVisFlag.mjVIS_RANGEFINDER + ] = False # Disables the lidar yellow lines. from stretch_mujoco.mujoco_server_passive import MujocoServerPassive diff --git a/stretch_mujoco/mujoco_server_managed.py b/stretch_mujoco/mujoco_server_managed.py index c6c09be..e65a2e3 100644 --- a/stretch_mujoco/mujoco_server_managed.py +++ b/stretch_mujoco/mujoco_server_managed.py @@ -1,16 +1,13 @@ import os -from stretch_mujoco.utils import override + import mujoco -import mujoco._functions import mujoco._callbacks -import mujoco._render -import mujoco._enums import mujoco.viewer from mujoco._structs import MjData, MjModel from stretch_mujoco.enums.stretch_cameras import StretchCameras from stretch_mujoco.mujoco_server import MujocoServer -from stretch_mujoco.mujoco_server_camera_manager import MujocoServerCameraManagerThreaded +from stretch_mujoco.utils import override class MujocoServerManaged(MujocoServer): diff --git a/stretch_mujoco/mujoco_server_passive.py b/stretch_mujoco/mujoco_server_passive.py index e2cb115..1280753 100644 --- a/stretch_mujoco/mujoco_server_passive.py +++ b/stretch_mujoco/mujoco_server_passive.py @@ -1,19 +1,18 @@ import threading import time -from stretch_mujoco.datamodels.status_command import StatusCommand -from stretch_mujoco.utils import Rx, Ry, Rz, override -import numpy as np import click import mujoco +import mujoco._enums import mujoco._functions import mujoco.viewer +import numpy as np from mujoco._enums import mjtGeom + +from stretch_mujoco.datamodels.status_command import StatusCommand from stretch_mujoco.enums.stretch_cameras import StretchCameras from stretch_mujoco.mujoco_server import MujocoServer -from stretch_mujoco.utils import FpsCounter - -import mujoco._enums +from stretch_mujoco.utils import FpsCounter, Rx, Ry, Rz, override class MujocoServerPassive(MujocoServer): @@ -55,17 +54,22 @@ def _run_ui_simulation(self, show_viewer_ui: bool): https://mujoco.readthedocs.io/en/stable/python.html#passive-viewer """ - self.viewer = mujoco.viewer.launch_passive( + self.viewer = mujoco.viewer.launch_passive( self.mjmodel, self.mjdata, show_left_ui=show_viewer_ui, show_right_ui=show_viewer_ui ) - self.viewer._opt.flags[mujoco._enums.mjtVisFlag.mjVIS_RANGEFINDER] = False # Disables the lidar yellow lines. + self.viewer._opt.flags[ + mujoco._enums.mjtVisFlag.mjVIS_RANGEFINDER + ] = False # Disables the lidar yellow lines. with self.viewer as viewer: physics_thread = threading.Thread( target=self._physics_loop, name="PhysicsThread", - args=(viewer.lock(), lambda: viewer.is_running() and not self._is_requested_to_stop()), + args=( + viewer.lock(), + lambda: viewer.is_running() and not self._is_requested_to_stop(), + ), daemon=True, ) physics_thread.start() @@ -77,12 +81,12 @@ def _run_ui_simulation(self, show_viewer_ui: bool): ) # 1/Hz.Put the UI thread to sleep so that the physics thread can do work, to mitigate `viewer.lock()` locking physics thread. click.secho( - f"Using the Mujoco Passive Viewer. Note: UI thread and camera rendering is capped to {1/UI_FPS_CAP_RATE}Hz to increase performance. You can set this rate using the `camera_rate` arugment.", + f"Using the Mujoco Passive Viewer. Note: UI thread and camera rendering is capped to {1/UI_FPS_CAP_RATE}Hz to increase performance. You can set this rate using the `camera_rate` argument.", fg="green", ) # Replace the camera_lock with the viewer lock so that we're not accessing mjdata at the same time as the physics thread. - self.camera_manager.camera_lock = viewer.lock() #type: ignore + self.camera_manager.camera_lock = viewer.lock() # type: ignore while viewer.is_running() and not self._is_requested_to_stop(): fps.tick() @@ -119,26 +123,29 @@ def _run_ui_simulation(self, show_viewer_ui: bool): click.secho("Mujoco viewer has terminated.", fg="blue") - def push_command(self, command_status:StatusCommand): + def push_command(self, command_status: StatusCommand): command_arrows = command_status.coordinate_frame_arrows_viz.copy() for arrows in command_arrows: if arrows.trigger: - self._add_axes_to_user_scn(self.viewer.user_scn, np.array(arrows.position) , arrows.rotation) + self._add_axes_to_user_scn( + self.viewer.user_scn, np.array(arrows.position), arrows.rotation + ) command_status.coordinate_frame_arrows_viz.remove(arrows) super().push_command(command_status) - @override - def _add_axes_to_user_scn(self, - user_scn, - origin: np.ndarray, - rotation: tuple[float,float,float], - length: float = 0.2, - radius: float = 0.006): + def _add_axes_to_user_scn( + self, + user_scn, + origin: np.ndarray, + rotation: tuple[float, float, float], + length: float = 0.2, + radius: float = 0.006, + ): """ Draw a right-handed RGB frame in `user_scn` using mjv_initGeom. @@ -146,29 +153,19 @@ def _add_axes_to_user_scn(self, * `origin` 3-vector in world frame * `R` 3×3 rotation matrix, columns are local x,y,z in world frame """ - colors = np.array([[1, 0, 0, 1], # +X - [0, 1, 0, 1], # +Y - [0, 0, 1, 1]]) # +Z - + colors = np.array([[1, 0, 0, 1], [0, 1, 0, 1], [0, 0, 1, 1]]) # +X # +Y # +Z + rot_matrix = Rx(rotation[0]) @ Ry(rotation[1]) @ Rz(rotation[2]) for axis in range(3): if axis == 0: # Rotate +Z to +X: -90° about Y-axis - R = np.array([ - [0, 0, 1], - [0, 1, 0], - [-1, 0, 0] - ]) - elif axis ==1: + R = np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]]) + elif axis == 1: # Rotate +Z to +Y: -90° about X-axis - R = np.array([ - [-1, 0, 0], - [0, 0, 1], - [0, 1, 0] - ]) - elif axis ==2: + R = np.array([[-1, 0, 0], [0, 0, 1], [0, 1, 0]]) + elif axis == 2: # No rotation needed - R = np.eye(3) + R = np.eye(3) R = rot_matrix @ R @@ -177,7 +174,7 @@ def _add_axes_to_user_scn(self, geom = user_scn.geoms[user_scn.ngeom] mujoco._functions.mjv_initGeom( geom, - type= mjtGeom.mjGEOM_ARROW, + type=mjtGeom.mjGEOM_ARROW, size=size, pos=origin, mat=np.array(R).flatten(), diff --git a/stretch_mujoco/mujoco_server_sensor_manager.py b/stretch_mujoco/mujoco_server_sensor_manager.py index ed9c159..3dd4be2 100644 --- a/stretch_mujoco/mujoco_server_sensor_manager.py +++ b/stretch_mujoco/mujoco_server_sensor_manager.py @@ -1,10 +1,11 @@ import threading import time from typing import TYPE_CHECKING + import numpy as np -from stretch_mujoco.enums.stretch_sensors import StretchSensors from stretch_mujoco.datamodels.status_stretch_sensors import StatusStretchSensors +from stretch_mujoco.enums.stretch_sensors import StretchSensors from stretch_mujoco.utils import FpsCounter if TYPE_CHECKING: diff --git a/stretch_mujoco/robocasa_gen.py b/stretch_mujoco/robocasa_gen.py index b7cd5cf..4a18f82 100644 --- a/stretch_mujoco/robocasa_gen.py +++ b/stretch_mujoco/robocasa_gen.py @@ -21,9 +21,7 @@ def get_styles() -> OrderedDict: - raw_styles = dict( - map(lambda item: (item.value, item.name.lower().capitalize()), StyleType) - ) + raw_styles = dict(map(lambda item: (item.value, item.name.lower().capitalize()), StyleType)) styles = OrderedDict() for k in sorted(raw_styles.keys()): if k < 0: @@ -97,17 +95,17 @@ def choose_option(options, option_name, show_keys=False, default=None, default_m # Return the chosen environment name return choice + def choose_layout(): - layout = choose_option( - layouts, "kitchen layout", default=-1, default_message="random layouts" - ) - + layout = choose_option(layouts, "kitchen layout", default=-1, default_message="random layouts") + if layout == -1: layout = np.random.choice(range(10)) print(colored(f"Randomly choosing layout... id: {layout}", "yellow")) - + return layout + def choose_style(): styles = get_styles() style = choose_option(styles, "kitchen style", default=-1, default_message="random styles") @@ -115,17 +113,20 @@ def choose_style(): if style == -1: style = np.random.choice(range(11)) print(colored(f"Randomly choosing style... id: {style}", "yellow")) - + return style -def layout_from_str(layout:str) -> int: + +def layout_from_str(layout: str) -> int: """Returns the index of the layout in the orderedDict""" return list(layouts.values()).index(layout) -def style_from_str(style:str) -> int: + +def style_from_str(style: str) -> int: """Returns the index of the style in the orderedDict""" return list(get_styles().values()).index(style) + def model_generation_wizard( task: str = "PnPCounterToCab", layout: int = None, diff --git a/stretch_mujoco/stretch_mujoco_simulator.py b/stretch_mujoco/stretch_mujoco_simulator.py index 4ccd265..d4c30c7 100644 --- a/stretch_mujoco/stretch_mujoco_simulator.py +++ b/stretch_mujoco/stretch_mujoco_simulator.py @@ -1,17 +1,24 @@ import atexit -from multiprocessing import Lock, Manager, Process - import multiprocessing import platform import signal import sys import threading import time +from multiprocessing import Lock, Manager, Process import click import numpy as np from mujoco._structs import MjModel +import stretch_mujoco.utils as utils +from stretch_mujoco.datamodels.status_command import ( + CommandBaseVelocity, + CommandCoordinateFrameArrowsViz, + CommandKeyframe, + CommandMove, + StatusCommand, +) from stretch_mujoco.datamodels.status_stretch_camera import StatusStretchCameras from stretch_mujoco.datamodels.status_stretch_joints import StatusStretchJoints from stretch_mujoco.datamodels.status_stretch_sensors import StatusStretchSensors @@ -20,15 +27,7 @@ from stretch_mujoco.mujoco_server import MujocoServer, MujocoServerProxies from stretch_mujoco.mujoco_server_managed import MujocoServerManaged from stretch_mujoco.mujoco_server_passive import MujocoServerPassive -from stretch_mujoco.datamodels.status_command import ( - CommandBaseVelocity, - CommandCoordinateFrameArrowsViz, - CommandKeyframe, - CommandMove, - StatusCommand, -) -import stretch_mujoco.utils as utils -from stretch_mujoco.utils import require_connection, block_until_check_succeeds +from stretch_mujoco.utils import block_until_check_succeeds, require_connection class StretchMujocoSimulator: @@ -275,7 +274,7 @@ def wait_until_at_setpoint( check=lambda: self.is_reached_set_position( actuator=actuator, position_tolerance=position_tolerance ) - == True, + is True, is_alive=self.is_running, ): pos = move_command.pos @@ -314,9 +313,7 @@ def check_if_moved(): Actuators.base_rotate, Actuators.base_translate, ]: - current_position = actuator.get_position_relative( - self.pull_status() - ) + current_position = actuator.get_position_relative(self.pull_status()) if actuator == Actuators.left_wheel_vel or actuator == Actuators.base_translate: current_position = current_position[0] elif actuator == Actuators.right_wheel_vel: @@ -326,7 +323,7 @@ def check_if_moved(): else: current_position = actuator.get_position(self.pull_status()) - if not actuator in self._last_movement_positions: + if actuator not in self._last_movement_positions: self._last_movement_positions[actuator] = current_position return True @@ -340,7 +337,7 @@ def check_if_moved(): if not block_until_check_succeeds( wait_timeout=timeout, - check=lambda: check_if_moved() == False, + check=lambda: check_if_moved() is False, is_alive=self.is_running, ): if timeout is not None: diff --git a/stretch_mujoco/utils.py b/stretch_mujoco/utils.py index f24d914..ee18646 100644 --- a/stretch_mujoco/utils.py +++ b/stretch_mujoco/utils.py @@ -1,25 +1,21 @@ import dataclasses +import importlib.resources import math import re import time import xml.etree.ElementTree as ET +from functools import wraps from typing import TYPE_CHECKING, Callable, Tuple import cv2 -import numpy as np -import importlib.resources -import urchin as urdf_loader - - -from functools import wraps - import mujoco -import mujoco._functions import mujoco._callbacks -import mujoco._render import mujoco._enums +import mujoco._functions +import mujoco._render import mujoco.viewer import numpy as np +import urchin as urdf_loader from mujoco._structs import MjModel from mujoco.glfw import GLContext as GlFwContext @@ -43,7 +39,7 @@ def require_connection(function): """Wraps class methods that need self""" - def wrapper_function(self:"StretchMujocoSimulator", *args, **kwargs): + def wrapper_function(self: "StretchMujocoSimulator", *args, **kwargs): if not self.is_running(): raise ConnectionError( "The Stretch Mujoco Simulator is not running. Use the start() method to start it." @@ -66,15 +62,27 @@ def Rx(theta): Rotation matrix about x-axis """ return np.matrix( - [[1,0,0], [0, math.cos(theta), -math.sin(theta)], [0, math.sin(theta), math.cos(theta)], ] + [ + [1, 0, 0], + [0, math.cos(theta), -math.sin(theta)], + [0, math.sin(theta), math.cos(theta)], + ] ) + + def Ry(theta): """ Rotation matrix about y-axis """ return np.matrix( - [[math.cos(theta), 0, math.sin(theta)], [0, 1, 0],[-math.sin(theta), 0, math.cos(theta)],] + [ + [math.cos(theta), 0, math.sin(theta)], + [0, 1, 0], + [-math.sin(theta), 0, math.cos(theta)], + ] ) + + def Rz(theta): """ Rotation matrix about z-axis @@ -143,15 +151,14 @@ def __init__(self): self.fps = 0 """The actual fps count""" - self.sim_to_real_ratio:float|None = None + self.sim_to_real_ratio: float | None = None """Sim time compared with real time""" self._last_sim_time = 0 - - def tick(self, sim_time:float|None = None): + def tick(self, sim_time: float | None = None): """ - Call this during step() to update the fps counter. + Call this during step() to update the fps counter. Pass sim_time to calculate sim-to-real time. """ @@ -161,22 +168,24 @@ def tick(self, sim_time:float|None = None): # When one second has passed, count: if elapsed >= 1.0: new_wall_time = time.perf_counter() - + if sim_time: - self.sim_to_real_ratio = (sim_time - self._last_sim_time)/(new_wall_time - self._wall_time) + self.sim_to_real_ratio = (sim_time - self._last_sim_time) / ( + new_wall_time - self._wall_time + ) self._last_sim_time = sim_time self.fps = self._fps_counter / elapsed self._wall_time = new_wall_time self._fps_counter = 0 - @property - def sim_to_real_time_ratio_msg(self): + def sim_to_real_time_ratio_msg(self): if self.sim_to_real_ratio is None: return "sim_to_real_ratio is not set. Call `tick(sim_time=)` with the sim_time to calculate it." return f"Sim is running {self.sim_to_real_ratio:.3f}x as fast as realtime" + class URDFmodel: def __init__(self) -> None: """ @@ -383,7 +392,7 @@ def dataclass_from_dict(klass, dict_data: dict): def block_until_check_succeeds( - wait_timeout: float|None, check: Callable[[], bool], is_alive: Callable[[], bool] + wait_timeout: float | None, check: Callable[[], bool], is_alive: Callable[[], bool] ) -> bool: """Blocks until the check callback succeeds""" @@ -392,7 +401,7 @@ def block_until_check_succeeds( if check(): return True return False - + start_time = time.time() while time.time() - start_time < wait_timeout: @@ -432,9 +441,10 @@ def switch_to_glfw_renderer(mjmodel: MjModel, renderer: mujoco.Renderer): # Only Python >12 has override. override = __import__("typing").override except: # noqa + def override(func): @wraps(func) def wrapper(*args, **kwargs): return func(*args, **kwargs) - return wrapper \ No newline at end of file + return wrapper