丰富抓取任务内容

src/Robot_arm_grasping_task/main.py

@@ -6,159 +6,225 @@
 import time
 from contextlib import suppress
 
-# ===================== 极简配置（剔除冗余，确保自动运行） =====================
+# ===================== 配置（已根据你的模型定制） =====================
 warnings.filterwarnings('ignore')
 CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
 MODEL_PATH = os.path.join(CURRENT_DIR, "robot.xml")
 
-# 核心参数（极简+强制）
-GRASP_FORCE = 3.8
-IK_GAIN = 1.0  # 极低增益，确保稳定
-JOINT_LIMITS = np.array([[-1.5, 1.5], [-1.2, 1.2], [-1.0, 1.0]])
-# 自动任务参数（极简流程）
-AUTO_TARGETS = [
-    np.array([0.2, 0.0, 0.08]),  # 物体位置
-    np.array([-0.1, 0.0, 0.08]),  # 放置位置
-    np.array([0.0, 0.0, 0.1])  # 归位位置
+# --- 1. 任务清单（已使用正确的物体名称 'target_object'） ---
+TASK_QUEUE = [
+    # 将名为 'target_object' 的物体移动到 (-0.3, 0, 0.05)
+    ["target_object", [-0.3, 0, 0.05]],
 ]
-STEP_PER_TARGET = 800  # 每个目标点执行步数（缩短，快速看到效果）
 
-# ===================== 全局变量（极简自动运行） =====================
-current_target_idx = 0  # 当前目标点索引
-task_step = 0  # 当前目标点内步数
-grasp_state = False  # 抓取状态
-viewer = None  # 全局viewer，确保可访问
+# --- 2. 核心控制参数 ---
+IK_GAIN = 1.5
+GRASP_FORCE = -8.0  # 夹爪闭合的力（负值表示向左/右）
+CLEARANCE_HEIGHT = 0.25  # 移动时的安全高度
+STEP_PER_MOVE = 1200  # 移动到一个新位置所需的步数
+STEP_PER_GRASP = 400  # 抓取/释放动作所需的步数
 
+# ===================== 全局状态机 =====================
+viewer = None
+current_task_index = 0
+task_step = 0
+
+
+class TaskState:
+    MOVE_TO_OBJECT_ABOVE = 1
+    MOVE_DOWN_TO_GRASP = 2
+    GRASP_OBJECT = 3
+    MOVE_UP_AFTER_GRASP = 4
+    MOVE_TO_TARGET_ABOVE = 5
+    MOVE_DOWN_TO_PLACE = 6
+    RELEASE_OBJECT = 7
+    MOVE_UP_AFTER_RELEASE = 8
+    FINISHED_ALL = 9
+
+
+current_state = TaskState.MOVE_TO_OBJECT_ABOVE
 
-# ===================== 核心逆运动学控制（极简版） =====================
-def simple_ik_control(model, data, ee_id, target_pos):
-    """极简逆运动学：只保留核心，确保不转圈+快速响应"""
-    # 获取当前末端位置
-    current_pos = data.site_xpos[ee_id] if ee_id >= 0 else np.array([0.0, 0.0, 0.1])
 
-    # 计算误差并限制
+# ===================== 核心功能函数 =====================
+def simple_ik_control(model, data, ee_id, target_pos):
+    """逆运动学控制，让末端执行器移动到目标位置"""
+    current_pos = data.site_xpos[ee_id]
     error = target_pos - current_pos
-    error = np.clip(error, -0.03, 0.03)
-
-    # 简易关节控制（直接映射，快速生效）
-    for i in range(min(3, model.njnt)):
-        # 直接更新关节角度（限制范围）
-        data.qpos[i] += error[i] * IK_GAIN * model.opt.timestep
-        data.qpos[i] = np.clip(data.qpos[i], JOINT_LIMITS[i][0], JOINT_LIMITS[i][1])
-
-    mujoco.mj_forward(model, data)
-
-
-# ===================== 强制自动运行逻辑（核心） =====================
-def run_auto_task(model, data, ee_id, obj_id):
-    """强制自动运行：启动即执行，无复杂判断"""
-    global current_target_idx, task_step, grasp_state
-
-    # 1. 执行当前目标点的控制
-    target = AUTO_TARGETS[current_target_idx]
-    simple_ik_control(model, data, ee_id, target)
-
-    # 2. 抓取/释放逻辑（极简）
-    if current_target_idx == 0 and task_step > STEP_PER_TARGET * 0.7:
-        # 到达物体位置，闭合夹爪
-        if model.nu >= 4:
-            data.ctrl[3] = min(data.ctrl[3] + 0.05, GRASP_FORCE)
-            data.ctrl[4] = max(data.ctrl[4] - 0.05, -GRASP_FORCE)
-        grasp_state = True
-    elif current_target_idx == 1 and task_step > STEP_PER_TARGET * 0.7:
-        # 到达放置位置，释放夹爪
-        if model.nu >= 4:
-            data.ctrl[3] = max(data.ctrl[3] - 0.05, 0.0)
-            data.ctrl[4] = min(data.ctrl[4] + 0.05, 0.0)
-        grasp_state = False
-
-    # 3. 切换目标点（步数到即切换）
-    task_step += 1
-    if task_step >= STEP_PER_TARGET:
-        print(f"✅ 完成目标点 {current_target_idx + 1}/{len(AUTO_TARGETS)}")
-        task_step = 0
-        current_target_idx += 1
+    error = np.clip(error, -0.05, 0.05)
+
+    jacp = np.zeros((3, model.nv))
+    mujoco.mj_jac(model, data, jacp, None, current_pos, ee_id)
+    jnt_vel = np.dot(jacp[:, :3].T, error * IK_GAIN)
+    jnt_vel = np.clip(jnt_vel, -0.5, 0.5)
+
+    # 注意：这里控制的是关节力矩（motor），而不是直接设置角度
+    for i in range(min(3, model.nu - 2)):  # 减去夹爪的两个控制
+        data.ctrl[i] = jnt_vel[i] * 100  # 乘以一个系数来放大控制信号
+
+
+def run_smart_grasp_task(model, data, ee_id):
+    """智能抓取任务的状态机逻辑"""
+    global current_task_index, task_step, current_state
+
+    if current_task_index >= len(TASK_QUEUE):
+        if current_state != TaskState.FINISHED_ALL:
+            print("\n🎉🎉🎉 所有抓取任务已成功完成！🎉🎉🎉")
+            current_state = TaskState.FINISHED_ALL
+        return False
+
+    obj_name, target_place_pos = TASK_QUEUE[current_task_index]
+    obj_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, obj_name)
+
+    if obj_id == -1:
+        print(f"❌ 错误：未在模型中找到物体 '{obj_name}'，请检查XML文件。")
+        current_task_index += 1
+        return True
+
+    # --- 状态机逻辑 ---
+    if current_state == TaskState.MOVE_TO_OBJECT_ABOVE:
+        if task_step == 0:
+            print(f"\n[任务 {current_task_index + 1}/{len(TASK_QUEUE)}] 开始处理物体: {obj_name}")
+            print("-> 状态: 移动到物体上方...")
+        target_pos = data.xpos[obj_id].copy()
+        target_pos[2] = CLEARANCE_HEIGHT
+        simple_ik_control(model, data, ee_id, target_pos)
+        if np.linalg.norm(data.site_xpos[ee_id] - target_pos) < 0.01:
+            task_step = 0
+            current_state = TaskState.MOVE_DOWN_TO_GRASP
+
+    elif current_state == TaskState.MOVE_DOWN_TO_GRASP:
+        if task_step == 0:
+            print("-> 状态: 下降以抓取物体...")
+        target_pos = data.xpos[obj_id].copy()
+        target_pos[2] += 0.05  # 停在物体表面上方一点
+        simple_ik_control(model, data, ee_id, target_pos)
+        if np.linalg.norm(data.site_xpos[ee_id] - target_pos) < 0.005:
+            task_step = 0
+            current_state = TaskState.GRASP_OBJECT
+
+    elif current_state == TaskState.GRASP_OBJECT:
+        if task_step == 0:
+            print("-> 状态: 正在抓取...")
+        # 闭合夹爪: 左爪左移(负), 右爪右移(正)
+        data.ctrl[3] = GRASP_FORCE
+        data.ctrl[4] = -GRASP_FORCE
+        if task_step > STEP_PER_GRASP:
+            task_step = 0
+            current_state = TaskState.MOVE_UP_AFTER_GRASP
+
+    elif current_state == TaskState.MOVE_UP_AFTER_GRASP:
+        if task_step == 0:
+            print("-> 状态: 抓取成功，上升...")
+        target_pos = data.site_xpos[ee_id].copy()
+        target_pos[2] = CLEARANCE_HEIGHT
+        simple_ik_control(model, data, ee_id, target_pos)
+        if np.linalg.norm(data.site_xpos[ee_id] - target_pos) < 0.01:
+            task_step = 0
+            current_state = TaskState.MOVE_TO_TARGET_ABOVE
+
+    elif current_state == TaskState.MOVE_TO_TARGET_ABOVE:
+        if task_step == 0:
+            print(f"-> 状态: 移动到目标放置区上方 {target_place_pos[:2]}...")
+        target_pos = np.array(target_place_pos)
+        target_pos[2] = CLEARANCE_HEIGHT
+        simple_ik_control(model, data, ee_id, target_pos)
+        if np.linalg.norm(data.site_xpos[ee_id] - target_pos) < 0.01:
+            task_step = 0
+            current_state = TaskState.MOVE_DOWN_TO_PLACE
+
+    elif current_state == TaskState.MOVE_DOWN_TO_PLACE:
+        if task_step == 0:
+            print("-> 状态: 下降以放置物体...")
+        target_pos = np.array(target_place_pos)
+        simple_ik_control(model, data, ee_id, target_pos)
+        if np.linalg.norm(data.site_xpos[ee_id] - target_pos) < 0.005:
+            task_step = 0
+            current_state = TaskState.RELEASE_OBJECT
+
+    elif current_state == TaskState.RELEASE_OBJECT:
+        if task_step == 0:
+            print("-> 状态: 正在释放物体...")
+        # 打开夹爪: 左右爪都回中
+        data.ctrl[3] = 0
+        data.ctrl[4] = 0
+        if task_step > STEP_PER_GRASP:
+            task_step = 0
+            current_state = TaskState.MOVE_UP_AFTER_RELEASE
+
+    elif current_state == TaskState.MOVE_UP_AFTER_RELEASE:
+        if task_step == 0:
+            print("-> 状态: 释放成功，上升并准备下一个任务...")
+        target_pos = data.site_xpos[ee_id].copy()
+        target_pos[2] = CLEARANCE_HEIGHT
+        simple_ik_control(model, data, ee_id, target_pos)
+        if np.linalg.norm(data.site_xpos[ee_id] - target_pos) < 0.01:
+            current_task_index += 1
+            task_step = 0
+            current_state = TaskState.MOVE_TO_OBJECT_ABOVE
 
