diff --git a/configurations/basic_configs/example/robot/airbots/mmk/airbot_com_mmk_bson.yaml b/configurations/basic_configs/example/robot/airbots/mmk/airbot_com_mmk_bson.yaml new file mode 100644 index 0000000..726eded --- /dev/null +++ b/configurations/basic_configs/example/robot/airbots/mmk/airbot_com_mmk_bson.yaml @@ -0,0 +1,47 @@ +_target_: robots.airbots.airbot_mmk.airbot_com_mmk2_bson.AIRBOTMMK2 + +ip: "192.168.11.200" + +default_action: null + +# The order is important +components: ["left_arm", "left_arm_eef", "right_arm", "right_arm_eef", "head", "spine", "base"] + +# USB cameras (only color image types) +cameras: + head_camera: + image_types: ["color"] + camera_type: "USB" + video_device: "/dev/video0" + image_width: "640" + image_height: "480" + # left_camera: + # image_types: ["color"] + # camera_type: "USB" + # video_device: "/dev/video2" + # image_width: "640" + # image_height: "480" + # right_camera: + # image_types: ["color"] + # camera_type: "USB" + # video_device: "/dev/video4" + # image_width: "640" + # image_height: "480" + +# # RealSense +# cameras: +# head_camera: +# image_types: ["color"] +# camera_type: "REALSENSE" +# serial_no: "'123456'" +# rgb_camera.color_profile: "640,480,30" +# left_camera: +# image_types: ["color"] +# camera_type: "REALSENSE" +# serial_no: "'123456'" +# rgb_camera.color_profile: "640,480,30" +# right_camera: +# image_types: ["color"] +# camera_type: "REALSENSE" +# serial_no: "'123456'" +# rgb_camera.color_profile: "640,480,30" diff --git a/configurations/basic_configs/example/robot/airbots/tok/airbot_tok_demonstration.yaml b/configurations/basic_configs/example/robot/airbots/tok/airbot_tok_demonstration.yaml new file mode 100644 index 0000000..24ba41c --- /dev/null +++ b/configurations/basic_configs/example/robot/airbots/tok/airbot_tok_demonstration.yaml @@ -0,0 +1,70 @@ +_target_: robots.airbots.airbot_mmk.airbot_com_mmk2_bson.AIRBOTMMK2 + +ip: "localhost" + +# NOTE: If you want to change the order of each part, be careful to change the order of each observation and action in the data conversion script simultaneously +# components: ["left_arm", "left_arm_eef"] +components: ["left_arm", "left_arm_eef", "right_arm", "right_arm_eef", "base"] + +# USB cameras (only color image types) +cameras: + head_camera: + image_types: ["color"] + camera_type: "USB" + video_device: "/dev/video0" + image_width: "640" + image_height: "480" + # left_camera: + # image_types: ["color"] + # camera_type: "USB" + # video_device: "/dev/video2" + # image_width: "640" + # image_height: "480" + # right_camera: + # image_types: ["color"] + # camera_type: "USB" + # video_device: "/dev/video4" + # image_width: "640" + # image_height: "480" + +# # RealSense +# cameras: +# head_camera: +# image_types: ["color"] +# camera_type: "REALSENSE" +# serial_no: "'123456'" +# rgb_camera.color_profile: "640,480,30" +# left_camera: +# image_types: ["color"] +# camera_type: "REALSENSE" +# serial_no: "'123456'" +# rgb_camera.color_profile: "640,480,30" +# right_camera: +# image_types: ["color"] +# camera_type: "REALSENSE" +# serial_no: "'123456'" +# rgb_camera.color_profile: "640,480,30" + +demonstrate: true + +default_action: + [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + ] + # [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] diff --git a/configurations/basic_configs/example/robot/airbots/tok/airbot_tok_ptk.yaml b/configurations/basic_configs/example/robot/airbots/tok/airbot_tok_ptk.yaml index b5742d9..f18d73b 100644 --- a/configurations/basic_configs/example/robot/airbots/tok/airbot_tok_ptk.yaml +++ b/configurations/basic_configs/example/robot/airbots/tok/airbot_tok_ptk.yaml @@ -3,29 +3,45 @@ _target_: robots.airbots.airbot_mmk.airbot_com_mmk2_bson.AIRBOTMMK2 ip: "localhost" # NOTE: If you want to change the order of each part, be careful to change the order of each observation and action in the data conversion script simultaneously -components: ["left_arm", "left_arm_eef", "right_arm", "right_arm_eef"] +components: ["left_arm", "left_arm_eef", "right_arm", "right_arm_eef", "base"] +# USB cameras (only color image types) cameras: - head_camera: ["color"] - # left_camera: ["color"] - # right_camera: ["color"] + head_camera: + image_types: ["color"] + camera_type: "USB" + video_device: "/dev/video0" + image_width: "640" + image_height: "480" + # left_camera: + # image_types: ["color"] + # camera_type: "USB" + # video_device: "/dev/video2" + # image_width: "640" + # image_height: "480" + # right_camera: + # image_types: ["color"] + # camera_type: "USB" + # video_device: "/dev/video4" + # image_width: "640" + # image_height: "480" -demonstrate: true +# # RealSense +# cameras: +# head_camera: +# image_types: ["color"] +# camera_type: "REALSENSE" +# serial_no: "'123456'" +# rgb_camera.color_profile: "640,480,30" +# left_camera: +# image_types: ["color"] +# camera_type: "REALSENSE" +# serial_no: "'123456'" +# rgb_camera.color_profile: "640,480,30" +# right_camera: +# image_types: ["color"] +# camera_type: "REALSENSE" +# serial_no: "'123456'" +# rgb_camera.color_profile: "640,480,30" -default_action: - [ - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - ] +ignore_base_action: true \ No newline at end of file diff --git a/configurations/basic_configs/example/robot/airbots/tok/airbot_tok_ptk_demonstration.yaml b/configurations/basic_configs/example/robot/airbots/tok/airbot_tok_ptk_demonstration.yaml deleted file mode 100644 index b5742d9..0000000 --- a/configurations/basic_configs/example/robot/airbots/tok/airbot_tok_ptk_demonstration.yaml +++ /dev/null @@ -1,31 +0,0 @@ -_target_: robots.airbots.airbot_mmk.airbot_com_mmk2_bson.AIRBOTMMK2 - -ip: "localhost" - -# NOTE: If you want to change the order of each part, be careful to change the order of each observation and action in the data conversion script simultaneously -components: ["left_arm", "left_arm_eef", "right_arm", "right_arm_eef"] - -cameras: - head_camera: ["color"] - # left_camera: ["color"] - # right_camera: ["color"] - -demonstrate: true - -default_action: - [ - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - ] diff --git a/configurations/task_configs/config_tools/basic_configer.py b/configurations/task_configs/config_tools/basic_configer.py index a3ac548..d0bf6dd 100644 --- a/configurations/task_configs/config_tools/basic_configer.py +++ b/configurations/task_configs/config_tools/basic_configer.py @@ -145,8 +145,8 @@ def get_all_config(args: dict, stage: str): # set start joint if all_config.get("start_joint", None) is None: init_states = get_init_states(all_config["load_data"], 0) - all_config["start_action"] = init_states[0] - all_config["start_joint"] = init_states[1] + all_config["start_joint"] = init_states[0] + all_config["start_action"] = init_states[1] # set augmentors all_config["load_data"]["augmentors"]["image"] = task_funcs["image_augmentor"] all_config["augmentors_flag"] = { diff --git a/control_robot_bson.py b/control_robot_bson.py index 1effe7e..6e4c8aa 100644 --- a/control_robot_bson.py +++ b/control_robot_bson.py @@ -6,6 +6,7 @@ """ from habitats.common.robot_devices.cameras.utils import prepare_cv2_imshow + prepare_cv2_imshow() import argparse @@ -570,154 +571,8 @@ def stop_recording(self): timestamp = 0 start_episode_t = time.perf_counter() # Record one episode - bson_dict: Dict[str, Dict[str, list]] = { - "id": "734ad1c8-66ee-4479-b3cb-41d16c9b2e22", - "timestamp": 1734076528859, - "metadata": { - "driver_version": "1.0.0", - "operator": "manual", - "station_id": "3784D4BA-87AF-47E7-B86D-42CA1904AA77", - "task": "example", - "topics": { - "/action/head/joint_state": { - "description": "", - "type": "jointstate", - "sn": "", - "firmware_version": "0.0.0", - }, - "/action/spine/joint_state": { - "description": "", - "type": "jointstate", - "sn": "", - "firmware_version": "0.0.0", - }, - "/action/left_arm/joint_state": { - "description": "replay", - "type": "jointstate", - "sn": "", - "firmware_version": "0.0.0", - }, - "/action/left_arm_eef/joint_state": { - "description": "replay", - "type": "jointstate", - "sn": "", - "firmware_version": "0.0.0", - }, - "/action/right_arm/joint_state": { - "description": "replay", - "type": "jointstate", - "sn": "", - "firmware_version": "0.0.0", - }, - "/action/right_arm_eef/joint_state": { - "description": "replay", - "type": "jointstate", - "sn": "", - "firmware_version": "0.0.0", - }, - # "/action/eef/pose": { - # "description": "", - # "type": "jointstate", - # "sn": "", - # "firmware_version": "0.0.0", - # }, - # "/action/base/joint_state": { - # "description": "slamtec-athena", - # "type": "jointstate", - # "sn": "", - # "firmware_version": "0.0.0", - # }, - "/observation/head/joint_state": { - "description": "", - "type": "jointstate", - "sn": "", - "firmware_version": "0.0.0", - }, - "/observation/spine/joint_state": { - "description": "", - "type": "jointstate", - "sn": "", - "firmware_version": "0.0.0", - }, - "/observation/left_arm/joint_state": { - "description": "airbot-play-short", - "type": "jointstate", - "sn": "", - "firmware_version": "0.0.0", - }, - "/observation/right_arm/joint_state": { - "description": "airbot-play-short", - "type": "jointstate", - "sn": "", - "firmware_version": "0.0.0", - }, - "/observation/left_arm_eef/joint_state": { - "description": "airbot-play-short", - "type": "jointstate", - "sn": "", - "firmware_version": "0.0.0", - }, - "/observation/right_arm_eef/joint_state": { - "description": "airbot-play-short", - "type": "jointstate", - "sn": "", - "firmware_version": "0.0.0", - }, - "/observation/left_arm/pose": { - "description": "", - "type": "pose", - }, - "/observation/right_arm/pose": { - "description": "", - "type": "pose", - }, - # "/observation/base/joint_state": { - # "description": "slamtec-athena", - # "type": "jointstate", - # "sn": "", - # "firmware_version": "0.0.0", - # }, - "/images/head_camera": { - "description": "DSJ-2062-309", - "type": "image", - "width": 640, - "height": 480, - "encoding": "H264", - "distortion_model": None, - "distortion_params": None, - "intrinsics": None, - "fov": 120.0, - "start_time": 1733377253041, - }, - # "/images/left_camera": { - # "description": "DSJ-2062-309", - # "type": "image", - # "width": 640, - # "height": 480, - # "encoding": "H264", - # "distortion_model": None, - # "distortion_params": None, - # "intrinsics": None, - # "fov": 120.0, - # "start_time": 1733377253041, - # }, - # "/images/right_camera": { - # "description": "DSJ-2062-309", - # "type": "image", - # "width": 640, - # "height": 480, - # "encoding": "H264", - # "distortion_model": None, - # "distortion_params": None, - # "intrinsics": None, - # "fov": 120.0, - # "start_time": 1733377253041, - # }, - }, - "version": "1.2.1", - }, - "data": {}, - } + bson_dict: Dict[str, Dict[str, list]] = robot.bson_dict + robot.update_data_meta(bson_dict, robot.capture_observation()) while timestamp < episode_time_s: start_loop_t = time.perf_counter() diff --git a/data_process/bson_to_hdf5.py b/data_process/bson_to_hdf5.py index 93c7ade..2124b2c 100644 --- a/data_process/bson_to_hdf5.py +++ b/data_process/bson_to_hdf5.py @@ -30,9 +30,11 @@ assert os.path.exists(task_dir), f"task_dir {task_dir} not exists" name_converter = { - f"/images/{raw_name}": f"/observations/images/{i}" - for i, raw_name in enumerate(camera_names) + f"/images/{raw_name}": f"/observations/images/{raw_name}" + for raw_name in camera_names } +print(f"name_converter: {name_converter}") +image_keys = [f"/images/{name}" for name in args.camera_names] if mode == "play": obs_keys_low_dim = ( @@ -48,6 +50,7 @@ "/observation/right_arm_eef/joint_state", "/observation/head/joint_state", "/observation/spine/joint_state", + "/observation/base/joint_state", ) act_keys = ( "/action/left_arm/joint_state", @@ -56,16 +59,25 @@ "/action/right_arm_eef/joint_state", "/action/head/joint_state", "/action/spine/joint_state", + "/action/base/joint_state", + ) +elif mode in ["tok", "ptk"]: + obs_keys_low_dim = ( + "/observation/left_arm/joint_state", + "/observation/left_arm_eef/joint_state", + "/observation/right_arm/joint_state", + "/observation/right_arm_eef/joint_state", + "/observation/base/joint_state", + ) + act_keys = ( + "/action/left_arm/joint_state", + "/action/left_arm_eef/joint_state", + "/action/right_arm/joint_state", + "/action/right_arm_eef/joint_state", + "/action/base/joint_state", ) - name_converter = { - f"/images/{raw_name}": f"/observations/images/{raw_name}" - for raw_name in camera_names - } else: raise ValueError(f"mode {mode} not supported") -image_keys = [f"/images/{name}" for name in args.camera_names] - -print(f"name_converter: {name_converter}") pre_process = { key: crd.Compresser("jpg", [int(cv2.IMWRITE_JPEG_QUALITY), 50], True).compress @@ -82,10 +94,12 @@ } ) -if mode == "mmk2": +if mode in ["mmk2", "tok", "ptk"]: key_filter = [ "/observation/left_arm/pose", "/observation/right_arm/pose", + # "/observation/base/joint_state", + # "/action/base/joint_state", # "action/eef/pose", # "/time", ] @@ -121,7 +135,7 @@ os.makedirs(target_dir, exist_ok=True) print(f"Try to find all episode files in {task_dir}...") -if mode == "mmk2": +if mode in ["mmk2", "tok", "ptk"]: episode_names = crd.get_files_name_by_suffix(task_dir, ".bson") elif mode == "play": episode_names = [f"{fd}/data.bson" for fd in os.listdir(task_dir)] diff --git a/data_process/mmk2_to_hdf5.py b/data_process/mmk2_to_hdf5.py deleted file mode 100644 index 7129b95..0000000 --- a/data_process/mmk2_to_hdf5.py +++ /dev/null @@ -1,134 +0,0 @@ -import sys, os - -sys.path.append(os.path.join(os.path.dirname(__file__), "..")) - -import data_process.convert_all as crd -import argparse - -parser = argparse.ArgumentParser(description="Convert mmk2 data to hdf5") - -parser.add_argument("--raw_dir", type=str, default="data/raw", help="input directory") -parser.add_argument("--output", type=str, default="data/hdf5", help="output directory") -parser.add_argument( - "-tn", "--task_name", type=str, default="example_task", help="task name" -) -parser.add_argument( - "--cameras", type=str, nargs="+", default=["0"], help="camera names" -) -parser.add_argument("-pad", "--padding", action="store_true", help="pad the hdf5 data") - -args = parser.parse_args() - -# get all low_dim data (head&spine velocity control) -task_name = args.task_name -raw_root_dir = args.raw_dir -raw_dir = f"{raw_root_dir}/{task_name}" -data = crd.raw_to_dict( - raw_dir, - ["low_dim.json"], - video_file_names=None, - flatten_mode="hdf5", - concatenater={ - "/observations/qpos": ( - "/observation/arm/left/joint_position", - "/observation/eef/left/joint_position", - "/observation/arm/right/joint_position", - "/observation/eef/right/joint_position", - ), - "/action": ( - "/action/arm/left/joint_position", - "/action/eef/left/joint_position", - "/action/arm/right/joint_position", - "/action/eef/right/joint_position", - ), - }, - key_filter=[ - "/observation/ts_diff_with_head_color_img", - "/observation/arm/left/joint_velocity", - "/observation/arm/right/joint_velocity", - "/observation/arm/left/joint_effort", - "/observation/arm/right/joint_effort", - "/observation/eef/left/joint_velocity", - "/observation/eef/right/joint_velocity", - "/observation/eef/left/joint_effort", - "/observation/eef/right/joint_effort", - "/observation/head/joint_position", - "/observation/head/joint_velocity", - "/observation/head/joint_effort", - "/observation/spine/joint_position", - "/observation/spine/joint_velocity", - "/observation/joint_states/time", - "/observation/time", - "/action/time", - "/action/arm/left/time", - "/action/arm/right/time", - "/action/head/color/time", - "/action/head/joint_position", - "/action/spine/joint_position", - "/action/base/velocity", - # "/action/head/joint_velocity", - # "/action/spine/joint_velocity" - ], -) - -import os -import cv2 - -# merge high_dim data and save -raw_dir -names = args.cameras -video_names = [f"{name}.mp4" for name in names] -target_root_dir = args.output -target_dir = f"{target_root_dir}/{task_name}" -low_dim_data = data -name_converter = {names[i]: f"/observations/images/{i}" for i in range(len(names))} -target_namer = lambda i: f"episode_{i}.hdf5" - -compresser = crd.Compresser("jpg", [int(cv2.IMWRITE_JPEG_QUALITY), 50], True) - -os.makedirs(target_dir, exist_ok=True) - -# get max episode length -episode_lens = [] -for key, low_d in low_dim_data.items(): - length = len(list(low_d.values())[0]) - episode_lens.append(length) - # if length < 200: - # print(f"{key} has length {length}") - -max_pad_length = max(episode_lens) if args.padding else None - -# save all data -episode_names = list(low_dim_data.keys()) -print(f"Episode lengths: {episode_lens}") -print(f"Max episode length: {max_pad_length}") -print(f"All episodes: {episode_names}") -print(f"episode number: {len(episode_names)}") -downsampling = 0 - - -def save_one(index, ep_name): - crd.merge_video_and_save( - low_dim_data[ep_name], - f"{raw_dir}/{ep_name}", - video_names, - crd.save_dict_to_hdf5, - name_converter, - compresser, - f"{target_dir}/" + target_namer(index), - max_pad_length, - downsampling, - ) - data.pop(ep_name) - - -# save all -from concurrent.futures import ThreadPoolExecutor - -futures = [] -with ThreadPoolExecutor(max_workers=25) as executor: - for index, ep_name in enumerate(episode_names): - # silent execution, no print - futures.append(executor.submit(save_one, index, ep_name)) - -print(f"All data saved to {target_dir}") diff --git a/data_process/raw_to_hdf5.py b/data_process/raw_to_hdf5.py index 413d71f..2c2fc28 100644 --- a/data_process/raw_to_hdf5.py +++ b/data_process/raw_to_hdf5.py @@ -96,6 +96,9 @@ def load_raw_mujoco_data(raw_dir, downsampling=0): episode_lens = [] for low_d in low_dim_data.values(): episode_lens.append(len(list(low_d.values())[0])) +print( + f"min episode length: {min(episode_lens)}, max episode length: {max(episode_lens)}" +) max_pad_length = max(episode_lens) if padding else None print(f"Episode flatten keys: {list(low_dim_data.values())[0].keys()}") @@ -134,4 +137,4 @@ def save_one(index, ep_name): # # save one data # index = 0 # ep_name = episode_names[index] -# save_one(index, ep_name) \ No newline at end of file +# save_one(index, ep_name) diff --git a/envs/airbot_com_mmk_env.py b/envs/airbot_com_mmk_env.py index 95250e5..e3e180a 100644 --- a/envs/airbot_com_mmk_env.py +++ b/envs/airbot_com_mmk_env.py @@ -1,7 +1,6 @@ from robots.airbots.airbot_mmk.airbot_com_mmk2 import AIRBOTMMK2 from robots.common import make_robot_from_yaml import time -import collections import dm_env import numpy as np @@ -13,9 +12,12 @@ def __init__(self, config_path: str): self._all_joints_num = self.robot.joint_num def set_reset_position(self, reset_position): - assert ( - len(reset_position) == self._all_joints_num - ), f"Expected {self._all_joints_num} joints, got {len(reset_position)}" + if not len(reset_position) == self._all_joints_num: + des = f"Expected {self._all_joints_num} joints, got {len(reset_position)}" + if self.robot.config.check_dim: + raise ValueError(des) + else: + print(f"Warning: {des}") self.robot.config.default_action = reset_position def reset(self, sleep_time=0): @@ -23,16 +25,15 @@ def reset(self, sleep_time=0): return self._get_obs() def _get_obs(self): - obs = collections.OrderedDict() + obs = {} obs["qpos"] = [] obs["images"] = {} raw_obs = self.robot.capture_observation() - low_dim = raw_obs["low_dim"] for comp in self.robot.components: - obs["qpos"].extend(low_dim[f"observation/{comp.value}/joint_position"]) - for camera in self.robot.cameras: + obs["qpos"].extend(raw_obs[f"/observation/{comp.value}/joint_state"]["data"]["pos"]) + for camera in self.robot.cameras_goal: assert camera not in obs["images"], f"Duplicate camera name: {camera}" - obs["images"][camera.value] = raw_obs[f"observation.images.{camera.value}"] + obs["images"][camera.value] = raw_obs[f"/images/{camera.value}"]["data"] return dm_env.TimeStep( step_type=dm_env.StepType.FIRST, @@ -47,13 +48,14 @@ def step( sleep_time=0, get_obs=True, ): + # TODO: require the arms to be first components joint_limits = ( - (-3.09, 2.04), - (-2.92, 0.12), - (-0.04, 3.09), - (-2.95, 2.95), # (-3.1, 3.1) - (-1.9, 1.9), # (-1.08, 1.08), - (-2.90, 2.90), # (-3.0, 3.0) + (-3.151, 2.09), # (-3.09, 2.04) + (-2.963, 0.181), # (-2.92, 0.12) + (-0.094, 3.161), # (-0.04, 3.09) + (-2.95, 2.95), # (-3.012, 3.012) + (-1.9, 1.9), # (-1.859, 1.859) + (-2.90, 2.90), # (-3.017, 3.017) (0, 1), ) diff --git a/robots/airbots/airbot_mmk/airbot_com_mmk2.py b/robots/airbots/airbot_mmk/airbot_com_mmk2.py index 1f65bae..7520de9 100644 --- a/robots/airbots/airbot_mmk/airbot_com_mmk2.py +++ b/robots/airbots/airbot_mmk/airbot_com_mmk2.py @@ -1,10 +1,10 @@ from airbot_py.airbot_mmk2 import AirbotMMK2 from mmk2_types.types import ( - MMK2Components, + RobotComponents, JointNames, ComponentTypes, TopicNames, - MMK2ComponentsGroup, + RobotComponentsGroup, ImageTypes, ControllerTypes, ) @@ -15,8 +15,11 @@ MoveServoParams, TrackingParams, ForwardPositionParams, + Pose3D, + Twist3D, + BaseControlParams, ) -from typing import Optional, Dict, List, Tuple +from typing import Optional, Dict, List, Tuple, Any from dataclasses import dataclass, replace, field import time import logging @@ -27,28 +30,30 @@ @dataclass -class AIRBOTMMK2Config(object): +class AIRBOTRobotConfig(object): name: str = "mmk2" domain_id: int = -1 ip: str = "192.168.11.200" port: int = 50055 default_action: Optional[List[float]] = None - cameras: Dict[str, List[str]] = field(default_factory=lambda: {}) + cameras: Dict[str, Dict[str, Any]] = field(default_factory=dict) components: List[str] = field( default_factory=lambda: [ - MMK2Components.LEFT_ARM.value, - MMK2Components.LEFT_ARM_EEF.value, - MMK2Components.RIGHT_ARM.value, - MMK2Components.RIGHT_ARM_EEF.value, + RobotComponents.LEFT_ARM.value, + RobotComponents.LEFT_ARM_EEF.value, + RobotComponents.RIGHT_ARM.value, + RobotComponents.RIGHT_ARM_EEF.value, ] ) demonstrate: bool = False + check_dim: bool = True + ignore_base_action: bool = False class AIRBOTMMK2(object): - def __init__(self, config: Optional[AIRBOTMMK2Config] = None, **kwargs) -> None: + def __init__(self, config: Optional[AIRBOTRobotConfig] = None, **kwargs) -> None: if config is None: - config = AIRBOTMMK2Config() + config = AIRBOTRobotConfig() self.config = replace(config, **kwargs) self.robot = AirbotMMK2( self.config.ip, @@ -57,46 +62,47 @@ def __init__(self, config: Optional[AIRBOTMMK2Config] = None, **kwargs) -> None: self.config.domain_id, ) self.joint_names = {} - self.cameras: Dict[MMK2Components, str] = {} - self.components: Dict[MMK2Components, ComponentTypes] = {} - all_joint_names = JointNames() + self.cameras_goal: Dict[RobotComponents, List[ImageTypes]] = {} + self.cameras_cfg: Dict[RobotComponents, Dict[str, str]] = {} + self.cameras = self.cameras_goal.keys() + self.components: Dict[RobotComponents, ComponentTypes] = {} self.joint_num = 0 - for k, types in self.config.cameras.items(): - self.cameras[MMK2Components(k)] = [ImageTypes(v) for v in types] + for k, cfg in self.config.cameras.items(): + comp = RobotComponents(k) + types = {ImageTypes(v) for v in cfg.pop("image_types")} + self.cameras_goal[comp] = types + if types != {ImageTypes.COLOR}: + cfg["enable_depth"] = "true" + if ImageTypes.ALIGNED_DEPTH_TO_COLOR in types: + cfg["align_depth.enable"] = "true" + self.cameras_cfg[comp] = cfg for comp_str in self.config.components: - comp = MMK2Components(comp_str) + comp = RobotComponents(comp_str) # TODO: get the type info from SDK self.components[comp] = ComponentTypes.UNKNOWN - names = all_joint_names.__dict__[comp_str] + names = JointNames[comp.name].value + if comp == RobotComponents.BASE: + # TODO: fix base control + names.append("base") self.joint_names[comp] = names self.joint_num += len(names) logger.info(f"Components: {self.components}") logger.info(f"Joint numbers: {self.joint_num}") - self.robot.enable_resources( - { - comp: { - "rgb_camera.color_profile": "640,480,30", - "enable_depth": "false", - } - for comp in self.cameras - } - ) + logger.info(f"enable resources cfg: {self.cameras_cfg}") + self.robot.enable_resources(self.cameras_cfg) # use stream to get images # self.robot.enable_stream(self.robot.get_image, self.cameras) + comp_action_topic = {} if self.config.demonstrate: - comp_action_topic = { - comp: TopicNames.tracking.format(component=comp.value) - for comp in MMK2ComponentsGroup.ARMS - } - comp_action_topic.update( - { - comp: TopicNames.controller_command.format( - component=comp.value, - controller=ControllerTypes.FORWARD_POSITION.value, + for comp in self.components: + if comp in RobotComponentsGroup.ARMS: + comp_action_topic[comp] = TopicNames.tracking.format( + component=comp.value + ) + elif comp in RobotComponentsGroup.HEAD_SPINE: + comp_action_topic[comp] = TopicNames.controller_command.format( + controller=f"/{comp.value}_{ControllerTypes.FORWARD_POSITION.value}_controller" ) - for comp in MMK2ComponentsGroup.HEAD_SPINE - } - ) self.robot.listen_to(list(comp_action_topic.values())) self._comp_action_topic = comp_action_topic self.logs = {} @@ -111,23 +117,17 @@ def __init__(self, config: Optional[AIRBOTMMK2Config] = None, **kwargs) -> None: self.reset() def _move_by_traj(self, goal: dict): - # goal.update( - # { - # MMK2Components.HEAD: JointState(position=[0, -1.0]), - # MMK2Components.SPINE: JointState(position=[0.15]), - # } - # ) if self.config.demonstrate: # TODO: since the arms and eefs are controlled by the teleop bag - for comp in MMK2ComponentsGroup.ARMS_EEFS: - goal.pop(comp) + for comp in RobotComponentsGroup.ARMS_EEFS: + goal.pop(comp, None) if goal: # start = time.time() # logger.info(f"Move by trajectory") self.robot.set_goal(goal, TrajectoryParams()) # logger.info(f"Move by trajectory time: {time.time() - start}") self.robot.set_goal(goal, ForwardPositionParams()) - + # logger.info(f"Move by trajectory: {goal}") return goal def reset(self, sleep_time=0): @@ -148,19 +148,26 @@ def send_action(self, action, wait=False): # logger.info(f"Send goal: {goal}") # param = MoveServoParams(header=self.robot.get_header()) + if self.config.ignore_base_action: + goal.pop(RobotComponents.BASE, None) if self.traj_mode: self._move_by_traj(goal) else: - param = ForwardPositionParams() + param = {} + for comp in goal: + if comp is RobotComponents.BASE: + param[comp] = BaseControlParams() + else: + param[comp] = ForwardPositionParams() # param = TrackingParams() # param = MoveServoParams(header=self.robot.get_header()) # param = { - # MMK2Components.LEFT_ARM: MoveServoParams(header=self.robot.get_header()), - # MMK2Components.RIGHT_ARM: MoveServoParams(header=self.robot.get_header()), - # MMK2Components.LEFT_ARM_EEF: TrajectoryParams(), - # MMK2Components.RIGHT_ARM_EEF: TrajectoryParams(), - # MMK2Components.HEAD: ForwardPositionParams(), - # MMK2Components.SPINE: ForwardPositionParams(), + # RobotComponents.LEFT_ARM: MoveServoParams(header=self.robot.get_header()), + # RobotComponents.RIGHT_ARM: MoveServoParams(header=self.robot.get_header()), + # RobotComponents.LEFT_ARM_EEF: TrajectoryParams(), + # RobotComponents.RIGHT_ARM_EEF: TrajectoryParams(), + # RobotComponents.HEAD: ForwardPositionParams(), + # RobotComponents.SPINE: ForwardPositionParams(), # } self.robot.set_goal(goal, param) @@ -174,7 +181,7 @@ def get_low_dim_data(self): all_joints, self.joint_names[comp] ) data[f"observation/{comp.value}/joint_position"] = joint_states - if comp == MMK2Components.BASE: + if comp == RobotComponents.BASE: base_pose = robot_state.base_state.pose base_vel = robot_state.base_state.velocity data_pose = [ @@ -191,7 +198,7 @@ def get_low_dim_data(self): data[f"action/{comp.value}/velocity"] = data_vel data[f"action/{comp.value}/joint_position"] = data_vel + data_pose if self.config.demonstrate: - if comp in MMK2ComponentsGroup.ARMS: + if comp in RobotComponentsGroup.ARMS: arm_jn = JointNames().__dict__[comp.value] comp_eef = comp.value + "_eef" eef_jn = JointNames().__dict__[comp_eef] @@ -200,7 +207,7 @@ def get_low_dim_data(self): data[f"action/{comp.value}/joint_position"] = jq[:-1] # the eef joint is in arms data[f"action/{comp_eef}/joint_position"] = jq[-1:] - elif comp in MMK2ComponentsGroup.HEAD_SPINE: + elif comp in RobotComponentsGroup.HEAD_SPINE: jq = list( self.robot.get_listened(self._comp_action_topic[comp]).data ) @@ -209,9 +216,9 @@ def get_low_dim_data(self): def _capture_images(self) -> Tuple[Dict[str, bytes], Dict[str, Time]]: images = {} - img_stamps: Dict[MMK2Components, Time] = {} + img_stamps: Dict[RobotComponents, Time] = {} before_camread_t = time.perf_counter() - comp_images = self.robot.get_image(self.cameras) + comp_images = self.robot.get_image(self.cameras_goal) for comp, image in comp_images.items(): # TODO: now only support for color image images[comp.value] = image.data[ImageTypes.COLOR] @@ -247,7 +254,7 @@ def low_dim_to_action(self, low_dim: dict, step: int) -> list: for comp in self.components: # action.extend(low_dim[f"action/{comp.value}/joint_position"][step]) # old version - if comp in MMK2ComponentsGroup.ARMS_EEFS: + if comp in RobotComponentsGroup.ARMS_EEFS: pos_comp = comp.value.split("_") key = f"{pos_comp[1]}/{pos_comp[0]}" else: @@ -263,13 +270,18 @@ def _action_check(self, action): len(action) == self.joint_num ), f"Invalid action {action} with length: {len(action)}" - def _action_to_goal(self, action) -> Dict[MMK2Components, JointState]: - self._action_check(action) + def _action_to_goal(self, action) -> Dict[RobotComponents, JointState]: + if self.config.check_dim: + self._action_check(action) goal = {} j_cnt = 0 for comp in self.components: end = j_cnt + len(self.joint_names[comp]) - goal[comp] = JointState(position=action[j_cnt:end]) + if comp is RobotComponents.BASE: + x, y, omega = action[j_cnt:end] + goal[comp] = Twist3D(x=x, y=y, omega=omega) + else: + goal[comp] = JointState(position=action[j_cnt:end]) j_cnt = end return goal diff --git a/robots/airbots/airbot_mmk/airbot_com_mmk2_bson.py b/robots/airbots/airbot_mmk/airbot_com_mmk2_bson.py index 61df8f4..52d64d9 100644 --- a/robots/airbots/airbot_mmk/airbot_com_mmk2_bson.py +++ b/robots/airbots/airbot_mmk/airbot_com_mmk2_bson.py @@ -1,13 +1,13 @@ from mmk2_types.types import ( JointNames, - MMK2Components, - MMK2ComponentsGroup, + RobotComponents, + RobotComponentsGroup, ) from mmk2_types.grpc_msgs import Time from typing import Optional, Dict import logging import numpy as np -from robots.airbots.airbot_mmk.airbot_com_mmk2 import AIRBOTMMK2Config +from robots.airbots.airbot_mmk.airbot_com_mmk2 import AIRBOTRobotConfig from robots.airbots.airbot_mmk.airbot_com_mmk2 import AIRBOTMMK2 as AIRBOTMMK2_BASE @@ -17,7 +17,7 @@ class AIRBOTMMK2(AIRBOTMMK2_BASE): - def __init__(self, config: Optional[AIRBOTMMK2Config] = None, **kwargs): + def __init__(self, config: Optional[AIRBOTRobotConfig] = None, **kwargs): self.images_ts: Dict[str, int] = {} super().__init__(config, **kwargs) @@ -68,7 +68,7 @@ def get_low_dim_data(self): all_joints = robot_state.joint_state for comp in self.components: stamp = all_joints.header.stamp - if comp != MMK2Components.BASE: + if comp != RobotComponents.BASE: # TODO: hand has no joint states names = self.joint_names[comp] if set(names) - set(all_joints.name): @@ -86,7 +86,7 @@ def get_low_dim_data(self): all_joints, names, "effort" ) # TODO: configure has-pose components - if comp in MMK2ComponentsGroup.ARMS: + if comp in RobotComponentsGroup.ARMS: poses = robot_state.robot_pose.robot_pose[comp.value] t = poses.position.x, poses.position.y, poses.position.z r = ( @@ -103,25 +103,18 @@ def get_low_dim_data(self): joint_pos = [base_pose.x, base_pose.y, base_pose.theta] joint_vel = [base_vel.x, base_vel.y, base_vel.omega] joint_eff = [0.0, 0.0, 0.0] - data.update( - self._get_joint_state( - "action", comp.value, stamp, joint_pos, joint_vel, joint_eff - ) - ) data.update( self._get_joint_state( "observation", comp.value, stamp, joint_pos, joint_vel, joint_eff ) ) if self.config.demonstrate: - if comp in MMK2ComponentsGroup.ARMS: - arm_jn = JointNames().__dict__[comp.value] + if comp in RobotComponentsGroup.ARMS: + arm_jn = JointNames[comp.name].value comp_eef = comp.value + "_eef" - eef_jn = JointNames().__dict__[comp_eef] + eef_jn = JointNames[RobotComponents(comp_eef).name].value js = self.robot.get_listened(self._comp_action_topic[comp]) - assert ( - js is not None - ), "The AIRBOT MMK2 should be in bag teleopration mode." + assert js is not None, "The robot should be in teleopration mode." jq = self.robot.get_joint_values_by_names(js, arm_jn + eef_jn) data.update( self._get_joint_state( @@ -133,7 +126,14 @@ def get_low_dim_data(self): "action", comp_eef, js.header.stamp, jq[-1:] ) ) - elif comp in MMK2ComponentsGroup.HEAD_SPINE: + elif comp == RobotComponents.BASE: + # TODO: now action and observation are the same for base + data.update( + self._get_joint_state( + "action", comp.value, stamp, joint_pos, joint_vel, joint_eff + ) + ) + elif comp in RobotComponentsGroup.HEAD_SPINE: result = self.robot.get_listened(self._comp_action_topic[comp]) assert ( result is not None @@ -153,6 +153,71 @@ def capture_observation(self): data.update(self._get_image(name, img_stamps[name], img)) return data + def update_data_meta(self, bson_dict: dict, observation: dict): + topics = bson_dict["metadata"]["topics"] + for camera in self.cameras: + image = observation[f"/images/{camera.value}"]["data"] + image_meta = topics[f"/images/{camera.value}"] + h, w = image.shape[:2] + image_meta["width"] = w + image_meta["height"] = h + + def low_dim_to_action(self, low_dim: dict, step: int) -> list: + action = [] + for comp in self.components: + action.extend(low_dim[f"/action/{comp.value}/joint_state"]["data"][step]) + return action + + @property + def bson_dict(self): + bson_dict = { + "id": "734ad1c8-66ee-4479-b3cb-41d16c9b2e22", + "timestamp": 1734076528859, + "metadata": { + "driver_version": "1.0.0", + "operator": "manual", + "station_id": "3784D4BA-87AF-47E7-B86D-42CA1904AA77", + "task": "example", + "version": "1.2.1", + "topics": {}, + }, + "data": {}, + } + topics = bson_dict["metadata"]["topics"] + for comp in self.components: + for tp in ["action", "observation"]: + topics[f"/{tp}/{comp.value}/joint_state"] = { + "description": "", + "type": "jointstate", + "sn": "", + "firmware_version": "0.0.0", + } + if comp in RobotComponentsGroup.ARMS: + topics[f"/observation/{comp.value}/pose"] = { + "description": "", + "type": "pose", + } + for camera in self.cameras: + topics[f"/images/{camera.value}"] = { + "description": "DSJ-2062-309", + "type": "image", + "width": 640, + "height": 480, + "encoding": "H264", + "distortion_model": None, + "distortion_params": None, + "intrinsics": None, + "fov": 120.0, + "start_time": 1733377253041, + } + # "/action/eef/pose": { + # "description": "", + # "type": "jointstate", + # "sn": "", + # "firmware_version": "0.0.0", + # }, + return bson_dict + def main(): robot = AIRBOTMMK2() diff --git a/robots/airbots/airbot_mmk/airbot_mmk2.py b/robots/airbots/airbot_mmk/airbot_mmk2.py index 3915ef5..c70fc82 100644 --- a/robots/airbots/airbot_mmk/airbot_mmk2.py +++ b/robots/airbots/airbot_mmk/airbot_mmk2.py @@ -1,10 +1,10 @@ -from mmk2_sdk.mmk2_client import AIRBOTMMK2 as AIRBOTMMK2Client +from mmk2_sdk.mmk2_client import AIRBOTMMK2 as AIRBOTRobotClient from mmk2_types.types import ( - MMK2Components, + RobotComponents, JointNames, ComponentTypes, TopicNames, - MMK2ComponentsGroup, + RobotComponentsGroup, ImageTypes, ControllerTypes, ) @@ -26,7 +26,7 @@ @dataclass -class AIRBOTMMK2Config(object): +class AIRBOTRobotConfig(object): name: str = "mmk2" domain_id: int = -1 ip: str = "192.168.11.200" @@ -35,35 +35,35 @@ class AIRBOTMMK2Config(object): cameras: Dict[str, str] = field(default_factory=lambda: {}) components: List[str] = field( default_factory=lambda: [ - MMK2Components.LEFT_ARM.value, - MMK2Components.LEFT_ARM_EEF.value, - MMK2Components.RIGHT_ARM.value, - MMK2Components.RIGHT_ARM_EEF.value, + RobotComponents.LEFT_ARM.value, + RobotComponents.LEFT_ARM_EEF.value, + RobotComponents.RIGHT_ARM.value, + RobotComponents.RIGHT_ARM_EEF.value, ] ) demonstrate: bool = False class AIRBOTMMK2(object): - def __init__(self, config: Optional[AIRBOTMMK2Config] = None, **kwargs) -> None: + def __init__(self, config: Optional[AIRBOTRobotConfig] = None, **kwargs) -> None: if config is None: - config = AIRBOTMMK2Config() + config = AIRBOTRobotConfig() self.config = replace(config, **kwargs) - self.robot = AIRBOTMMK2Client( + self.robot = AIRBOTRobotClient( self.config.ip, self.config.port, self.config.name, self.config.domain_id, ) self.joint_names = {} - self.cameras: Dict[MMK2Components, str] = {} - self.components: Dict[MMK2Components, ComponentTypes] = {} + self.cameras: Dict[RobotComponents, str] = {} + self.components: Dict[RobotComponents, ComponentTypes] = {} all_joint_names = JointNames() self.joint_num = 0 for k, v in self.config.cameras.items(): - self.cameras[MMK2Components(k)] = ImageTypes(v) + self.cameras[RobotComponents(k)] = ImageTypes(v) for comp_str in self.config.components: - comp = MMK2Components(comp_str) + comp = RobotComponents(comp_str) # TODO: get the type info from SDK self.components[comp] = ComponentTypes.UNKNOWN names = all_joint_names.__dict__[comp_str] @@ -85,7 +85,7 @@ def __init__(self, config: Optional[AIRBOTMMK2Config] = None, **kwargs) -> None: if self.config.demonstrate: comp_action_topic = { comp: TopicNames.tracking.format(component=comp.value) - for comp in MMK2ComponentsGroup.ARMS + for comp in RobotComponentsGroup.ARMS } comp_action_topic.update( { @@ -93,7 +93,7 @@ def __init__(self, config: Optional[AIRBOTMMK2Config] = None, **kwargs) -> None: component=comp.value, controller=ControllerTypes.FORWARD_POSITION.value, ) - for comp in MMK2ComponentsGroup.HEAD_SPINE + for comp in RobotComponentsGroup.HEAD_SPINE } ) self.robot.listen_to(list(comp_action_topic.values())) @@ -112,13 +112,13 @@ def __init__(self, config: Optional[AIRBOTMMK2Config] = None, **kwargs) -> None: def _move_by_traj(self, goal: dict): # goal.update( # { - # MMK2Components.HEAD: JointState(position=[0, -1.0]), - # MMK2Components.SPINE: JointState(position=[0.15]), + # RobotComponents.HEAD: JointState(position=[0, -1.0]), + # RobotComponents.SPINE: JointState(position=[0.15]), # } # ) if self.config.demonstrate: # TODO: since the arms and eefs are controlled by the teleop bag - for comp in MMK2ComponentsGroup.ARMS_EEFS: + for comp in RobotComponentsGroup.ARMS_EEFS: goal.pop(comp) if goal: self.robot.set_goal(goal, TrajectoryParams()) @@ -150,12 +150,12 @@ def send_action(self, action, wait=False): param = ForwardPositionParams() # param = MoveServoParams(header=self.robot.get_header()) # param = { - # MMK2Components.LEFT_ARM: MoveServoParams(header=self.robot.get_header()), - # MMK2Components.RIGHT_ARM: MoveServoParams(header=self.robot.get_header()), - # MMK2Components.LEFT_ARM_EEF: TrajectoryParams(), - # MMK2Components.RIGHT_ARM_EEF: TrajectoryParams(), - # MMK2Components.HEAD: ForwardPositionParams(), - # MMK2Components.SPINE: ForwardPositionParams(), + # RobotComponents.LEFT_ARM: MoveServoParams(header=self.robot.get_header()), + # RobotComponents.RIGHT_ARM: MoveServoParams(header=self.robot.get_header()), + # RobotComponents.LEFT_ARM_EEF: TrajectoryParams(), + # RobotComponents.RIGHT_ARM_EEF: TrajectoryParams(), + # RobotComponents.HEAD: ForwardPositionParams(), + # RobotComponents.SPINE: ForwardPositionParams(), # } self.robot.set_goal(goal, param) @@ -169,7 +169,7 @@ def get_low_dim_data(self): ) data[f"observation/{comp.value}/joint_position"] = joint_states if self.config.demonstrate: - if comp in MMK2ComponentsGroup.ARMS: + if comp in RobotComponentsGroup.ARMS: arm_jn = JointNames().__dict__[comp.value] comp_eef = comp.value + "_eef" eef_jn = JointNames().__dict__[comp_eef] @@ -178,7 +178,7 @@ def get_low_dim_data(self): data[f"action/{comp.value}/joint_position"] = jq[:-1] # the eef joint is in arms data[f"action/{comp_eef}/joint_position"] = jq[-1:] - elif comp in MMK2ComponentsGroup.HEAD_SPINE: + elif comp in RobotComponentsGroup.HEAD_SPINE: jq = list( self.robot.get_listened(self._comp_action_topic[comp]).data ) @@ -187,7 +187,7 @@ def get_low_dim_data(self): def _capture_images(self) -> Tuple[Dict[str, bytes], Dict[str, Time]]: images = {} - img_stamps: Dict[MMK2Components, Time] = {} + img_stamps: Dict[RobotComponents, Time] = {} before_camread_t = time.perf_counter() comp_images = self.robot.get_image(self.cameras) for comp, image in comp_images.items(): @@ -232,7 +232,7 @@ def low_dim_to_action(self, low_dim: dict, step: int) -> list: for comp in self.components: # action.extend(low_dim[f"action/{comp.value}/joint_position"][step]) # old version - if comp in MMK2ComponentsGroup.ARMS_EEFS: + if comp in RobotComponentsGroup.ARMS_EEFS: pos_comp = comp.value.split("_") key = f"{pos_comp[1]}/{pos_comp[0]}" else: @@ -248,7 +248,7 @@ def _action_check(self, action): len(action) == self.joint_num ), f"Invalid action {action} with length: {len(action)}" - def _action_to_goal(self, action) -> Dict[MMK2Components, JointState]: + def _action_to_goal(self, action) -> Dict[RobotComponents, JointState]: self._action_check(action) goal = {} j_cnt = 0 diff --git a/robots/airbots/airbot_mmk/airbot_mmk2_bson.py b/robots/airbots/airbot_mmk/airbot_mmk2_bson.py index 31e5398..27f9dcc 100644 --- a/robots/airbots/airbot_mmk/airbot_mmk2_bson.py +++ b/robots/airbots/airbot_mmk/airbot_mmk2_bson.py @@ -1,12 +1,12 @@ from mmk2_types.types import ( JointNames, - MMK2ComponentsGroup, + RobotComponentsGroup, ) from mmk2_types.grpc_msgs import Time from typing import Optional, Dict import logging import numpy as np -from robots.airbots.airbot_mmk.airbot_mmk2 import AIRBOTMMK2Config +from robots.airbots.airbot_mmk.airbot_mmk2 import AIRBOTRobotConfig from robots.airbots.airbot_mmk.airbot_mmk2 import AIRBOTMMK2 as AIRBOTMMK2_BASE @@ -16,7 +16,7 @@ class AIRBOTMMK2(AIRBOTMMK2_BASE): - def __init__(self, config: Optional[AIRBOTMMK2Config] = None, **kwargs): + def __init__(self, config: Optional[AIRBOTRobotConfig] = None, **kwargs): self.images_ts: Dict[str, int] = {} super().__init__(config, **kwargs) @@ -69,7 +69,7 @@ def get_low_dim_data(self): ) ) if self.config.demonstrate: - if comp in MMK2ComponentsGroup.ARMS: + if comp in RobotComponentsGroup.ARMS: arm_jn = JointNames().__dict__[comp.value] comp_eef = comp.value + "_eef" eef_jn = JointNames().__dict__[comp_eef] @@ -86,7 +86,7 @@ def get_low_dim_data(self): "action", comp_eef, js.header.stamp, jq[-1:] ) ) - elif comp in MMK2ComponentsGroup.HEAD_SPINE: + elif comp in RobotComponentsGroup.HEAD_SPINE: result = self.robot.get_listened(self._comp_action_topic[comp]) assert result is not None, "The AIRBOT MMK2 should be in bag teleopration sync mode." jq = list(result.data) diff --git a/robots/airbots/airbot_play/airbot_play_5.py b/robots/airbots/airbot_play/airbot_play_5.py new file mode 100644 index 0000000..77dfa55 --- /dev/null +++ b/robots/airbots/airbot_play/airbot_play_5.py @@ -0,0 +1,173 @@ +from dataclasses import dataclass, field, replace +import time +from habitats.common.robot_devices.cameras.utils import Camera +from typing import Dict, Optional, List +from airbot_py.arm import AirbotPlay + + + +@dataclass +class AIRBOTPlayConfig(object): + port: int = 50000 + start_arm_joint_position: List[float] = field( + default_factory=lambda: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + ) + start_eef_joint_position: float = 0.0 + # ONLINE_TRAJ, ONLINE_IDLE, ONLINE_SERVO, DEMONSTRATE_PREP + default_robot_mode: str = "ONLINE_IDLE" + cameras: dict = field(default_factory=lambda: {}) + display: bool = False + + def __post_init__(self): + assert self.default_robot_mode in [ + "ONLINE_TRAJ", + "ONLINE_IDLE", + "ONLINE_SERVO", + "DEMONSTRATE_PREP", + ] + + +class AIRBOTPlay(object): + def __init__(self, config: Optional[AIRBOTPlayConfig] = None, **kwargs) -> None: + if config is None: + config = AIRBOTPlayConfig() + self.config = replace(config, **kwargs) + self.cameras: Dict[str, Camera] = self.config.cameras + self.logs = {} + self.__init() + self._state_mode = "active" + self._exited = False + + def __init(self): + args = self.config + # Connect the cameras + for name in self.cameras: + self.cameras[name].connect() + # Connect the robot + self.robot = AirbotPlay( + port = args.port + ) + args.arm_type = self.robot.params["arm_type"] + time.sleep(0.3) + self.reset() + + def reset(self): + args = self.config + robot = self.robot + # set to traj mode + if args.arm_type != "replay" and robot.get_current_state() != "ONLINE_TRAJ": + assert robot.online_mode(), "online idle mode failed" + # assert robot.online_traj_mode(), "online traj mode failed" + time.sleep(0.3) + # go to start position + if args.arm_type != "replay": + if args.start_arm_joint_position is not None: + assert robot.set_target_joint_q( + args.start_arm_joint_position, blocking=True + ), "set target joint q failed" + if args.start_eef_joint_position is not None and robot.params["eef_type"] not in ["none", "E2B"]: + assert robot.set_target_end( + args.start_eef_joint_position, blocking=False + ), "set target end failed" + # enter default mode + if args.default_robot_mode == "ONLINE_TRAJ": + self.enter_traj_mode() + elif args.default_robot_mode == "ONLINE_IDLE": + self.enter_active_mode() + elif args.default_robot_mode == "ONLINE_SERVO": + self.enter_servo_mode() + else: + raise ValueError( + f"Invalid default robot mode: {args.default_robot_mode}" + ) + + self._state_mode = "active" + + def enter_traj_mode(self): + self.enter_active_mode() + if self.config.arm_type == "replay": + return + # else: + # assert self.robot.online_traj_mode(), "online traj mode failed" + time.sleep(0.5) + self._state_mode = "active" + + def enter_active_mode(self): + if self.config.arm_type == "replay": + return + else: + assert self.robot.online_mode(), "online idle mode failed" + self._state_mode = "active" + + def enter_passive_mode(self): + if self.config.arm_type == "replay": + return + else: + assert self.robot.manual_mode(), "demonstrate start mode failed" + self._state_mode = "passive" + + def enter_servo_mode(self): + self.enter_active_mode() + if self.config.arm_type == "replay": + return + # else: + # assert self.robot.online_servo_mode(), "online_servo_mode mode failed" # 没有online_servo_mode + self._state_mode = "active" + + def send_action(self, action, wait=False): + assert self._state_mode == "active", "Robot is not in active mode" + if self.config.arm_type == "replay": + return + else: + assert self.robot.set_target_joint_q( + action[:6], blocking=wait, vel=0.5, use_planning=False + ), "set target joint q failed" + if self.robot.params["eef_type"] not in ["none", "E2B"]: + # assert self.robot.set_target_end(action[6]), "set target end failed" + self.robot.set_target_end(action[6], blocking=wait) + return action + + def get_low_dim_data(self): + data = {} + data["/time"] = time.time() + pose = self.robot.get_current_pose() + data["observation/arm/joint_position"] = list(self.robot.get_current_joint_q()) + data["observation/eef/joint_position"] = [self.robot.get_current_end()] + data["observation/eef/pose"] = pose[0] + pose[1] # xyz + quat(xyzw) + return data + + def capture_observation(self): + """The returned observations do not have a batch dimension.""" + obs_act_dict = {} + # Capture images from cameras + images = {} + for name in self.cameras: + before_camread_t = time.perf_counter() + images[name] = self.cameras[name].async_read() + # images[name] = torch.from_numpy(images[name]) + obs_act_dict[f"/time/{name}"] = time.time() + self.logs[f"read_camera_{name}_dt_s"] = self.cameras[name].logs[ + "delta_timestamp_s" + ] + self.logs[f"async_read_camera_{name}_dt_s"] = ( + time.perf_counter() - before_camread_t + ) + + low_dim_data = self.get_low_dim_data() + + # Populate output dictionnaries and format to pytorch + obs_act_dict["low_dim"] = low_dim_data + for name in self.cameras: + obs_act_dict[f"observation.images.{name}"] = images[name] + return obs_act_dict + + def exit(self): + assert not self._exited, "Robot already exited" + for name in self.cameras: + self.cameras[name].disconnect() + self.robot.shutdown() + self._exited = True + print("Robot exited") + + def get_state_mode(self): + return self._state_mode diff --git a/robots/common.py b/robots/common.py index 09f693c..0f00d41 100644 --- a/robots/common.py +++ b/robots/common.py @@ -38,6 +38,7 @@ def enter_servo_mode(self): ... def enter_traj_mode(self): ... def get_low_dim_data(self): ... def capture_observation(self): ... + def update_data_meta(self, bson_dict: dict, observation:dict): ... def send_action(self, action): ... def reset(self): ... def exit(self): ...