diff --git a/src/box/README.md b/src/box/README.md index 81ecd8e5a8..342f85b47e 100644 --- a/src/box/README.md +++ b/src/box/README.md @@ -1,3 +1,4 @@ +<<<<<<< HEAD # box — 仿真与强化学习实验箱 ## 概述 @@ -82,3 +83,64 @@ python tests/test_simulator.py - 查看目录下的具体脚本与模块顶部注释,通常包含使用示例与参数说明; - 若需要,我可以为 `src/box` 中的主要文件生成更详细的文档或示例运行脚本。 +======= +**box — 仿真与强化学习实验箱** + +简介 +- `src/box` 目录包含基于 Gymnasium 和 MuJoCo 的仿真环境与相关辅助脚本,用于开发和测试生物力学/机器人仿真、感知模块与强化学习任务。 + +目录结构(示例) +- `simulator.py`:仿真环境核心(通常继承 `gym.Env`)。 +- `test_simulator.py`:示例运行脚本,用于启动仿真并可视化。 +- `main.py`:辅助脚本(例如证书或配置检查)。 +- `README.md`:本文件,说明目录用途与快速上手指南。 + +快速上手 +1. 创建并激活虚拟环境(以 Windows 为例): + +```powershell +cd <项目根目录> +python -m venv venv --python=3.9 +.\\venv\\Scripts\\Activate.ps1 +``` + +2. 安装依赖(建议使用清华镜像加速): + +```powershell +pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple +``` + +如果仓库没有完整的 `requirements.txt`,可参考下列核心库: + +```text +gymnasium +mujoco +stable-baselines3 +pygame +opencv-python +numpy +scipy +matplotlib +ruamel.yaml +certifi +``` + +运行示例 +- 启动仿真: + +```powershell +python test_simulator.py +``` + +运行后应弹出可视化窗口(若使用 Pygame/SDL),并在终端输出仿真日志。 + +贡献与问题反馈 +- 若需添加说明或示例,请提交 Pull Request。 +- 遇到环境或依赖问题,请在 Issue 中描述操作系统、Python 版本与错误日志。 + +更多信息 +- 若目录中包含更详细的子模块文档,请参阅相应文件(如 `simulator.py` 顶部注释或同目录下的文档)。 + +--- +(此 README 为目录概览,具体实现与文件名以代码库为准) +>>>>>>> ffff1b2d (更新README.md) diff --git a/src/box/main.py b/src/box/main.py index 96f58810d7..d29e8f472d 100644 --- a/src/box/main.py +++ b/src/box/main.py @@ -1,15 +1,477 @@ -import argparse -from certifi import where - -parser = argparse.ArgumentParser() -parser.add_argument("-c", "--contents", action="store_true", help="查看证书文件内容") -args = parser.parse_args() - -if args.contents: - # 读取 certifi 证书文件内容 - cert_path = where() - with open(cert_path, "r", encoding="utf-8") as f: - print(f.read()) -else: - # 打印证书文件路径 - print(where()) \ No newline at end of file +#!/usr/bin/env python3 +""" +双机械臂协同操作仿真系统主程序 +""" + +import numpy as np +import matplotlib.pyplot as plt +import yaml +import os +import sys +from pathlib import Path +import time +from datetime import datetime + +# 添加当前目录到路径 +sys.path.append(os.path.dirname(os.path.abspath(__file__))) + +# 导入自定义模块 +from dual_arm_bm_model import DualArmBMModel +from cooperative_task import CooperativeTransportTask +from perception_module import DualEndEffectorPerception +from visualization import DualArmVisualizer + +class DualArmSimulator: + """双机械臂仿真器""" + + def __init__(self, config_path: str = "config.yaml"): + """初始化仿真器""" + # 加载配置 + with open(config_path, 'r') as f: + self.config = yaml.safe_load(f) + + print("=" * 60) + print("双机械臂协同操作仿真系统") + print("=" * 60) + + # 初始化各模块 + self.bm_model = DualArmBMModel(self.config['simulation']['bm_model']['kwargs']) + self.task = CooperativeTransportTask(self.config['simulation']['task']['kwargs']) + self.perception = DualEndEffectorPerception(self.config['simulation']['perception_modules'][0]['kwargs']) + + # 仿真参数 + self.dt = self.config['simulation']['run_parameters']['dt'] + self.max_steps = self.config['simulation']['task']['kwargs']['max_steps'] + + # 数据记录 + self.states = [] + self.actions = [] + self.rewards = [] + + # 结果目录 + self.results_dir = Path("results") + self.results_dir.mkdir(exist_ok=True) + (self.results_dir / "frames").mkdir(exist_ok=True) + (self.results_dir / "videos").mkdir(exist_ok=True) + + print("✓ 系统初始化完成") + print(f"✓ 时间步长: {self.dt}秒") + print(f"✓ 最大步数: {self.max_steps}") + print(f"✓ 结果目录: {self.results_dir.absolute()}") + + def reset(self): + """重置仿真""" + self.bm_model.reset() + self.task.reset() + self.perception.reset() + + self.states = [] + self.actions = [] + self.rewards = [] + + # 记录初始状态 + initial_state = { + 'left_arm': self.bm_model.left_arm, + 'right_arm': self.bm_model.right_arm, + 'object': self.task.state.object_position + } + self.states.append(initial_state) + + print("✓ 仿真已重置") + + def step(self, left_action: np.ndarray, right_action: np.ndarray): + """执行一步仿真""" + # 限制动作范围 + left_action = np.clip(left_action, -1.0, 1.0) + right_action = np.clip(right_action, -1.0, 1.0) + + # 更新生物力学模型 + self.bm_model.update(left_action, right_action, self.dt) + + # 获取当前末端位置 + left_pos = self.bm_model.left_arm.end_effector_pos + right_pos = self.bm_model.right_arm.end_effector_pos + + # 更新任务状态 + reward, terminated, info = self.task.update( + left_pos, right_pos, left_action, right_action, self.dt + ) + + # 获取感知观测 + observation = self.perception.get_observation( + left_pos, right_pos, self.task.state.object_position + ) + + # 记录数据 + self.actions.append((left_action.copy(), right_action.copy())) + self.rewards.append(reward) + + current_state = { + 'left_arm': self.bm_model.left_arm, + 'right_arm': self.bm_model.right_arm, + 'object': self.task.state.object_position.copy(), + 'observation': observation, + 'info': info + } + self.states.append(current_state) + + return observation, reward, terminated, info + + def run_simulation(self, policy_type: str = "sinusoidal"): + """运行仿真""" + print("\n" + "=" * 60) + print("开始双机械臂协同操作仿真") + print("=" * 60) + + self.reset() + + # 定义控制策略 + def sinusoidal_policy(step: int) -> Tuple[np.ndarray, np.ndarray]: + """正弦波控制策略(用于演示)""" + t = step * self.dt + freq = 1.0 # 频率 + + # 左臂动作 + left_action = np.array([ + 0.5 * np.sin(2 * np.pi * freq * t), # 关节1 + 0.3 * np.sin(2 * np.pi * freq * t + np.pi/3), # 关节2 + 0.2 * np.sin(2 * np.pi * freq * t + 2*np.pi/3) # 关节3 + ]) + + # 右臂动作(相位相反) + right_action = np.array([ + 0.5 * np.sin(2 * np.pi * freq * t + np.pi), # 关节1 + 0.3 * np.sin(2 * np.pi * freq * t + np.pi + np.pi/3), # 关节2 + 0.2 * np.sin(2 * np.pi * freq * t + np.pi + 2*np.pi/3) # 关节3 + ]) + + return left_action, right_action + + def target_tracking_policy(step: int) -> Tuple[np.ndarray, np.ndarray]: + """目标跟踪策略""" + t = step * self.dt + + # 计算目标位置(随时间移动) + target_left = self.task.target_left + np.array([ + 0.1 * np.sin(2 * np.pi * 0.2 * t), + 0.0, + 0.05 * np.sin(2 * np.pi * 0.3 * t) + ]) + + target_right = self.task.target_right + np.array([ + -0.1 * np.sin(2 * np.pi * 0.2 * t), + 0.0, + 0.05 * np.sin(2 * np.pi * 0.3 * t) + ]) + + # 计算当前末端位置误差 + current_left = self.bm_model.left_arm.end_effector_pos + current_right = self.bm_model.right_arm.end_effector_pos + + # PD控制器 + left_error = target_left - current_left + right_error = target_right - current_right + + # 简单比例控制 + kp = 2.0 + left_action = kp * left_error[:3] # 只取前三个分量(位置) + right_action = kp * right_error[:3] + + return left_action, right_action + + # 选择策略 + if policy_type == "sinusoidal": + policy = sinusoidal_policy + elif policy_type == "tracking": + policy = target_tracking_policy + else: + raise ValueError(f"未知策略类型: {policy_type}") + + # 运行仿真循环 + terminated = False + step = 0 + total_reward = 0.0 + + print(f"\n使用策略: {policy_type}") + print(f"{'Step':>6} {'Left Pos':>20} {'Right Pos':>20} {'Reward':>10} {'Terminated':>10}") + print("-" * 80) + + while not terminated and step < self.max_steps: + # 生成动作 + left_action, right_action = policy(step) + + # 执行一步 + observation, reward, terminated, info = self.step(left_action, right_action) + + total_reward += reward + + # 每100步打印一次状态 + if step % 100 == 0: + left_pos = self.bm_model.left_arm.end_effector_pos + right_pos = self.bm_model.right_arm.end_effector_pos + print(f"{step:6d} {str(left_pos.round(2)):>20} {str(right_pos.round(2)):>20} " + f"{reward:10.3f} {str(terminated):>10}") + + step += 1 + + print("-" * 80) + print(f"仿真完成!") + print(f"总步数: {step}") + print(f"总奖励: {total_reward:.3f}") + print(f"是否抓取成功: {self.task.state.is_grasped}") + print(f"物体最终位置: {self.task.state.object_position.round(3)}") + + return step, total_reward + + def analyze_results(self): + """分析仿真结果""" + print("\n" + "=" * 60) + print("仿真结果分析") + print("=" * 60) + + # 提取数据 + left_positions = np.array([s['left_arm'].end_effector_pos for s in self.states]) + right_positions = np.array([s['right_arm'].end_effector_pos for s in self.states]) + object_positions = np.array([s['object'] for s in self.states]) + rewards = np.array(self.rewards) + + # 计算统计量 + left_path_length = np.sum(np.linalg.norm(np.diff(left_positions, axis=0), axis=1)) + right_path_length = np.sum(np.linalg.norm(np.diff(right_positions, axis=0), axis=1)) + + left_max_speed = np.max(np.linalg.norm(np.diff(left_positions, axis=0) / self.dt, axis=1)) + right_max_speed = np.max(np.linalg.norm(np.diff(right_positions, axis=0) / self.dt, axis=1)) + + # 协同度指标 + coordination_index = self._calculate_coordination_index(left_positions, right_positions) + + print(f"左机械臂路径长度: {left_path_length:.3f} m") + print(f"右机械臂路径长度: {right_path_length:.3f} m") + print(f"左机械臂最大速度: {left_max_speed:.3f} m/s") + print(f"右机械臂最大速度: {right_max_speed:.3f} m/s") + print(f"协同度指标: {coordination_index:.3f}") + print(f"平均奖励: {np.mean(rewards):.3f}") + print(f"总奖励: {np.sum(rewards):.3f}") + + # 保存统计数据 + stats = { + 'left_path_length': float(left_path_length), + 'right_path_length': float(right_path_length), + 'left_max_speed': float(left_max_speed), + 'right_max_speed': float(right_max_speed), + 'coordination_index': float(coordination_index), + 'mean_reward': float(np.mean(rewards)), + 'total_reward': float(np.sum(rewards)), + 'total_steps': len(self.states), + 'success': bool(self.task.state.is_grasped), + 'final_object_position': self.task.state.object_position.tolist(), + 'timestamp': datetime.now().isoformat() + } + + stats_path = self.results_dir / "simulation_stats.yaml" + with open(stats_path, 'w') as f: + yaml.dump(stats, f, default_flow_style=False) + + print(f"✓ 统计数据已保存至: {stats_path}") + + return stats + + def _calculate_coordination_index(self, left_pos: np.ndarray, right_pos: np.ndarray) -> float: + """计算协同度指标""" + # 计算双手距离的稳定性 + distances = np.linalg.norm(left_pos - right_pos, axis=1) + distance_std = np.std(distances) + + # 计算运动方向的相似性 + left_vel = np.diff(left_pos, axis=0) + right_vel = np.diff(right_pos, axis=0) + + if len(left_vel) > 0: + # 计算速度方向余弦相似度 + cos_similarities = [] + for lv, rv in zip(left_vel, right_vel): + if np.linalg.norm(lv) > 0.001 and np.linalg.norm(rv) > 0.001: + cos_sim = np.dot(lv, rv) / (np.linalg.norm(lv) * np.linalg.norm(rv)) + cos_similarities.append(cos_sim) + + if cos_similarities: + mean_cos_sim = np.mean(cos_similarities) + else: + mean_cos_sim = 0 + else: + mean_cos_sim = 0 + + # 综合协同度指标(距离稳定性 + 运动相似性) + coordination = 0.5 * (1.0 / (1.0 + distance_std)) + 0.5 * (mean_cos_sim + 1) / 2 + + return coordination + + def visualize_all(self): + """生成所有可视化结果""" + print("\n" + "=" * 60) + print("生成可视化结果") + print("=" * 60) + + # 1. 轨迹图 + trajectory_plot_path = self.results_dir / "trajectory_plot.png" + self.bm_model.plot_trajectory(str(trajectory_plot_path)) + + # 2. 任务可视化 + task_viz_path = self.results_dir / "task_visualization.png" + self.task.visualize(self.bm_model.trajectory, str(task_viz_path)) + + # 3. 创建动画 + animation_path = self.results_dir / "videos" / "dual_arm_animation.mp4" + self.task.create_animation(self.bm_model.trajectory, str(animation_path)) + + # 4. 性能图表 + self._plot_performance() + + print(f"✓ 轨迹图: {trajectory_plot_path}") + print(f"✓ 任务可视化: {task_viz_path}") + print(f"✓ 动画视频: {animation_path}") + + def _plot_performance(self): + """绘制性能图表""" + fig = plt.figure(figsize=(15, 10)) + + # 奖励曲线 + ax1 = fig.add_subplot(221) + rewards = np.array(self.rewards) + cumulative_rewards = np.cumsum(rewards) + + ax1.plot(rewards, 'b-', alpha=0.7, label='Instant Reward') + ax1.plot(cumulative_rewards, 'r-', linewidth=2, label='Cumulative Reward') + ax1.axhline(y=0, color='k', linestyle='-', alpha=0.3) + ax1.set_xlabel('Time Step') + ax1.set_ylabel('Reward') + ax1.set_title('Reward Over Time') + ax1.legend() + ax1.grid(True) + + # 末端位置误差 + ax2 = fig.add_subplot(222) + left_errors = self.task.history['left_errors'] + right_errors = self.task.history['right_errors'] + + ax2.plot(left_errors, 'r-', label='Left Hand Error') + ax2.plot(right_errors, 'b-', label='Right Hand Error') + ax2.axhline(y=self.task.grasp_distance, color='g', linestyle='--', + label='Grasp Threshold') + ax2.set_xlabel('Time Step') + ax2.set_ylabel('Distance to Object (m)') + ax2.set_title('Position Errors') + ax2.legend() + ax2.grid(True) + + # 关节角度 + ax3 = fig.add_subplot(223) + left_joints = np.array([s['left_arm'].joint_positions for s in self.states]) + right_joints = np.array([s['right_arm'].joint_positions for s in self.states]) + + for i in range(3): + ax3.plot(left_joints[:, i], f'C{i}-', alpha=0.7, label=f'Left Joint {i+1}') + ax3.plot(right_joints[:, i], f'C{i}--', alpha=0.7, label=f'Right Joint {i+1}') + + ax3.set_xlabel('Time Step') + ax3.set_ylabel('Joint Angle (rad)') + ax3.set_title('Joint Angles Over Time') + ax3.legend(ncol=2) + ax3.grid(True) + + # 协同度分析 + ax4 = fig.add_subplot(224) + + # 计算双手距离 + left_pos = np.array([s['left_arm'].end_effector_pos for s in self.states]) + right_pos = np.array([s['right_arm'].end_effector_pos for s in self.states]) + hand_distances = np.linalg.norm(left_pos - right_pos, axis=1) + + ax4.plot(hand_distances, 'purple', linewidth=2) + ax4.axhline(y=0.2, color='g', linestyle='--', label='Ideal Distance (0.2m)') + ax4.set_xlabel('Time Step') + ax4.set_ylabel('Distance Between Hands (m)') + ax4.set_title('Inter-Hand Distance (Coordination Metric)') + ax4.legend() + ax4.grid(True) + + plt.suptitle('Dual-Arm Cooperative Transport Performance Analysis', fontsize=14) + plt.tight_layout() + + performance_path = self.results_dir / "performance_analysis.png" + plt.savefig(str(performance_path), dpi=150, bbox_inches='tight') + plt.show() + + print(f"✓ 性能分析图: {performance_path}") + +def main(): + """主函数""" + print("\n" + "=" * 60) + print("双机械臂协同操作仿真系统") + print("=" * 60) + + try: + # 创建仿真器 + simulator = DualArmSimulator() + + # 运行仿真(可选择不同策略) + print("\n请选择控制策略:") + print("1. 正弦波控制 (演示)") + print("2. 目标跟踪控制") + + choice = input("请输入选择 (1 或 2): ").strip() + + if choice == "1": + policy_type = "sinusoidal" + elif choice == "2": + policy_type = "tracking" + else: + print("使用默认策略: 正弦波控制") + policy_type = "sinusoidal" + + # 运行仿真 + steps, total_reward = simulator.run_simulation(policy_type) + + # 分析结果 + stats = simulator.analyze_results() + + # 生成可视化 + simulator.visualize_all() + + # 保存仿真数据 + data_path = simulator.results_dir / "simulation_data.npz" + np.savez_compressed( + str(data_path), + states=simulator.states, + actions=simulator.actions, + rewards=simulator.rewards, + config=simulator.config + ) + + print(f"\n✓ 仿真数据已保存至: {data_path}") + + print("\n" + "=" * 60) + print("仿真完成!") + print(f"总步数: {steps}") + print(f"总奖励: {total_reward:.3f}") + print(f"协同度: {stats['coordination_index']:.3f}") + print(f"任务成功: {'是' if stats['success'] else '否'}") + print("=" * 60) + + # 显示关键结果文件 + print("\n生成的结果文件:") + for file in simulator.results_dir.rglob("*"): + if file.is_file(): + size_mb = file.stat().st_size / (1024 * 1024) + print(f" - {file.relative_to(simulator.results_dir)} ({size_mb:.2f} MB)") + + except Exception as e: + print(f"\n❌ 运行出错: {e}") + import traceback + traceback.print_exc() + return 1 + + return 0 + +if __name__ == "__main__": + sys.exit(main()) \ No newline at end of file diff --git a/src/box/ros/README.md b/src/box/ros/README.md new file mode 100644 index 0000000000..5ba769c8c0 --- /dev/null +++ b/src/box/ros/README.md @@ -0,0 +1,82 @@ +### ROS 2 机械臂感知+数据获取+虚拟运动可视化 完整封装步骤(最终版) +「感知模块发布关节数据 + 数据获取模块保存数据 + RViz2 虚拟运动可视化」的核心步骤,每一步可落地、无冗余,适配 ROS 2 Humble + WSL2 环境: + +#### 一、基础环境与工作空间准备 +1. **确认依赖安装**(确保ROS 2核心依赖齐全): + ``` + sudo apt update && sudo apt install -y ros-humble-rclpy ros-humble-sensor-msgs ros-humble-robot-state-publisher ros-humble-rviz2 + ``` +2. **创建并初始化ROS 2工作空间**: + ``` + # 创建目录 + mkdir -p ~/ros2_ws/src + cd ~/ros2_ws + # 初始化编译(生成基础配置) + colcon build --symlink-install + # 加载环境变量(每次新开终端需执行,或添加到~/.bashrc) + source install/setup.bash + ``` + +#### 二、创建核心功能包(感知+数据获取) +1. **创建感知模块功能包**(发布机械臂关节数据): + ```bash + cd ~/ros2_ws/src + ros2 pkg create --build-type ament_python perception_module --dependencies rclpy std_msgs sensor_msgs + ``` +2. **创建数据获取模块功能包**(订阅并保存关节数据): + ```bash + ros2 pkg create --build-type ament_python data_acquisition_module --dependencies rclpy std_msgs sensor_msgs + ``` + +#### 三、编写核心代码(感知+数据获取) +##### 1. 感知模块代码(发布关节数据) +- 编辑文件:`~/ros2_ws/src/perception_module/perception_module/perception_node.py` + 写入可运行的关节数据发布代码(核心逻辑:定时发布6轴机械臂关节角度,模拟joint2抬升运动); +- 配置编译入口:修改`~/ros2_ws/src/perception_module/setup.py`,在`entry_points`中添加: + ```python + 'console_scripts': ['arm_perception_node = perception_module.perception_node:main'] + ``` + +##### 2. 数据获取模块代码(保存数据到CSV) +- 编辑文件:`~/ros2_ws/src/data_acquisition_module/data_acquisition_module/acquisition_node.py` + 写入订阅关节话题、保存数据到CSV的代码(核心逻辑:订阅`/arm/joint_states`,按时间戳保存关节名+角度); +- 配置编译入口:修改`~/ros2_ws/src/data_acquisition_module/setup.py`,在`entry_points`中添加: + ```python + 'console_scripts': ['arm_acquisition_node = data_acquisition_module.acquisition_node:main'] + ``` + +#### 四、编写URDF模型+可视化配置(虚拟运动) +1. **创建URDF文件**(定义6轴机械臂结构): + ```bash + mkdir -p ~/ros2_ws/src/perception_module/urdf + nano ~/ros2_ws/src/perception_module/urdf/6dof_arm.urdf + ``` + 写入6轴机械臂URDF代码(定义底座、连杆、关节、颜色,适配`joint1~joint6`); +2. **配置URDF打包**:修改`~/ros2_ws/src/perception_module/setup.py`,在`data_files`中添加URDF文件路径: + ```python + ('share/' + package_name + '/urdf', ['urdf/6dof_arm.urdf']) + ``` + +#### 五、编译工作空间(使代码/配置生效) +```bash +cd ~/ros2_ws +colcon build --symlink-install # --symlink-install:修改代码无需重新编译 +source install/setup.bash # 重新加载环境变量 +``` + +#### 六、启动节点+验证功能(分终端操作,最稳定) +| 终端序号 | 执行命令(核心操作)| 功能说明 | +|----------|---------------------------------------------|--------------------------| +| 终端1 | `ros2 run perception_module arm_perception_node` | 启动感知模块,发布关节数据 | +| 终端2 | `ros2 run data_acquisition_module arm_acquisition_node` | 启动数据获取模块,保存CSV | +| 终端3 | `ros2 run robot_state_publisher robot_state_publisher --ros-args -p robot_description:="$(cat ~/ros2_ws/src/perception_module/urdf/6dof_arm.urdf)"` | 解析URDF,发布坐标系(TF) | +| 终端4 | `rviz2` | 启动RViz2,手动配置可视化 | + +##### RViz2 手动配置(必做,确保模型显示) +1. 点击「Add」→ 选择「RobotModel」→ 「OK」; +2. RobotModel 面板配置: + - `Description Source` → `Topic`; + - `Description Topic` → `/robot_description`; + - `Joint State Topic` → `/arm/joint_states`; +3. 顶部「Global Options」→ `Fixed Frame` → `base_link`。 + diff --git a/src/box/ros/acquisition_node.py b/src/box/ros/acquisition_node.py new file mode 100644 index 0000000000..b6a421011c --- /dev/null +++ b/src/box/ros/acquisition_node.py @@ -0,0 +1,61 @@ +# 导入ROS 2核心库 +import rclpy +from rclpy.node import Node +# 导入关节消息类型 +from sensor_msgs.msg import JointState +# 导入文件操作相关库 +import csv +from datetime import datetime +import os + +class ArmDataAcquisitionNode(Node): + """数据获取模块节点:订阅关节数据并保存为CSV""" + def __init__(self): + super().__init__('arm_data_acquisition_node') + + # 创建数据保存目录(带时间戳,避免重名) + self.save_dir = os.path.expanduser(f"~/ros2_arm_data/{datetime.now().strftime('%Y%m%d_%H%M%S')}") + os.makedirs(self.save_dir, exist_ok=True) + + # 创建CSV文件并写入表头 + self.csv_file_path = os.path.join(self.save_dir, 'arm_joint_data.csv') + with open(self.csv_file_path, 'w', newline='') as f: + writer = csv.writer(f) + writer.writerow(['时间戳(秒)', '关节名', '关节角度(弧度)']) + + # 创建订阅者:订阅/arm/joint_states话题,回调函数处理数据 + self.joint_subscriber = self.create_subscription( + JointState, + '/arm/joint_states', + self.joint_data_callback, + 10 # 队列大小 + ) + + # 日志提示:节点启动成功 + self.get_logger().info(f"数据获取模块已启动!数据保存路径:{self.csv_file_path}") + + def joint_data_callback(self, msg): + """订阅回调函数:处理接收到的关节数据""" + # 计算时间戳(秒,精确到小数点后2位) + timestamp = msg.header.stamp.sec + msg.header.stamp.nanosec / 1e9 + timestamp = round(timestamp, 2) + + # 将每个关节的角度写入CSV文件 + with open(self.csv_file_path, 'a', newline='') as f: + writer = csv.writer(f) + for joint_name, joint_angle in zip(msg.name, msg.position): + writer.writerow([timestamp, joint_name, round(joint_angle, 2)]) + + # 日志输出:确认数据保存 + self.get_logger().info(f"保存数据:时间戳={timestamp},joint2角度={round(msg.position[1], 2)}") + +def main(args=None): + """主函数:启动数据获取节点""" + rclpy.init(args=args) + node = ArmDataAcquisitionNode() + rclpy.spin(node) + node.destroy_node() + rclpy.shutdown() + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/src/box/ros/perception_node.py b/src/box/ros/perception_node.py new file mode 100644 index 0000000000..b7aa46b503 --- /dev/null +++ b/src/box/ros/perception_node.py @@ -0,0 +1,63 @@ +import rclpy +from rclpy.node import Node +from sensor_msgs.msg import JointState +from geometry_msgs.msg import PointStamped +import numpy as np +# 导入仓库中的感知模块 + +class PerceptionNode(Node): + def __init__(self): + super().__init__('perception_node') + # 初始化仿真器(假设模型路径已知) + self.model_path = "~/ros2_ws/src/my-test-repo/mobl_arms_bimanual/bm_model.xml" + self.simulator = MoblArmsSimulator(model_path=self.model_path) + + # 初始化感知模块(配置末端执行器,参考仓库定义) + self.end_effector = [["site", "right_hand_site"], ["site", "left_hand_site"]] # 从bm_model.xml中获取的末端执行器名称 + self.perception = BasicWithEndEffectorPosition( + model=self.simulator.model, + data=self.simulator.data, + bm_model=self.simulator.bimanual_model, + end_effector=self.end_effector + ) + + # 创建ROS话题发布者 + self.joint_state_pub = self.create_publisher(JointState, '/joint_states', 10) + self.end_effector_pub = self.create_publisher(PointStamped, '/end_effector_position', 10) + + # 定时器(10Hz发布感知数据) + self.timer = self.create_timer(0.1, self.publish_perception_data) + self.get_logger().info("感知模块启动成功") + + def publish_perception_data(self): + # 1. 发布关节状态(角度、速度) + joint_msg = JointState() + joint_msg.header.stamp = self.get_clock().now().to_msg() + # 从仿真器获取关节名称和角度(参考仓库中_joint_id映射) + joint_names = ["shoulder_elv_r", "shoulder_elv_l", "elbow_flexion_r", "elbow_flexion_l"] + joint_msg.name = joint_names + joint_msg.position = [self.simulator.data.qpos[self.simulator._get_joint_id(name)] for name in joint_names] + joint_msg.velocity = [self.simulator.data.qvel[self.simulator._get_joint_id(name)] for name in joint_names] + self.joint_state_pub.publish(joint_msg) + + # 2. 发布末端执行器位置(从感知模块获取) + # 假设感知模块输出包含末端执行器坐标(需根据仓库代码调整) + end_pos = self.perception._get_end_effector_positions() # 需在BasicWithEndEffectorPosition中实现该方法 + for i, pos in enumerate(end_pos): + point_msg = PointStamped() + point_msg.header.stamp = self.get_clock().now().to_msg() + point_msg.header.frame_id = f"end_effector_{i}" + point_msg.point.x = pos[0] + point_msg.point.y = pos[1] + point_msg.point.z = pos[2] + self.end_effector_pub.publish(point_msg) + +def main(args=None): + rclpy.init(args=args) + node = PerceptionNode() + rclpy.spin(node) + node.destroy_node() + rclpy.shutdown() + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/src/box/BasicWithEndEffectorPosition.py b/src/box/simulators/arm_simulation/BasicWithEndEffectorPosition.py similarity index 100% rename from src/box/BasicWithEndEffectorPosition.py rename to src/box/simulators/arm_simulation/BasicWithEndEffectorPosition.py diff --git a/src/box/eye.py b/src/box/simulators/arm_simulation/eye.py similarity index 100% rename from src/box/eye.py rename to src/box/simulators/arm_simulation/eye.py diff --git a/src/box/jy.py b/src/box/simulators/arm_simulation/jy.py similarity index 100% rename from src/box/jy.py rename to src/box/simulators/arm_simulation/jy.py diff --git a/src/box/simulators/arm_simulation/llc_controller.py b/src/box/simulators/arm_simulation/llc_controller.py new file mode 100644 index 0000000000..63d88a07d4 --- /dev/null +++ b/src/box/simulators/arm_simulation/llc_controller.py @@ -0,0 +1,266 @@ +import os +import numpy as np +np.set_printoptions(precision=2, suppress=True) +from stable_baselines3 import PPO +import re +import argparse +import scipy.ndimage +from collections import defaultdict +import matplotlib.pyplot as pp +import matplotlib +matplotlib.use('TkAgg') +import cv2 +from time import sleep +import queue +from pynput.keyboard import Listener, Key + +from uitb.utils.logger import StateLogger, ActionLogger +from uitb.simulator import Simulator + + +def natural_sort(l): + convert = lambda text: int(text) if text.isdigit() else text.lower() + alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)] + return sorted(l, key=alphanum_key) + +def grab_pip_image(simulator): + # Grab an image from both 'for_testing' camera and 'oculomotor' camera, and display them 'picture-in-picture' + + # Visualise muscle actuation + simulator._model.tendon_rgba[:, 0] = 0.3 + simulator._data.ctrl * 0.7 + + # Grab images + img, _ = simulator._GUI_camera.render() + + ocular_img = None + for module in simulator.perception.perception_modules: + if module.modality == "vision": + # TODO would be better to have a class function that returns "human-viewable" rendering of the observation; + # e.g. in case the vision model has two cameras, or returns a combination of rgb + depth images etc. + ocular_img, _ = module._camera.render() + + # Set back to default + simulator._model.tendon_rgba[:, 0] = 0.95 + + if ocular_img is not None: + + # Resample + resample_factor = 2 + resample_height = ocular_img.shape[0]*resample_factor + resample_width = ocular_img.shape[1]*resample_factor + resampled_img = np.zeros((resample_height, resample_width, 3), dtype=np.uint8) + for channel in range(3): + resampled_img[:, :, channel] = scipy.ndimage.zoom(ocular_img[:, :, channel], resample_factor, order=0) + + # Embed ocular image into free image + i = simulator._GUI_camera.height - resample_height + j = simulator._GUI_camera.width - resample_width + img[i:, j:] = resampled_img + + return img + +# def read_keyboard(queue): +# # Wait for input from user +# input = + +def on_press(key): + try: + + # Exit with esc + if key == Key.esc: + return False + + # Set to initial position with space + elif key == Key.space: + q.put(("set", initial_position.copy())) + + # First joint + elif key.char == "1": + q.put(("add", np.array([0.5, 0, 0, 0, 0]))) + elif key.char == "q": + q.put(("add", np.array([0.1, 0, 0, 0, 0]))) + elif key.char == "a": + q.put(("add", np.array([-0.1, 0, 0, 0, 0]))) + elif key.char == "z": + q.put(("add", np.array([-0.5, 0, 0, 0, 0]))) + + # Second joint + elif key.char == "2": + q.put(("add", np.array([0, 0.5, 0, 0, 0]))) + elif key.char == "w": + q.put(("add", np.array([0, 0.1, 0, 0, 0]))) + elif key.char == "s": + q.put(("add", np.array([0, -0.1, 0, 0, 0]))) + elif key.char == "x": + q.put(("add", np.array([0, -0.5, 0, 0, 0]))) + + # Third joint + elif key.char == "3": + q.put(("add", np.array([0, 0, 0.5, 0, 0]))) + elif key.char == "e": + q.put(("add", np.array([0, 0, 0.1, 0, 0]))) + elif key.char == "d": + q.put(("add", np.array([0, 0, -0.1, 0, 0]))) + elif key.char == "c": + q.put(("add", np.array([0, 0, -0.5, 0, 0]))) + + # Fourth joint + elif key.char == "4": + q.put(("add", np.array([0, 0, 0, 0.5, 0]))) + elif key.char == "r": + q.put(("add", np.array([0, 0, 0, 0.1, 0]))) + elif key.char == "f": + q.put(("add", np.array([0, 0, 0, -0.1, 0]))) + elif key.char == "v": + q.put(("add", np.array([0, 0, 0, -0.5, 0]))) + + # Fifth joint + elif key.char == "5": + q.put(("add", np.array([0, 0, 0, 0, 0.5]))) + elif key.char == "t": + q.put(("add", np.array([0, 0, 0, 0, 0.1]))) + elif key.char == "g": + q.put(("add", np.array([0, 0, 0, 0, -0.1]))) + elif key.char == "b": + q.put(("add", np.array([0, 0, 0, 0, -0.5]))) + + else: + return + + except AttributeError: + pass + +if __name__=="__main__": + + parser = argparse.ArgumentParser(description='Evaluate a policy.') + parser.add_argument('simulator_folder', type=str, + help='the simulation folder') + parser.add_argument('--action_sample_freq', type=float, default=100, + help='action sample frequency (how many times per second actions are sampled from policy, default: 20)') + parser.add_argument('--checkpoint', type=str, default=None, + help='filename of a specific checkpoint (default: None, latest checkpoint is used)') + parser.add_argument('--logging', action='store_true', help='enable logging') + parser.add_argument('--state_log_file', default='controller_state_log', + help='output file for state log if logging is enabled (default: ./controller_state_log)') + parser.add_argument('--action_log_file', default='controller_action_log', + help='output file for action log if logging is enabled (default: ./controller_action_log)') + args = parser.parse_args() + + # Define directories + checkpoint_dir = os.path.join(args.simulator_folder, 'checkpoints') + evaluate_dir = os.path.join(args.simulator_folder, 'evaluate') + + # Make sure output dir exists + os.makedirs(evaluate_dir, exist_ok=True) + + # Override run parameters + run_params = dict() + run_params["action_sample_freq"] = args.action_sample_freq + run_params["evaluate"] = True + + # Use deterministic actions? + deterministic = True + + # Initialise simulator + simulator = Simulator.get(args.simulator_folder, run_parameters=run_params) + + print(f"run parameters are: {simulator.run_parameters}\n") + + # Load latest model if filename not given + if args.checkpoint is not None: + model_file = args.checkpoint + else: + files = natural_sort(os.listdir(checkpoint_dir)) + model_file = files[-1] + + # Load policy TODO should create a load method for uitb.rl.BaseRLModel + print(f'Loading model: {os.path.join(checkpoint_dir, model_file)}\n') + model = PPO.load(os.path.join(checkpoint_dir, model_file)) + + # Set callbacks to match the value used for this training point (if the simulator had any) + simulator.update_callbacks(model.num_timesteps) + + if args.logging: + + # Initialise log + state_logger = StateLogger(args.num_episodes, keys=simulator.get_state().keys()) + + # Actions are logged separately to make things easier + action_logger = ActionLogger(args.num_episodes) + + # Visualise evaluations + statistics = defaultdict(list) + imgs = [] + + # Reset environment + simulator.reset() + simulator.task._target_qpos[:] = simulator.task._qpos.copy() + initial_position = simulator.task._qpos.copy() + obs = simulator.get_observation() + + + if args.logging: + state = simulator.get_state() + state_logger.log(0, state) + + # Initialise a queue so we can share data between threads + q = queue.Queue() + + # Start a keyboard press reader in a thread + # keyboard_reader = threading.Thread(target=read_keyboard, args=(q,), daemon=True) + # keyboard_reader.start() + # listener = Listener(on_press=lambda event: on_press(event, q)) + listener = Listener(on_press=on_press) + listener.start() + + # Run simulation here in main thread + while True: + + # If listener has exited (user pressed 'esc'), break + if not listener.is_alive(): + break + + # Read queue + if not q.empty(): + action = q.get() + if action[0] == "add": + simulator.task._target_qpos = np.clip(simulator.task._target_qpos + action[1], -1, 1) + elif action[0] == "set": + simulator.task._target_qpos = action[1] + obs = simulator.get_observation() + print(simulator.task._target_qpos) + + # Get actions from policy + action, _states = model.predict(obs, deterministic=deterministic) + + # Take a step + obs, r, terminated, truncated, info = simulator.step(action) + + if args.logging: + action_logger.log(0, {"steps": state["steps"], "timestep": state["timestep"], "action": action.copy(), + "reward": r}) + state = simulator.get_state() + state.update(info) + state_logger.log(0, state) + + # Do some rendering + img = grab_pip_image(simulator) + qpos = simulator.task._target_qpos + cv2.putText(img, f'{qpos[0]:.2f} elevation angle', (700, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2, cv2.LINE_AA) + cv2.putText(img, f'{qpos[1]:.2f} shoulder elevation', (700, 100), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2, cv2.LINE_AA) + cv2.putText(img, f'{qpos[2]:.2f} shoulder rotation', (700, 150), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2, cv2.LINE_AA) + cv2.putText(img, f'{qpos[3]:.2f} elbow flexion', (700, 200), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2, cv2.LINE_AA) + cv2.putText(img, f'{qpos[4]:.2f} pronation/supination', (700, 250), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2, cv2.LINE_AA) + cv2.imshow("simulation", np.flip(img, axis=2)/255) + cv2.waitKey(1) + # sleep(0.01) + + # Run keyboard press reader in another thread + + + if args.logging: + # Output log + state_logger.save(os.path.join(evaluate_dir, args.state_log_file)) + action_logger.save(os.path.join(evaluate_dir, args.action_log_file)) + print(f'Log files have been saved files {os.path.join(evaluate_dir, args.state_log_file)}.pickle and ' + f'{os.path.join(evaluate_dir, args.action_log_file)}.pickle')