-        # 所有目标点完成，退出
-        if current_target_idx >= len(AUTO_TARGETS):
-            print("\n🎉 所有自动任务强制完成！")
-            return False  # 任务完成，返回False
-    return True  # 任务继续
+    task_step += 1
+    return True
 
 
-# ===================== 初始化+主程序（强制自动） =====================
+# ===================== 主程序 =====================
 def init():
-    """极简初始化：确保快速启动"""
     global viewer
-    # 检查模型文件
     if not os.path.exists(MODEL_PATH):
-        raise FileNotFoundError(f"请确保robot.xml在当前目录：{MODEL_PATH}")
+        raise FileNotFoundError(f"请确保 'robot.xml' 文件在当前目录: {MODEL_PATH}")
 
-    # 加载模型
     model = mujoco.MjModel.from_xml_path(MODEL_PATH)
     data = mujoco.MjData(model)
 
-    # 初始化关节到中间位置（避免初始转圈）
-    for i in range(min(3, model.njnt)):
-        data.qpos[i] = (JOINT_LIMITS[i][0] + JOINT_LIMITS[i][1]) / 2
-    mujoco.mj_forward(model, data)
-
-    # 初始化Viewer（强制显示）
     viewer = mujoco_viewer.MujocoViewer(model, data, hide_menus=True)
