import yaml import cv2 import numpy as np import collections import dm_env import argparse from typing import Dict, List, Any, Optional from collections import deque import rospy from cv_bridge import CvBridge from std_msgs.msg import Header from sensor_msgs.msg import Image, JointState from nav_msgs.msg import Odometry from rosrobot import Robot import torch import time class AgilexRobot(Robot): def get_frame(self) -> Optional[Dict[str, Any]]: """ 获取同步帧数据 返回: 包含同步数据的字典，或None如果同步失败 """ # 检查基本数据可用性 # print(self.sync_img_queues.values()) if any(len(q) == 0 for q in self.sync_img_queues.values()): print("camera has not get image data") return None if self.use_depth_image and any(len(q) == 0 for q in self.sync_depth_queues.values()): return None # print(self.sync_arm_queues.values()) # if any(len(q) == 0 for q in self.sync_arm_queues.values()): # print("2") # if len(self.sync_arm_queues['master_left']) == 0: # print("can not get data from master topic") # if len(self.sync_arm_queues['puppet_left']) == 0 or len(self.sync_arm_queues['puppet_right']) == 0: # print("can not get data from puppet topic") # return None if len(self.sync_arm_queues['puppet_left']) == 0 or len(self.sync_arm_queues['puppet_right']) == 0: print("can not get data from puppet topic") return None # 计算最小时间戳 timestamps = [ q[-1].header.stamp.to_sec() for q in list(self.sync_img_queues.values()) + list(self.sync_arm_queues.values()) ] if self.use_depth_image: timestamps.extend(q[-1].header.stamp.to_sec() for q in self.sync_depth_queues.values()) if self.use_robot_base and len(self.sync_base_queue) > 0: timestamps.append(self.sync_base_queue[-1].header.stamp.to_sec()) min_time = min(timestamps) # 检查数据同步性 for queue in list(self.sync_img_queues.values()) + list(self.sync_arm_queues.values()): if queue[-1].header.stamp.to_sec() < min_time: return None if self.use_depth_image: for queue in self.sync_depth_queues.values(): if queue[-1].header.stamp.to_sec() < min_time: return None if self.use_robot_base and len(self.sync_base_queue) > 0: if self.sync_base_queue[-1].header.stamp.to_sec() < min_time: return None # 提取同步数据 frame_data = { 'images': {}, 'arms': {}, 'timestamp': min_time } # 图像数据 for cam_name, queue in self.sync_img_queues.items(): while queue[0].header.stamp.to_sec() < min_time: queue.popleft() frame_data['images'][cam_name] = self.bridge.imgmsg_to_cv2(queue.popleft(), 'passthrough') # 深度数据 if self.use_depth_image: frame_data['depths'] = {} for cam_name, queue in self.sync_depth_queues.items(): while queue[0].header.stamp.to_sec() < min_time: queue.popleft() depth_img = self.bridge.imgmsg_to_cv2(queue.popleft(), 'passthrough') # 保持原有的边界填充 frame_data['depths'][cam_name] = cv2.copyMakeBorder( depth_img, 40, 40, 0, 0, cv2.BORDER_CONSTANT, value=0 ) # 机械臂数据 for arm_name, queue in self.sync_arm_queues.items(): while queue[0].header.stamp.to_sec() < min_time: queue.popleft() frame_data['arms'][arm_name] = queue.popleft() # 基座数据 if self.use_robot_base and len(self.sync_base_queue) > 0: while self.sync_base_queue[0].header.stamp.to_sec() < min_time: self.sync_base_queue.popleft() frame_data['base'] = self.sync_base_queue.popleft() return frame_data def teleop_step(self) -> Optional[tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]]: """ 获取同步帧数据，输出格式与 teleop_step 一致 返回: (obs_dict, action_dict) 或 None 如果同步失败 """ # 检查基本数据可用性 if any(len(q) == 0 for q in self.sync_img_queues.values()): return None, None if self.use_depth_image and any(len(q) == 0 for q in self.sync_depth_queues.values()): return None, None if any(len(q) == 0 for q in self.sync_arm_queues.values()): return None, None # 计算最小时间戳 timestamps = [ q[-1].header.stamp.to_sec() for q in list(self.sync_img_queues.values()) + list(self.sync_arm_queues.values()) ] if self.use_depth_image: timestamps.extend(q[-1].header.stamp.to_sec() for q in self.sync_depth_queues.values()) if self.use_robot_base and len(self.sync_base_queue) > 0: timestamps.append(self.sync_base_queue[-1].header.stamp.to_sec()) min_time = min(timestamps) # 检查数据同步性 for queue in list(self.sync_img_queues.values()) + list(self.sync_arm_queues.values()): if queue[-1].header.stamp.to_sec() < min_time: return None, None if self.use_depth_image: for queue in self.sync_depth_queues.values(): if queue[-1].header.stamp.to_sec() < min_time: return None, None if self.use_robot_base and len(self.sync_base_queue) > 0: if self.sync_base_queue[-1].header.stamp.to_sec() < min_time: return None, None # 初始化输出字典 obs_dict = {} action_dict = {} # 处理图像数据 for cam_name, queue in self.sync_img_queues.items(): while queue[0].header.stamp.to_sec() < min_time: queue.popleft() img = self.bridge.imgmsg_to_cv2(queue.popleft(), 'passthrough') obs_dict[f"observation.images.{cam_name}"] = torch.from_numpy(img) # 处理深度数据 if self.use_depth_image: for cam_name, queue in self.sync_depth_queues.items(): while queue[0].header.stamp.to_sec() < min_time: queue.popleft() depth_img = self.bridge.imgmsg_to_cv2(queue.popleft(), 'passthrough') depth_img = cv2.copyMakeBorder( depth_img, 40, 40, 0, 0, cv2.BORDER_CONSTANT, value=0 ) obs_dict[f"observation.images.depth_{cam_name}"] = torch.from_numpy(depth_img).unsqueeze(-1) # 处理机械臂观测数据 arm_states = [] arm_velocity = [] arm_effort = [] actions = [] for arm_name, queue in self.sync_arm_queues.items(): while queue[0].header.stamp.to_sec() < min_time: queue.popleft() arm_data = queue.popleft() # np.array(arm_data.position), # np.array(arm_data.velocity), # np.array(arm_data.effort) # 如果是从臂(puppet)，作为观测 if arm_name.startswith('puppet'): arm_states.append(np.array(arm_data.position, dtype=np.float32)) arm_velocity.append(np.array(arm_data.velocity, dtype=np.float32)) arm_effort.append(np.array(arm_data.effort, dtype=np.float32)) # 如果是主臂(master)，作为动作 if arm_name.startswith('master'): # action_dict[f"action.{arm_name}"] = torch.from_numpy(np.array(arm_data.position)) actions.append(np.array(arm_data.position, dtype=np.float32)) if arm_states: obs_dict["observation.state"] = torch.tensor(np.concatenate(arm_states).reshape(-1)) # 先转Python列表 if arm_velocity: obs_dict["observation.velocity"] = torch.tensor(np.concatenate(arm_velocity).reshape(-1)) if arm_effort: obs_dict["observation.effort"] = torch.tensor(np.concatenate(arm_effort).reshape(-1)) if actions: action_dict["action"] = torch.tensor(np.concatenate(actions).reshape(-1)) # action_dict["action"] = np.concatenate(actions).squeeze() # 处理基座数据 if self.use_robot_base and len(self.sync_base_queue) > 0: while self.sync_base_queue[0].header.stamp.to_sec() < min_time: self.sync_base_queue.popleft() base_data = self.sync_base_queue.popleft() obs_dict["observation.base_vel"] = torch.tensor([ base_data.twist.twist.linear.x, base_data.twist.twist.angular.z ], dtype=torch.float32) # 添加时间戳 # obs_dict["observation.timestamp"] = torch.tensor(min_time, dtype=torch.float64) return obs_dict, action_dict def capture_observation(self): """Capture observation data from ROS topics without batch dimension. Returns: dict: Observation dictionary containing state and images. Raises: RobotDeviceNotConnectedError: If robot is not connected. """ # Initialize observation dictionary obs_dict = {} # Get synchronized frame data frame_data = self.get_frame() if frame_data is None: # raise RuntimeError("Failed to capture synchronized observation data") return None # Process arm state data (from puppet arms) arm_states = [] arm_velocity = [] arm_effort = [] for arm_name, joint_state in frame_data['arms'].items(): if arm_name.startswith('puppet'): # Record timing for performance monitoring before_read_t = time.perf_counter() # Get position data and convert to tensor pos = torch.from_numpy(np.array(joint_state.position, dtype=np.float32)) arm_states.append(pos) velocity = torch.from_numpy(np.array(joint_state.velocity, dtype=np.float32)) arm_velocity.append(velocity) effort = torch.from_numpy(np.array(joint_state.effort, dtype=np.float32)) arm_effort.append(effort) # Log timing information # self.logs[f"read_arm_{arm_name}_pos_dt_s"] = time.perf_counter() - before_read_t print(f"read_arm_{arm_name}_pos_dt_s is", time.perf_counter() - before_read_t) # Combine all arm states into single tensor if arm_states: obs_dict["observation.state"] = torch.cat(arm_states) if arm_velocity: obs_dict["observation.velocity"] = torch.cat(arm_velocity) if arm_effort: obs_dict["observation.effort"] = torch.cat(arm_effort) # Process image data for cam_name, img in frame_data['images'].items(): # Record timing for performance monitoring before_camread_t = time.perf_counter() # Convert image to tensor img_tensor = torch.from_numpy(img) obs_dict[f"observation.images.{cam_name}"] = img_tensor # Log timing information # self.logs[f"read_camera_{cam_name}_dt_s"] = time.perf_counter() - before_camread_t print(f"read_camera_{cam_name}_dt_s is", time.perf_counter() - before_camread_t) # Process depth data if enabled if self.use_depth_image and 'depths' in frame_data: for cam_name, depth_img in frame_data['depths'].items(): before_depthread_t = time.perf_counter() # Convert depth image to tensor and add channel dimension depth_tensor = torch.from_numpy(depth_img).unsqueeze(-1) obs_dict[f"observation.images.depth_{cam_name}"] = depth_tensor # self.logs[f"read_depth_{cam_name}_dt_s"] = time.perf_counter() - before_depthread_t print(f"read_depth_{cam_name}_dt_s is", time.perf_counter() - before_depthread_t) # Process base velocity if enabled if self.use_robot_base and 'base' in frame_data: base_data = frame_data['base'] obs_dict["observation.base_vel"] = torch.tensor([ base_data.twist.twist.linear.x, base_data.twist.twist.angular.z ], dtype=torch.float32) return obs_dict def send_action(self, action: torch.Tensor) -> torch.Tensor: """ Send joint position commands to the puppet arms via ROS. Args: action: Tensor containing concatenated goal positions for all puppet arms Shape should match the action space defined in features["action"] Returns: The actual action that was sent (may be clipped if safety checks are implemented) """ # if not hasattr(self, 'puppet_arm_publishers'): # # Initialize publishers on first call # self._init_action_publishers() last_velocity = [-0.010990142822265625, -0.010990142822265625, -0.03296661376953125, 0.010990142822265625, -0.010990142822265625, -0.010990142822265625, -0.010990142822265625, -0.010990142822265625, -0.010990142822265625, -0.03296661376953125, -0.010990142822265625, -0.010990142822265625, -0.03296661376953125, -0.03296661376953125] last_effort = [-0.021978378295898438, 0.2417583465576172, 4.320878982543945, 3.6527481079101562, -0.013187408447265625, -0.013187408447265625, 0.0, -0.010990142822265625, -0.010990142822265625, -0.03296661376953125, -0.010990142822265625, -0.010990142822265625, -0.03296661376953125, -0.03296661376953125] # Convert tensor to numpy array if needed if isinstance(action, torch.Tensor): action = action.detach().cpu().numpy() # Split action into individual arm commands based on config from_idx = 0 to_idx = 0 action_sent = [] for arm_name, arm_config in self.arms.items(): # 主臂topic是否存在 if not "master" in arm_name: continue # Get number of joints for this arm num_joints = len(arm_config.get('motors', [])) to_idx += num_joints # Extract this arm's portion of the action arm_action = action[from_idx:to_idx] arm_velocity = last_velocity[from_idx:to_idx] arm_effort = last_effort[from_idx:to_idx] from_idx = to_idx # Apply safety checks if configured if 'max_relative_target' in self.config: # Get current position from the queue if len(self.sync_arm_queues[arm_name]) > 0: current_state = self.sync_arm_queues[arm_name][-1] current_pos = np.array(current_state.position) # Clip the action to stay within max relative target max_delta = self.config['max_relative_target'] clipped_action = np.clip(arm_action, current_pos - max_delta, current_pos + max_delta) arm_action = clipped_action action_sent.append(arm_action) # Create and publish JointState message joint_state = JointState() joint_state.header = Header() joint_state.header.stamp = rospy.Time.now() joint_state.name = ['joint0', 'joint1', 'joint2', 'joint3', 'joint4', 'joint5', ''] joint_state.position = arm_action.tolist() joint_state.velocity = arm_velocity joint_state.effort = arm_effort # Publish to the corresponding topic self.publishers[f"arm_{arm_name}"].publish(joint_state) return torch.from_numpy(np.concatenate(action_sent)) if action_sent else torch.tensor([]) # def _init_action_publishers(self) -> None: # """Initialize ROS publishers for puppet arms""" # self.puppet_arm_publishers = {} # # rospy.init_node("replay_node") # for arm_name, arm_config in self.arms.items(): # if not "puppet" in arm_name: # # if not "master" in arm_name: # continue # if 'topic_name' not in arm_config: # rospy.logwarn(f"No puppet topic defined for arm {arm_name}") # continue # self.puppet_arm_publishers[arm_name] = rospy.Publisher( # arm_config['topic_name'], # JointState, # queue_size=10 # ) # # Wait for publisher to connect # rospy.sleep(0.1) # rospy.loginfo("Initialized puppet arm publishers") # def get_arguments() -> argparse.Namespace: # """获取运行时参数""" # parser = argparse.ArgumentParser() # parser.add_argument('--fps', type=int, help='Frame rate', default=30) # parser.add_argument('--max_timesteps', type=int, help='Max timesteps', default=500) # parser.add_argument('--episode_idx', type=int, help='Episode index', default=0) # parser.add_argument('--use_depth', action='store_true', help='Use depth images') # parser.add_argument('--use_base', action='store_true', help='Use robot base') # return parser.parse_args() # if __name__ == "__main__": # # 示例用法 # import json # args = get_arguments() # robot = AgilexRobot(config_file="/home/jgl20/LYT/work/collect_data_lerobot/1.yaml", args=args) # print(json.dumps(robot.features, indent=4)) # robot.warmup_record() # count = 0 # print_flag = True # rate = rospy.Rate(args.fps) # while (count < args.max_timesteps + 1) and not rospy.is_shutdown(): # a, b = robot.teleop_step() # if a is None or b is None: # if print_flag: # print("syn fail\n") # print_flag = False # rate.sleep() # continue # else: # print(a) # timesteps, actions = robot.process() # print(timesteps) print()