From fa7eb860ea179f3ccb9ce1251af9af6c13df3e20 Mon Sep 17 00:00:00 2001 From: steve Date: Tue, 19 Nov 2024 12:36:28 +0100 Subject: [PATCH] cameras seems to work in Gazebo. Cleaner disconnect --- .../common/robot_devices/cameras/reachy2.py | 77 +++++++++++++ .../common/robot_devices/robots/reachy2.py | 107 ++++++++++-------- lerobot/scripts/test_reachy.py | 6 +- 3 files changed, 138 insertions(+), 52 deletions(-) create mode 100644 lerobot/common/robot_devices/cameras/reachy2.py diff --git a/lerobot/common/robot_devices/cameras/reachy2.py b/lerobot/common/robot_devices/cameras/reachy2.py new file mode 100644 index 00000000..f7b4f11c --- /dev/null +++ b/lerobot/common/robot_devices/cameras/reachy2.py @@ -0,0 +1,77 @@ +""" +Wrapper for Reachy2 camera from sdk +""" + +import argparse +import concurrent.futures +import math +import platform +import shutil +import threading +import time +from dataclasses import dataclass, replace +from pathlib import Path +from threading import Thread + +import numpy as np +from PIL import Image + +from lerobot.common.robot_devices.utils import ( + RobotDeviceAlreadyConnectedError, + RobotDeviceNotConnectedError, + busy_wait, +) +from lerobot.common.utils.utils import capture_timestamp_utc + +from reachy2_sdk.media.camera import CameraView +from reachy2_sdk.media.camera_manager import CameraManager + + +@dataclass +class ReachyCameraConfig: + + fps: int | None = None + width: int | None = None + height: int | None = None + color_mode: str = "rgb" + rotation: int | None = None + mock: bool = False + + +class ReachyCamera: + def __init__(self, host: str, port: int,name: str, image_type: str, config: ReachyCameraConfig | None = None, **kwargs): + self.host=host + self.port=port + self.image_type = image_type + self.name = name + self.config = config + self.cam_manager=None + self.is_connected = False + self.logs = {} + + def connect(self): + if not self.is_connected: + self.cam_manager=CameraManager(host=self.host, port=self.port) + self.cam_manager.initialize_cameras()#FIXME: maybe we should not re-initialize + self.is_connected=True + + def read(self) -> np.ndarray: + + if not self.is_connected: + self.connect() + + if self.name=='teleop' and hasattr(self.cam_manager,'teleop'): + if self.image_type=='left': + return self.cam_manager.teleop.get_frame(CameraView.LEFT) + elif self.image_type=='right': + return self.cam_manager.teleop.get_frame(CameraView.RIGHT) + else: + return None + elif self.name=='depth' and hasattr(self.cam_manager,'depth'): + + if self.image_type=='depth': + return self.cam_manager.depth.get_depth_frame() + elif self.image_type=='rgb': + return self.cam_manager.depth.get_frame() + else: + return None diff --git a/lerobot/common/robot_devices/robots/reachy2.py b/lerobot/common/robot_devices/robots/reachy2.py index e7cf3901..9a24e55a 100644 --- a/lerobot/common/robot_devices/robots/reachy2.py +++ b/lerobot/common/robot_devices/robots/reachy2.py @@ -45,16 +45,17 @@ class ReachyRobot(): self.robot_type = self.config.robot_type self.cameras = self.config.cameras - self.is_connected = False #at init Reachy2 is in fact connected... + self.is_connected = False self.teleop = None self.logs = {} self.reachy=ReachySDK(host=config.ip_address) - + self.is_connected=True #at init Reachy2 is in fact connected... self.state_keys = None self.action_keys = None def connect(self) -> None: + print("Connecting to Reachy") self.reachy.is_connected = self.reachy.connect() if not self.is_connected: print(f'Cannot connect to Reachy at address {self.config.ip_address}. Maybe a connection already exists.') @@ -81,64 +82,70 @@ class ReachyRobot(): def get_state(self) -> dict: - return { - "neck_yaw.pos": np.radians(self.reachy.head.neck.yaw.present_position), - "neck_pitch.pos": np.radians(self.reachy.head.neck.pitch.present_position), - "neck_roll.pos": np.radians(self.reachy.head.neck.roll.present_position), - "r_shoulder_pitch.pos": np.radians(self.reachy.r_arm.shoulder.pitch.present_position), - "r_shoulder_roll.pos": np.radians(self.reachy.r_arm.shoulder.roll.present_position), - "r_elbow_yaw.pos": np.radians(self.reachy.r_arm.elbow.yaw.present_position), - "r_elbow_pitch.pos": np.radians(self.reachy.r_arm.elbow.pitch.present_position), - "r_wrist_roll.pos": np.radians(self.reachy.r_arm.wrist.roll.present_position), - "r_wrist_pitch.pos": np.radians(self.reachy.r_arm.wrist.pitch.present_position), - "r_wrist_yaw.pos": np.radians(self.reachy.r_arm.wrist.yaw.present_position), - "r_gripper.pos": np.radians(self.reachy.r_arm.gripper.present_position), - "l_shouldel_pitch.pos": np.radians(self.reachy.l_arm.shoulder.pitch.present_position), - "l_shouldel_roll.pos": np.radians(self.reachy.l_arm.shoulder.roll.present_position), - "l_elbow_yaw.pos": np.radians(self.reachy.l_arm.elbow.yaw.present_position), - "l_elbow_pitch.pos": np.radians(self.reachy.l_arm.elbow.pitch.present_position), - "l_wrist_roll.pos": np.radians(self.reachy.l_arm.wrist.roll.present_position), - "l_wrist_pitch.pos": np.radians(self.reachy.l_arm.wrist.pitch.present_position), - "l_wrist_yaw.pos": np.radians(self.reachy.l_arm.wrist.yaw.present_position), - "l_gripper.pos": np.radians(self.reachy.l_arm.gripper.present_position), - #TODO mobile base - } + if self.is_connected: + return { + "neck_yaw.pos": np.radians(self.reachy.head.neck.yaw.present_position), + "neck_pitch.pos": np.radians(self.reachy.head.neck.pitch.present_position), + "neck_roll.pos": np.radians(self.reachy.head.neck.roll.present_position), + "r_shoulder_pitch.pos": np.radians(self.reachy.r_arm.shoulder.pitch.present_position), + "r_shoulder_roll.pos": np.radians(self.reachy.r_arm.shoulder.roll.present_position), + "r_elbow_yaw.pos": np.radians(self.reachy.r_arm.elbow.yaw.present_position), + "r_elbow_pitch.pos": np.radians(self.reachy.r_arm.elbow.pitch.present_position), + "r_wrist_roll.pos": np.radians(self.reachy.r_arm.wrist.roll.present_position), + "r_wrist_pitch.pos": np.radians(self.reachy.r_arm.wrist.pitch.present_position), + "r_wrist_yaw.pos": np.radians(self.reachy.r_arm.wrist.yaw.present_position), + "r_gripper.pos": np.radians(self.reachy.r_arm.gripper.present_position), + "l_shouldel_pitch.pos": np.radians(self.reachy.l_arm.shoulder.pitch.present_position), + "l_shouldel_roll.pos": np.radians(self.reachy.l_arm.shoulder.roll.present_position), + "l_elbow_yaw.pos": np.radians(self.reachy.l_arm.elbow.yaw.present_position), + "l_elbow_pitch.pos": np.radians(self.reachy.l_arm.elbow.pitch.present_position), + "l_wrist_roll.pos": np.radians(self.reachy.l_arm.wrist.roll.present_position), + "l_wrist_pitch.pos": np.radians(self.reachy.l_arm.wrist.pitch.present_position), + "l_wrist_yaw.pos": np.radians(self.reachy.l_arm.wrist.yaw.present_position), + "l_gripper.pos": np.radians(self.reachy.l_arm.gripper.present_position), + #TODO mobile base + } + else: + return {} def capture_observation(self) -> dict: - before_read_t = time.perf_counter() - state = self.get_state() - self.logs["read_pos_dt_s"] = time.perf_counter() - before_read_t + if self.is_connected: - if self.state_keys is None: - self.state_keys = list(state) + before_read_t = time.perf_counter() + state = self.get_state() + self.logs["read_pos_dt_s"] = time.perf_counter() - before_read_t - state = torch.as_tensor(list(state.values())) + if self.state_keys is None: + self.state_keys = list(state) - # Capture images from cameras - images = {} - for name in self.cameras: - before_camread_t = time.perf_counter() - images[name] = self.cameras[name].read() #Reachy cameras read() is not blocking? - print(f'name: {name} img: {images[name]}') - if images[name] is not None: - images[name] = torch.from_numpy(copy(images[name][0])) #seems like I need to copy? - self.logs[f"read_camera_{name}_dt_s"] = images[name][1] #full timestamp, TODO dt + state = torch.as_tensor(list(state.values())) - # Populate output dictionnaries - obs_dict = {} - obs_dict["observation.state"] = state - for name in self.cameras: - obs_dict[f"observation.images.{name}"] = images[name] + # Capture images from cameras + images = {} + for name in self.cameras: + before_camread_t = time.perf_counter() + images[name] = self.cameras[name].read() #Reachy cameras read() is not blocking? + # print(f'name: {name} img: {images[name]}') + if images[name] is not None: + images[name] = torch.from_numpy(copy(images[name][0])) #seems like I need to copy? + self.logs[f"read_camera_{name}_dt_s"] = images[name][1] #full timestamp, TODO dt - return obs_dict + # Populate output dictionnaries + obs_dict = {} + obs_dict["observation.state"] = state + for name in self.cameras: + obs_dict[f"observation.images.{name}"] = images[name] + return obs_dict + else: + return {} def send_action(self, action: torch.Tensor) -> torch.Tensor: if not self.is_connected: raise ConnectionError() - #TODO + #TODO: what shape is the action tensor? return action def print_logs(self) -> None: @@ -146,10 +153,10 @@ class ReachyRobot(): def disconnect(self) -> None: + print("Disconnecting") self.is_connected = False + print("Turn off") self.reachy.turn_off_smoothly() # self.reachy.turn_off() - time.sleep(6) - - def __del__(self): - self.disconnect() + print("\t turn off done") + self.reachy.disconnect() diff --git a/lerobot/scripts/test_reachy.py b/lerobot/scripts/test_reachy.py index 97e1cca8..79cb10d8 100644 --- a/lerobot/scripts/test_reachy.py +++ b/lerobot/scripts/test_reachy.py @@ -34,6 +34,8 @@ if __name__ == '__main__': robot_cfg = init_hydra_config(robot_path, robot_overrides) robot = make_robot(robot_cfg) - print(robot.get_state()) + print(robot.is_connected) + # print(robot.get_state()) print(robot.capture_observation()) - time.sleep(6) + time.sleep(5) + robot.disconnect()