-    viewer.cam.distance = 1.5
-    viewer.cam.elevation = 20
-    viewer.cam.azimuth = 70
-    viewer.cam.lookat = [0.1, 0.0, 0.1]
+    viewer.cam.distance = 2.0
+    viewer.cam.elevation = -20
+    viewer.cam.azimuth = 90
+    viewer.cam.lookat = [0.2, 0.0, 0.1]
 
-    # 极简ID识别（只找关键ID）
     ee_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, "ee_site")
-    obj_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "target_object")
+    if ee_id == -1:
+        raise ValueError("模型中必须包含一个名为 'ee_site' 的site。")
 
-    # 打印强制启动提示
-    print("=" * 50)
-    print("🚨 强制自动运行模式启动！")
-    print("📌 无需任何按键，立刻执行抓取任务")
-    print("🎯 目标点：物体位置→放置位置→归位")
-    print("=" * 50)
-    return model, data, ee_id, obj_id
+    print("=" * 60)
+    print("🚀 全自动智能抓取系统启动！")
+    print(f"📋 任务清单: 共 {len(TASK_QUEUE)} 个物体需要处理。")
+    print("💡 正在连接到模型 'simple_arm'...")
+    print("=" * 60)
+    return model, data, ee_id
 
 
 def main():
     global viewer
     try:
-        # 初始化
-        model, data, ee_id, obj_id = init()
+        model, data, ee_id = init()
 
-        # 强制自动运行核心循环（无任何按键依赖）
         while viewer.is_alive:
-            # 执行自动任务，返回False则退出
-            if not run_auto_task(model, data, ee_id, obj_id):
+            if not run_smart_grasp_task(model, data, ee_id):
                 break
 
-            # 仿真步进（快速渲染）
             mujoco.mj_step(model, data)
             viewer.render()
             time.sleep(0.005)
 
-        # 任务完成后，保持窗口3秒
-        print("\n⏳ 任务完成，3秒后自动退出...")
-        for _ in range(3):
+        print("\n⏳ 所有任务已完成，窗口将在5秒后自动关闭。")
+        for _ in range(5):
             viewer.render()
             time.sleep(1)
 
     except Exception as e:
-        print(f"\n❌ 错误：{e}")
+        print(f"\n❌ 程序发生错误: {e}")
+        import traceback
+        traceback.print_exc()
     finally:
         with suppress(Exception):
             viewer.close()
-        print("🔚 强制自动运行结束")
+        print("🔚 程序已退出。")
 
 
 if __name__ == "__main__":
-    # 强制检查依赖并启动
     try:
         import mujoco, mujoco_viewer
     except ImportError:
-        print("❌ 缺少依赖！执行：pip install mujoco mujoco-viewer numpy")
+        print("❌ 缺少依赖！请运行: pip install mujoco mujoco-viewer numpy")
         exit(1)
     main()