add realman shadow src

realman_src/realman_aloha/shadow_rm_act/test/test_camera.py

@@ -0,0 +1,163 @@
+from shadow_camera.realsense import RealSenseCamera
+from shadow_rm_robot.realman_arm import RmArm
+import yaml
+import time
+import multiprocessing
+import numpy as np
+import collections
+import logging
+import dm_env
+import tracemalloc
+
+
+class DeviceAloha:
+    def __init__(self, aloha_config):
+        """
+        初始化设备
+
+        Args:
+            device_name (str): 设备名称
+        """
+        config_left_arm = aloha_config["rm_left_arm"]
+        config_right_arm = aloha_config["rm_right_arm"]
+        config_head_camera = aloha_config["head_camera"]
+        config_bottom_camera = aloha_config["bottom_camera"]
+        config_left_camera = aloha_config["left_camera"]
+        config_right_camera = aloha_config["right_camera"]
+        self.init_left_arm_angle = aloha_config["init_left_arm_angle"]
+        self.init_right_arm_angle = aloha_config["init_right_arm_angle"]
+        self.arm_left = RmArm(config_left_arm)
+        self.arm_right = RmArm(config_right_arm)
+        self.camera_left = RealSenseCamera(config_head_camera, False)
+        self.camera_right = RealSenseCamera(config_bottom_camera, False)
+        self.camera_bottom = RealSenseCamera(config_left_camera, False)
+        self.camera_top = RealSenseCamera(config_right_camera, False)
+        self.camera_left.start_camera()
+        self.camera_right.start_camera()
+        self.camera_bottom.start_camera()
+        self.camera_top.start_camera()
+
+    def close(self):
+        """
+        关闭摄像头
+        """
+        self.camera_left.close()
+        self.camera_right.close()
+        self.camera_bottom.close()
+        self.camera_top.close()
+
+    def get_qps(self):
+        """
+        获取关节角度
+
+        Returns:
+            np.array: 关节角度
+        """
+        left_slave_arm_angle = self.arm_left.get_joint_angle()
+        left_joint_angles_array = np.array(list(left_slave_arm_angle.values()))
+        right_slave_arm_angle = self.arm_right.get_joint_angle()
+        right_joint_angles_array = np.array(list(right_slave_arm_angle.values()))
+        return np.concatenate([left_joint_angles_array, right_joint_angles_array])
+
+    def get_qvel(self):
+        """
+        获取关节速度
+
+        Returns:
+            np.array: 关节速度
+        """
+        left_slave_arm_velocity = self.arm_left.get_joint_velocity()
+        left_joint_velocity_array = np.array(list(left_slave_arm_velocity.values()))
+        right_slave_arm_velocity = self.arm_right.get_joint_velocity()
+        right_joint_velocity_array = np.array(list(right_slave_arm_velocity.values()))
+        return np.concatenate([left_joint_velocity_array, right_joint_velocity_array])
+
+    def get_effort(self):
+        """
+        获取关节力
+
+        Returns:
+            np.array: 关节力
+        """
+        left_slave_arm_effort = self.arm_left.get_joint_effort()
+        left_joint_effort_array = np.array(list(left_slave_arm_effort.values()))
+        right_slave_arm_effort = self.arm_right.get_joint_effort()
+        right_joint_effort_array = np.array(list(right_slave_arm_effort.values()))
+        return np.concatenate([left_joint_effort_array, right_joint_effort_array])
+
+    def get_images(self):
+        """
+        获取图像
+
+        Returns:
+            dict: 图像字典
+        """
+        top_image, _, _, _ = self.camera_top.read_frame(True, False, False, False)
+        bottom_image, _, _, _ = self.camera_bottom.read_frame(True, False, False, False)
+        left_image, _, _, _ = self.camera_left.read_frame(True, False, False, False)
+        right_image, _, _, _ = self.camera_right.read_frame(True, False, False, False)
+        return {
+            "cam_high": top_image,
+            "cam_low": bottom_image,
+            "cam_left": left_image,
+            "cam_right": right_image,
+        }
+
+    def get_observation(self):
+        obs = collections.OrderedDict()
+        obs["qpos"] = self.get_qps()
+        obs["qvel"] = self.get_qvel()
+        obs["effort"] = self.get_effort()
+        obs["images"] = self.get_images()
+        # self.get_images()
+        return obs
+
+    def reset(self):
+        logging.info("Resetting the environment")
+        _ = self.arm_left.set_joint_position(self.init_left_arm_angle[0:6])
+        _ = self.arm_right.set_joint_position(self.init_right_arm_angle[0:6])
+        self.arm_left.set_gripper_position(0)
+        self.arm_right.set_gripper_position(0)
+        return dm_env.TimeStep(
+            step_type=dm_env.StepType.FIRST,
+            reward=0,
+            discount=None,
+            observation=self.get_observation(),
+        )
+    
+    def step(self, target_angle):
+        self.arm_left.set_joint_canfd_position(target_angle[0:6])
+        self.arm_right.set_joint_canfd_position(target_angle[7:13])
+        self.arm_left.set_gripper_position(target_angle[6])
+        self.arm_right.set_gripper_position(target_angle[13])
+        return dm_env.TimeStep(
+            step_type=dm_env.StepType.MID,
+            reward=0,
+            discount=None,
+            observation=self.get_observation(),
+        )
+    
+
+if __name__ == '__main__':
+    with open("/home/rm/code/shadow_act/config/config.yaml", "r") as f:
+        config = yaml.safe_load(f)
+    aloha_config = config["robot_env"]
+    device = DeviceAloha(aloha_config)
+    device.reset()
+    image_list = [] 
+    tager_angle = np.concatenate([device.init_left_arm_angle, device.init_right_arm_angle])
+    while True:
+        tracemalloc.start()  # 启动内存跟踪
+
+        tager_angle = np.array([angle + 0.1 if i not in [6, 13] else angle for i, angle in enumerate(tager_angle)])
+        time_step = time.time()
+        timestep = device.step(tager_angle)
+        logging.info(f"Time: {time.time() - time_step}")
+        image_list.append(timestep.observation["images"])
+        snapshot = tracemalloc.take_snapshot()
+        top_stats = snapshot.statistics('lineno')
+        # del timestep
+        print("[ Top 10 ]")
+        for stat in top_stats[:10]:
+            print(stat)
+        # logging.info(f"Images: {obs}")

realman_src/realman_aloha/shadow_rm_act/test/test_h5.py

@@ -0,0 +1,32 @@
+import os
+# import time
+import yaml
+import torch
+import pickle
+import dm_env
+import logging
+import collections
+import numpy as np
+from einops import rearrange
+import matplotlib.pyplot as plt
+from torchvision import transforms
+from shadow_rm_robot.realman_arm import RmArm
+from shadow_camera.realsense import RealSenseCamera
+from shadow_act.models.latent_model import Latent_Model_Transformer
+from shadow_act.network.policy import ACTPolicy, CNNMLPPolicy, DiffusionPolicy
+from shadow_act.utils.utils import (
+    load_data,
+    sample_box_pose,
+    sample_insertion_pose,
+    compute_dict_mean,
+    set_seed,
+    detach_dict,
+)
+logging.basicConfig(level=logging.DEBUG)
+
+
+
+
+
+
+print('daasdas')