373 lines
14 KiB
Python
373 lines
14 KiB
Python
import yaml
|
|
from collections import deque
|
|
import rospy
|
|
from cv_bridge import CvBridge
|
|
from typing import Dict, Any, Optional, List
|
|
from sensor_msgs.msg import Image, JointState
|
|
from nav_msgs.msg import Odometry
|
|
import argparse
|
|
|
|
|
|
class Robot:
|
|
def __init__(self, config_file: str, args: Optional[argparse.Namespace] = None):
|
|
"""
|
|
机器人基类，处理通用初始化逻辑
|
|
Args:
|
|
config_file: YAML配置文件路径
|
|
args: 运行时参数
|
|
"""
|
|
self._load_config(config_file)
|
|
self._merge_runtime_args(args)
|
|
self._init_components()
|
|
self._init_data_structures()
|
|
self.init_ros()
|
|
self.init_features()
|
|
self.warmup()
|
|
|
|
def _load_config(self, config_file: str) -> None:
|
|
"""加载YAML配置文件"""
|
|
with open(config_file, 'r') as f:
|
|
self.config = yaml.safe_load(f)
|
|
|
|
def _merge_runtime_args(self, args: Optional[argparse.Namespace]) -> None:
|
|
"""合并运行时参数到配置"""
|
|
if args is None:
|
|
return
|
|
|
|
runtime_params = {
|
|
'frame_rate': getattr(args, 'fps', None),
|
|
'max_timesteps': getattr(args, 'max_timesteps', None),
|
|
'episode_idx': getattr(args, 'episode_idx', None),
|
|
'use_depth_image': getattr(args, 'use_depth_image', None),
|
|
'use_robot_base': getattr(args, 'use_base', None),
|
|
'video': getattr(args, 'video', False),
|
|
'control_type': getattr(args, 'control_type', False),
|
|
}
|
|
|
|
for key, value in runtime_params.items():
|
|
if value is not None:
|
|
self.config[key] = value
|
|
|
|
def _init_components(self) -> None:
|
|
"""初始化核心组件"""
|
|
self.bridge = CvBridge()
|
|
self.subscribers = {}
|
|
self.publishers = {}
|
|
self._validate_config()
|
|
|
|
def _validate_config(self) -> None:
|
|
"""验证配置完整性"""
|
|
required_sections = ['cameras', 'arm']
|
|
for section in required_sections:
|
|
if section not in self.config:
|
|
raise ValueError(f"Missing required config section: {section}")
|
|
|
|
def _init_data_structures(self) -> None:
|
|
"""初始化数据结构模板方法"""
|
|
# 相机数据
|
|
self.cameras = self.config.get('cameras', {})
|
|
self.sync_img_queues = {name: deque(maxlen=2000) for name in self.cameras}
|
|
|
|
# 深度数据
|
|
self.use_depth_image = self.config.get('use_depth_image', False)
|
|
if self.use_depth_image:
|
|
self.sync_depth_queues = {
|
|
name: deque(maxlen=2000)
|
|
for name, cam in self.cameras.items()
|
|
if 'depth_topic_name' in cam
|
|
}
|
|
|
|
# 机械臂数据
|
|
self.arms = self.config.get('arm', {})
|
|
if self.config.get('control_type', '') != 'record':
|
|
# 如果不是录制模式，则仅初始化从机械臂数据队列
|
|
self.sync_arm_queues = {name: deque(maxlen=2000) for name in self.arms if 'puppet' in name}
|
|
else:
|
|
self.sync_arm_queues = {name: deque(maxlen=2000) for name in self.arms}
|
|
|
|
# 机器人基座数据
|
|
self.use_robot_base = self.config.get('use_robot_base', False)
|
|
if self.use_robot_base:
|
|
self.sync_base_queue = deque(maxlen=2000)
|
|
|
|
def init_ros(self) -> None:
|
|
"""初始化ROS订阅的模板方法"""
|
|
rospy.init_node(
|
|
f"{self.config.get('ros_node_name', 'generic_robot_node')}",
|
|
anonymous=True
|
|
)
|
|
|
|
self._setup_camera_subscribers()
|
|
self._setup_arm_subscribers_publishers()
|
|
self._setup_base_subscriber()
|
|
self._log_ros_status()
|
|
|
|
def init_features(self):
|
|
"""
|
|
根据YAML配置自动生成features结构
|
|
"""
|
|
self.features = {}
|
|
|
|
# 初始化相机特征
|
|
self._init_camera_features()
|
|
|
|
# 初始化机械臂特征
|
|
self._init_state_features()
|
|
|
|
self._init_action_features()
|
|
|
|
# 初始化基座特征（如果启用）
|
|
if self.use_robot_base:
|
|
self._init_base_features()
|
|
import pprint
|
|
pprint.pprint(self.features, indent=4)
|
|
|
|
|
|
def _init_camera_features(self):
|
|
"""处理所有相机特征"""
|
|
for cam_name, cam_config in self.cameras.items():
|
|
# 普通图像
|
|
self.features[f"observation.images.{cam_name}"] = {
|
|
"dtype": "video" if self.config.get("video", False) else "image",
|
|
"shape": cam_config.get("rgb_shape", [480, 640, 3]),
|
|
"names": ["height", "width", "channel"],
|
|
# "video_info": {
|
|
# "video.fps": cam_config.get("fps", 30.0),
|
|
# "video.codec": cam_config.get("codec", "av1"),
|
|
# "video.pix_fmt": cam_config.get("pix_fmt", "yuv420p"),
|
|
# "video.is_depth_map": False,
|
|
# "has_audio": False
|
|
# }
|
|
}
|
|
|
|
if self.config.get("use_depth_image", False):
|
|
self.features[f"observation.images.depth_{cam_name}"] = {
|
|
"dtype": "uint16",
|
|
"shape": (cam_config.get("width", 480), cam_config.get("height", 640), 1),
|
|
"names": ["height", "width"],
|
|
}
|
|
|
|
|
|
def _init_state_features(self):
|
|
state = self.config.get('state', {})
|
|
# 状态特征
|
|
self.features["observation.state"] = {
|
|
"dtype": "float32",
|
|
"shape": (len(state.get('motors', "")),),
|
|
"names": {"motors": state.get('motors', "")}
|
|
}
|
|
|
|
if self.config.get('velocity'):
|
|
velocity = self.config.get('velocity', "")
|
|
self.features["observation.velocity"] = {
|
|
"dtype": "float32",
|
|
"shape": (len(velocity.get('motors', "")),),
|
|
"names": {"motors": velocity.get('motors', "")}
|
|
}
|
|
|
|
if self.config.get('effort'):
|
|
effort = self.config.get('effort', "")
|
|
self.features["observation.effort"] = {
|
|
"dtype": "float32",
|
|
"shape": (len(effort.get('motors', "")),),
|
|
"names": {"motors": effort.get('motors', "")}
|
|
}
|
|
|
|
|
|
|
|
def _init_action_features(self):
|
|
action = self.config.get('action', {})
|
|
# 状态特征
|
|
self.features["action"] = {
|
|
"dtype": "float32",
|
|
"shape": (len(action.get('motors', "")),),
|
|
"names": {"motors": action.get('motors', "")}
|
|
}
|
|
|
|
def _init_base_features(self):
|
|
"""处理基座特征"""
|
|
self.features["observation.base_vel"] = {
|
|
"dtype": "float32",
|
|
"shape": (2,),
|
|
"names": ["linear_x", "angular_z"]
|
|
}
|
|
|
|
|
|
def _setup_camera_subscribers(self) -> None:
|
|
"""设置相机订阅者"""
|
|
for cam_name, cam_config in self.cameras.items():
|
|
if 'img_topic_name' in cam_config:
|
|
self.subscribers[f"camera_{cam_name}"] = rospy.Subscriber(
|
|
cam_config['img_topic_name'],
|
|
Image,
|
|
self._make_camera_callback(cam_name, is_depth=False),
|
|
queue_size=1000,
|
|
tcp_nodelay=True
|
|
)
|
|
|
|
if self.use_depth_image and 'depth_topic_name' in cam_config:
|
|
self.subscribers[f"depth_{cam_name}"] = rospy.Subscriber(
|
|
cam_config['depth_topic_name'],
|
|
Image,
|
|
self._make_camera_callback(cam_name, is_depth=True),
|
|
queue_size=1000,
|
|
tcp_nodelay=True
|
|
)
|
|
|
|
def _setup_arm_subscribers_publishers(self) -> None:
|
|
"""设置机械臂订阅者"""
|
|
# 当为record模式时,主从机械臂都需要订阅
|
|
# 否则只订阅从机械臂,但向主机械臂发布
|
|
if self.config.get('control_type', '') == 'record':
|
|
for arm_name, arm_config in self.arms.items():
|
|
if 'topic_name' in arm_config:
|
|
self.subscribers[f"arm_{arm_name}"] = rospy.Subscriber(
|
|
arm_config['topic_name'],
|
|
JointState,
|
|
self._make_arm_callback(arm_name),
|
|
queue_size=1000,
|
|
tcp_nodelay=True
|
|
)
|
|
else:
|
|
for arm_name, arm_config in self.arms.items():
|
|
if 'puppet' in arm_name:
|
|
self.subscribers[f"arm_{arm_name}"] = rospy.Subscriber(
|
|
arm_config['topic_name'],
|
|
JointState,
|
|
self._make_arm_callback(arm_name),
|
|
queue_size=1000,
|
|
tcp_nodelay=True
|
|
)
|
|
if 'master' in arm_name:
|
|
self.publishers[f"arm_{arm_name}"] = rospy.Publisher(
|
|
arm_config['topic_name'],
|
|
JointState,
|
|
queue_size=10
|
|
)
|
|
|
|
def _setup_base_subscriber(self) -> None:
|
|
"""设置基座订阅者"""
|
|
if self.use_robot_base and 'robot_base' in self.config:
|
|
self.subscribers['base'] = rospy.Subscriber(
|
|
self.config['robot_base']['topic_name'],
|
|
Odometry,
|
|
self.robot_base_callback,
|
|
queue_size=1000,
|
|
tcp_nodelay=True
|
|
)
|
|
|
|
def _log_ros_status(self) -> None:
|
|
"""记录ROS状态"""
|
|
rospy.loginfo("\n=== ROS订阅状态 ===")
|
|
rospy.loginfo(f"已初始化节点: {rospy.get_name()}")
|
|
rospy.loginfo("活跃的订阅者:")
|
|
for topic, sub in self.subscribers.items():
|
|
rospy.loginfo(f" - {topic}: {'活跃' if sub.impl else '未连接'}")
|
|
rospy.loginfo("=================")
|
|
|
|
def _make_camera_callback(self, cam_name: str, is_depth: bool = False):
|
|
"""生成相机回调函数工厂方法"""
|
|
def callback(msg):
|
|
try:
|
|
target_queue = (
|
|
self.sync_depth_queues[cam_name]
|
|
if is_depth
|
|
else self.sync_img_queues[cam_name]
|
|
)
|
|
if len(target_queue) >= 2000:
|
|
target_queue.popleft()
|
|
target_queue.append(msg)
|
|
except Exception as e:
|
|
rospy.logerr(f"Camera {cam_name} callback error: {str(e)}")
|
|
return callback
|
|
|
|
def _make_arm_callback(self, arm_name: str):
|
|
"""生成机械臂回调函数工厂方法"""
|
|
def callback(msg):
|
|
try:
|
|
if len(self.sync_arm_queues[arm_name]) >= 2000:
|
|
self.sync_arm_queues[arm_name].popleft()
|
|
self.sync_arm_queues[arm_name].append(msg)
|
|
except Exception as e:
|
|
rospy.logerr(f"Arm {arm_name} callback error: {str(e)}")
|
|
return callback
|
|
|
|
def robot_base_callback(self, msg):
|
|
"""基座回调默认实现"""
|
|
if len(self.sync_base_queue) >= 2000:
|
|
self.sync_base_queue.popleft()
|
|
self.sync_base_queue.append(msg)
|
|
|
|
def warmup(self, timeout: float = 10.0) -> bool:
|
|
"""Wait until all data queues have at least 20 messages.
|
|
|
|
Args:
|
|
timeout: Maximum time to wait in seconds before giving up
|
|
|
|
Returns:
|
|
bool: True if warmup succeeded, False if timed out
|
|
"""
|
|
start_time = rospy.Time.now().to_sec()
|
|
rate = rospy.Rate(10) # Check at 10Hz
|
|
|
|
rospy.loginfo("Starting warmup process...")
|
|
|
|
while not rospy.is_shutdown():
|
|
# Check if timeout has been reached
|
|
current_time = rospy.Time.now().to_sec()
|
|
if current_time - start_time > timeout:
|
|
rospy.logwarn("Warmup timed out before all queues were filled")
|
|
return False
|
|
|
|
# Check all required queues
|
|
all_ready = True
|
|
|
|
# Check camera image queues
|
|
for cam_name in self.cameras:
|
|
if len(self.sync_img_queues[cam_name]) < 50:
|
|
rospy.loginfo(f"Waiting for camera {cam_name} (current: {len(self.sync_img_queues[cam_name])}/50)")
|
|
all_ready = False
|
|
break
|
|
|
|
# Check depth queues if enabled
|
|
if self.use_depth_image:
|
|
for cam_name in self.sync_depth_queues:
|
|
if len(self.sync_depth_queues[cam_name]) < 50:
|
|
rospy.loginfo(f"Waiting for depth camera {cam_name} (current: {len(self.sync_depth_queues[cam_name])}/50)")
|
|
all_ready = False
|
|
break
|
|
|
|
# # Check arm queues
|
|
# for arm_name in self.arms:
|
|
# if len(self.sync_arm_queues[arm_name]) < 20:
|
|
# rospy.loginfo(f"Waiting for arm {arm_name} (current: {len(self.sync_arm_queues[arm_name])}/20)")
|
|
# all_ready = False
|
|
# break
|
|
|
|
# Check base queue if enabled
|
|
if self.use_robot_base:
|
|
if len(self.sync_base_queue) < 20:
|
|
rospy.loginfo(f"Waiting for base (current: {len(self.sync_base_queue)}/20)")
|
|
all_ready = False
|
|
|
|
# If all queues are ready, return success
|
|
if all_ready:
|
|
rospy.loginfo("Warmup completed successfully")
|
|
return True
|
|
|
|
rate.sleep()
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
|
|
def get_frame(self) -> Optional[Dict[str, Any]]:
|
|
"""获取同步帧数据的模板方法"""
|
|
raise NotImplementedError("Subclasses must implement get_frame()")
|
|
|
|
def process(self) -> tuple:
|
|
"""主处理循环的模板方法"""
|
|
raise NotImplementedError("Subclasses must implement process()")
|