diff --git a/src/Neuro_Mujoco/main.py b/src/Neuro_Mujoco/main.py index ad13854ffc..fee7c7d74f 100644 --- a/src/Neuro_Mujoco/main.py +++ b/src/Neuro_Mujoco/main.py @@ -1,3 +1,4 @@ +#! /usr/bin/env python import os import sys import time @@ -20,7 +21,6 @@ ROS_AVAILABLE = False logging.warning("未检测到 ROS 环境，ROS 功能已禁用（如需启用，请安装 ROS 1 Noetic 并配置环境）") - # 配置日志系统 logging.basicConfig( level=logging.INFO, @@ -62,11 +62,9 @@ def load_model(model_path: str) -> Tuple[Optional[mujoco.MjModel], Optional[mujo def convert_model(input_path: str, output_path: str) -> bool: """ 转换模型格式（XML↔MJB） - 参数: input_path: 输入模型路径 output_path: 输出模型路径（需指定扩展名.xml或.mjb） - 返回: 转换成功返回True，失败返回False """ @@ -166,176 +164,162 @@ def simulate_thread(thread_id: int) -> float: logger.info(f"实时因子: {realtime_factor:.2f}x") logger.info(f"线程平均耗时: {np.mean(thread_durations):.2f}秒 (±{np.std(thread_durations):.2f})") - -# ===================== ROS管理器类（新增模块化ROS功能）===================== -class ROSManager: - """ROS管理器类（模块化ROS功能）""" +# ===================== 可视化函数（仅优化ROS关节状态发布逻辑）===================== +def visualize(model_path: str, use_ros: bool = False) -> None: + """ + 可视化模型并运行模拟（支持ROS 1模式） - def __init__(self, model: mujoco.MjModel): + 参数: + model_path: 模型文件路径 + use_ros: 是否启用ROS模式（默认False） + """ + model, data = load_model(model_path) + if not model: + return + + # ===================== ROS 1 初始化（新增）===================== + ros_publishers = None + ros_subscribers = None + ros_rate = None + ctrl_cmd = None + joint_msg = None + # 新增：存储非自由关节的ID和对应的qpos/qvel索引（解决索引错位+支持多自由度） + joint_ids = [] + joint_qpos_idxs = [] + joint_qvel_idxs = [] + + if use_ros: if not ROS_AVAILABLE: - raise RuntimeError("ROS环境未就绪，无法启用ROS模式") - - self.model = model - self.ctrl_cmd = None - self.joint_msg = None - self.njnt = 0 - self.initialized = False - - # ROS发布者 - self.joint_state_pub = None - self.pose_pub = None - self.ros_rate = None + logger.error("ROS 环境未就绪，无法启用 ROS 模式（请检查ROS安装和环境配置）") + return - def initialize(self): - """初始化ROS节点和发布者/订阅者""" + # 初始化ROS节点 rospy.init_node("mujoco_ros_node", anonymous=True) - self.ros_rate = rospy.Rate(100) # 100Hz发布频率 - - # 创建发布者 - self.joint_state_pub = rospy.Publisher( + ros_rate = rospy.Rate(100) # 100Hz发布频率（与MuJoCo默认步长0.01s匹配） + logger.info("="*60) + logger.info("ROS 1 模式已启用！") + logger.info(f"发布话题：/mujoco/joint_states（关节状态）、/mujoco/pose（基座姿态）") + logger.info(f"订阅话题：/mujoco/ctrl_cmd（控制指令，长度={model.nu}）") + logger.info("="*60) + + # 1. 创建ROS发布者 + joint_state_pub = rospy.Publisher( "/mujoco/joint_states", JointState, - queue_size=10 + queue_size=10 # 消息队列大小 ) - - self.pose_pub = rospy.Publisher( + pose_pub = rospy.Publisher( "/mujoco/pose", PoseStamped, queue_size=10 ) - - # 初始化关节状态消息 - self.joint_msg = JointState() - joint_names = [] - - for i in range(self.model.njnt): - joint_type = self.model.joint(i).type + ros_publishers = (joint_state_pub, pose_pub) + + # 2. 初始化关节状态消息（精准映射非自由关节的索引，支持多自由度） + joint_msg = JointState() + joint_msg.name = [] + for i in range(model.njnt): + joint_type = model.joint(i).type if joint_type != mujoco.mjtJoint.mjJNT_FREE: - joint_names.append(self.model.joint(i).name) + joint_msg.name.append(model.joint(i).name) + joint_ids.append(i) + # 获取该关节在qpos中的起始索引（mjJNT_FREE=7维, mjJNT_BALL=3维, mjJNT_HINGE/SLIDE=1维） + joint_qpos_idxs.append(model.jnt_qposadr[i]) + # 获取该关节在qvel中的起始索引 + joint_qvel_idxs.append(model.jnt_dofadr[i]) - self.joint_msg.name = joint_names - self.njnt = len(joint_names) + njnt = len(joint_msg.name) + logger.info(f"ROS将发布 {njnt} 个非自由关节状态：{joint_msg.name}") + if njnt > 0: + logger.debug(f"关节qpos索引映射：{dict(zip(joint_msg.name, joint_qpos_idxs))}") + + # 3. 创建ROS订阅者（接收控制指令） + ctrl_cmd = np.zeros(model.nu) if model.nu > 0 else None + def ctrl_callback(msg: Float32MultiArray): + nonlocal ctrl_cmd + if model.nu == len(msg.data): + ctrl_cmd = np.array(msg.data) + logger.debug(f"收到ROS控制指令：{ctrl_cmd[:5]}...") # 只打印前5个值，避免日志冗余 + else: + logger.warning(f"控制指令长度不匹配！期望 {model.nu} 个值，实际收到 {len(msg.data)} 个") - # 初始化控制命令 - if self.model.nu > 0: - self.ctrl_cmd = np.zeros(self.model.nu) - rospy.Subscriber( + if model.nu > 0: + ros_subscribers = rospy.Subscriber( "/mujoco/ctrl_cmd", Float32MultiArray, - self._ctrl_callback, + ctrl_callback, queue_size=5 ) - - self.initialized = True - return self - - def _ctrl_callback(self, msg: Float32MultiArray): - """控制指令回调函数""" - if self.model.nu == len(msg.data): - self.ctrl_cmd = np.array(msg.data) - logger.debug(f"收到ROS控制指令，前5个值: {self.ctrl_cmd[:5]}...") - - def apply_control(self, data: mujoco.MjData): - """将ROS控制指令应用到MuJoCo数据""" - if self.ctrl_cmd is not None and self.model.nu > 0: - data.ctrl[:] = self.ctrl_cmd - - def publish_states(self, data: mujoco.MjData): - """发布关节状态和姿态""" - if not self.initialized: - return - - # 发布关节状态 - self.joint_msg.header.stamp = rospy.Time.now() - self.joint_msg.position = data.qpos[:self.njnt].tolist() - self.joint_msg.velocity = data.qvel[:self.njnt].tolist() - self.joint_state_pub.publish(self.joint_msg) - - # 发布基座姿态 - pose_msg = PoseStamped() - pose_msg.header.stamp = rospy.Time.now() - pose_msg.header.frame_id = "world" - - if self.model.nq >= 1: - pose_msg.pose.position.x = data.qpos[0] - if self.model.nq >= 2: - pose_msg.pose.position.y = data.qpos[1] - if self.model.nq >= 3: - pose_msg.pose.position.z = data.qpos[2] - - if self.model.nq >= 4: - pose_msg.pose.orientation.x = data.qpos[3] - if self.model.nq >= 5: - pose_msg.pose.orientation.y = data.qpos[4] - if self.model.nq >= 6: - pose_msg.pose.orientation.z = data.qpos[5] - if self.model.nq >= 7: - pose_msg.pose.orientation.w = data.qpos[6] - - self.pose_pub.publish(pose_msg) - - # 控制发布频率 - self.ros_rate.sleep() - - -# ===================== 可视化函数（使用ROSManager类）===================== -def visualize(model_path: str, use_ros: bool = False) -> None: - """ - 可视化模型并运行模拟（支持ROS 1模式） - - 参数: - model_path: 模型文件路径 - use_ros: 是否启用ROS模式（默认False） - """ - model, data = load_model(model_path) - if not model: - return - - # ROS管理器初始化 - ros_manager = None - if use_ros: - if not ROS_AVAILABLE: - logger.error("ROS 环境未就绪，无法启用 ROS 模式（请检查ROS安装和环境配置）") - return - - try: - ros_manager = ROSManager(model).initialize() - logger.info("="*60) - logger.info("ROS 1 模式已启用！") - logger.info(f"发布话题：/mujoco/joint_states（{ros_manager.njnt}个非自由关节）") - logger.info(f"发布话题：/mujoco/pose（基座姿态）") - logger.info(f"订阅话题：/mujoco/ctrl_cmd（控制指令，长度={model.nu}）") - logger.info("="*60) - except Exception as e: - logger.error(f"ROS初始化失败: {str(e)}") - return + else: + logger.warning("模型无控制输入（nu=0），不订阅控制指令话题") - # ===================== 可视化主循环 ===================== + # ===================== 可视化主循环（原有逻辑+优化ROS发布）===================== logger.info("启动可视化窗口（按ESC键退出，鼠标可交互：拖拽旋转、滚轮缩放）") try: with viewer.launch_passive(model, data) as v: while v.is_running() and (not use_ros or not rospy.is_shutdown()): - # ROS模式：应用控制指令 - if ros_manager: - ros_manager.apply_control(data) - - # 执行MuJoCo模拟步 + # ROS模式：应用控制指令（新增） + if use_ros and ctrl_cmd is not None: + data.ctrl[:] = ctrl_cmd + + # 执行MuJoCo模拟步（原有逻辑） mujoco.mj_step(model, data) v.sync() - - # ROS模式：发布状态 - if ros_manager: - ros_manager.publish_states(data) - + + # ===================== ROS 消息发布（仅优化关节状态部分）===================== + if use_ros and ros_publishers is not None: + joint_state_pub, pose_pub = ros_publishers + + # 1. 发布关节状态（位置、速度）- 优化后：精准映射+支持多自由度 + joint_msg.header.stamp = rospy.Time.now() + joint_msg.position = [] + joint_msg.velocity = [] + for idx, (joint_id, qpos_idx, qvel_idx) in enumerate(zip(joint_ids, joint_qpos_idxs, joint_qvel_idxs)): + joint_type = model.joint(joint_id).type + # 球关节（3自由度）：补充3维位置/速度 + if joint_type == mujoco.mjtJoint.mjJNT_BALL: + joint_msg.position.extend(data.qpos[qpos_idx:qpos_idx+3]) + joint_msg.velocity.extend(data.qvel[qvel_idx:qvel_idx+3]) + # 铰链/滑动关节（1自由度）：仅取1维 + elif joint_type in [mujoco.mjtJoint.mjJNT_HINGE, mujoco.mjtJoint.mjJNT_SLIDE]: + joint_msg.position.append(data.qpos[qpos_idx]) + joint_msg.velocity.append(data.qvel[qvel_idx]) + + joint_state_pub.publish(joint_msg) + + # 2. 发布基座姿态（原有逻辑不变） + pose_msg = PoseStamped() + pose_msg.header.stamp = rospy.Time.now() + pose_msg.header.frame_id = "world" # 坐标系名称（可自定义） + + # 位置信息（x,y,z） + if model.nq >= 1: + pose_msg.pose.position.x = data.qpos[0] + if model.nq >= 2: + pose_msg.pose.position.y = data.qpos[1] + if model.nq >= 3: + pose_msg.pose.position.z = data.qpos[2] + + # 姿态信息（四元数 qx,qy,qz,qw） + if model.nq >= 4: + pose_msg.pose.orientation.x = data.qpos[3] + if model.nq >= 5: + pose_msg.pose.orientation.y = data.qpos[4] + if model.nq >= 6: + pose_msg.pose.orientation.z = data.qpos[5] + if model.nq >= 7: + pose_msg.pose.orientation.w = data.qpos[6] + + pose_pub.publish(pose_msg) + + # 按ROS频率休眠，确保消息发布稳定 + ros_rate.sleep() + logger.info("可视化窗口已关闭") except Exception as e: logger.error(f"可视化过程出错: {str(e)}", exc_info=True) - finally: - if ros_manager: - logger.info("ROS管理器已关闭") - -# ===================== 主函数（新增--ros选项）===================== +# ===================== 主函数（完全保持原有逻辑不变）===================== def main() -> None: parser = argparse.ArgumentParser( description="MuJoCo功能整合工具（支持ROS 1消息封装）", @@ -345,7 +329,7 @@ def main() -> None: # 1. 可视化命令（新增--ros选项） viz_parser = subparsers.add_parser("visualize", help="可视化模型并运行模拟") - viz_parser.add_argument("model", help="模型文件路径") + viz_parser.add_argument("model", help="/home/lan/桌面/nn/mujoco_menagerie/anybotics_anymal_b") viz_parser.add_argument( "--ros", action="store_true", @@ -364,7 +348,8 @@ def main() -> None: convert_parser.add_argument("input", help="输入模型路径") convert_parser.add_argument("output", help="输出模型路径（需指定.xml或.mjb扩展名）") - args = parser.parse_args() + args, unknown = parser.parse_known_args() + # 命令映射（更新visualize，支持use_ros参数） command_handlers: Dict[str, callable] = { @@ -384,5 +369,6 @@ def main() -> None: sys.exit(1) + if __name__ == "__main__": main() \ No newline at end of file