# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

from __future__ import annotations

import torch
from typing import TYPE_CHECKING

from isaaclab.assets import Articulation, RigidObject
from isaaclab.managers import SceneEntityCfg
from isaaclab.utils.math import wrap_to_pi, quat_error_magnitude, quat_apply, normalize

if TYPE_CHECKING:
    from isaaclab.envs import ManagerBasedRLEnv


def joint_pos_target_l2(env: ManagerBasedRLEnv, target: float, asset_cfg: SceneEntityCfg) -> torch.Tensor:
    """Penalize joint position deviation from a target value."""
    # extract the used quantities (to enable type-hinting)
    asset: Articulation = env.scene[asset_cfg.name]
    # wrap the joint positions to (-pi, pi)
    joint_pos = wrap_to_pi(asset.data.joint_pos[:, asset_cfg.joint_ids])
    # compute the reward
    return torch.sum(torch.square(joint_pos - target), dim=1)


def lid_is_lifted(env: ManagerBasedRLEnv, minimal_height: float,
                  lid_cfg: SceneEntityCfg = SceneEntityCfg("DexCube")) -> torch.Tensor:
    """Reward the agent for lifting the lid above the minimal height."""
    lid: RigidObject = env.scene[lid_cfg.name]
    # root_pos_w shape: (num_envs, 3) - extract z-coordinate: (num_envs,)
    height = lid.data.root_pos_w[:, 2]
    return torch.where(height > minimal_height, 1.0, 0.0)

def gripper_lid_distance(env: ManagerBasedRLEnv, std: float,
                         lid_cfg: SceneEntityCfg = SceneEntityCfg("DexCube"),
                         robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
                         gripper_body_name: str = "left_hand_body") -> torch.Tensor:
    """Reward the agent for reaching the lid using tanh-kernel."""
    # extract the used quantities (to enable type-hinting)
    lid: RigidObject = env.scene[lid_cfg.name]
    robot: Articulation = env.scene[robot_cfg.name]

    # float 化，防止成为 (3,) 向量
    if not isinstance(std, float):
        if torch.is_tensor(std):
            std = std.item()
        else:
            std = float(std)

    # Target lid position: (num_envs, 3)
    lid_pos_w = lid.data.root_pos_w[:, :3]
    
    # Gripper position: (num_envs, 3)
    # body_idx = robot.find_bodies(gripper_body_name)[0]
    # gripper_pos_w = robot.data.body_pos_w[:, body_idx, :].squeeze(dim=1)
    body_ids, _ = robot.find_bodies([gripper_body_name], preserve_order=True)
    gripper_pos_w = robot.data.body_pos_w[:, body_ids[0], :]
    # Distance of the gripper to the lid: (num_envs,)
    distance = torch.norm(lid_pos_w - gripper_pos_w, dim=1) + 1e-6
    
    return 1 - torch.tanh(distance / std)


def lid_grasped(env: ManagerBasedRLEnv,
                distance_threshold: float = 0.05,
                lid_cfg: SceneEntityCfg = SceneEntityCfg("DexCube"),
                robot_cfg: SceneEntityCfg = SceneEntityCfg     ("Mindbot"),
                gripper_body_name: str = "left_hand_body") -> torch.Tensor:
    """Reward the agent for grasping the lid (close and lifted)."""
    # extract the used quantities (to enable type-hinting)
    lid: RigidObject = env.scene[lid_cfg.name]
    robot: Articulation = env.scene[robot_cfg.name]
    
    # Target lid position: (num_envs, 3)
    lid_pos_w = lid.data.root_pos_w[:, :3]
    
    # Gripper position: (num_envs, 3)
    body_idx = robot.find_bodies(gripper_body_name)[0]
    gripper_pos_w = robot.data.body_pos_w[:, body_idx, :]
    
    # Distance of the gripper to the lid: (num_envs,)
    distance = torch.norm(lid_pos_w - gripper_pos_w, dim=1)
    
    # Check if close and lifted
    is_close = distance < distance_threshold
    is_lifted = lid.data.root_pos_w[:, 2] > 0.1
    
    return torch.where(is_close & is_lifted, 1.0, 0.0)

def gripper_lid_distance_combined(env: ManagerBasedRLEnv, std: float, minimal_height: float,
                                  lid_cfg: SceneEntityCfg = SceneEntityCfg("DexCube"),
                                  robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
                                  gripper_body_name: str = "left_hand_body") -> torch.Tensor:
    """Combined reward for reaching the lid AND lifting it."""
    # Get reaching reward
    reach_reward = gripper_lid_distance(env, std, lid_cfg, robot_cfg, gripper_body_name)
    # Get lifting reward
    lift_reward = lid_is_lifted(env, minimal_height, lid_cfg)
    # Combine rewards multiplicatively
    return reach_reward * lift_reward


def gripper_handle_position_alignment(env: ManagerBasedRLEnv,
                                    dryingbox_cfg: SceneEntityCfg = SceneEntityCfg("dryingbox"),
                                    robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
                                    left_gripper_name: str = "right_hand_l",
                                    right_gripper_name: str = "right_hand__r",
                                    handle_body_name: str = "Equipment_BB_13_C",
                                    height_offset: float = 0.02,
                                    std: float = 0.1) -> torch.Tensor:
    """奖励函数：夹爪相对于lid的位置对齐（简化版）。
    
    计算夹爪中心到目标点（lid上方height_offset处）的距离奖励。
    
    Args:
        env: 环境实例
        dringbox_cfg: lid的配置
        robot_cfg: 机器人的配置
        left_gripper_name: 左侧夹爪body名称
        right_gripper_name: 右侧夹爪body名称
        handle_body_name: 把手body名称
        height_offset: Y方向的期望高度偏移（目标点在lid上方）
        std: tanh核函数的标准差
        
    Returns:
        位置对齐奖励 (num_envs,)
    """
    dryingbox: Articulation = env.scene[dryingbox_cfg.name]
    robot: Articulation = env.scene[robot_cfg.name]

    db_ids,_=dryingbox.find_bodies([handle_body_name],preserve_order=True)
    handle_pos_w=dryingbox.data.body_pos_w[:, db_ids[0], :]
    # lid_pos_w = lid.data.root_pos_w[:, :3]
    
    # 目标点：lid位置 + height_offset（Z方向）
    target_pos =handle_pos_w.clone()
    target_pos[:, 1] += height_offset  # 在lid上方height_offset处
    # import pdb; pdb.set_trace()
    # print(target_pos[0])
    target_pos[:, 2] -= height_offset  # 在lid上方height_offset处
    # print(target_pos[0])
    # 获取两个夹爪的位置
    left_body_ids, _ = robot.find_bodies([left_gripper_name], preserve_order=True)
    right_body_ids, _ = robot.find_bodies([right_gripper_name], preserve_order=True)
    
    left_gripper_pos = robot.data.body_pos_w[:, left_body_ids[0], :]  # (num_envs, 3)
    right_gripper_pos = robot.data.body_pos_w[:, right_body_ids[0], :]  # (num_envs, 3)
    
    # 计算夹爪中心位置
    gripper_center = (left_gripper_pos + right_gripper_pos) / 2.0  # (num_envs, 3)
    # print(gripper_center[0])
    # 计算夹爪中心到目标点的3D距离
    distance = torch.norm(gripper_center - target_pos, dim=1) + 1e-6
    
    # 距离奖励：距离越小，奖励越大
    position_reward = 1.0 - torch.tanh(distance / std)
    
    return position_reward


# def gripper_lid_orientation_alignment(
#     env,
#     lid_cfg: SceneEntityCfg = SceneEntityCfg("lid"),
#     robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
#     gripper_body_name: str = "left_hand_body",
#     gripper_forward_axis: tuple = (0.0, 0.0, 1.0),  # 修正为 Z
#     gripper_width_axis: tuple = (0.0, 1.0, 0.0),    # 修正为 Y (用于与把手平行)
#     lid_handle_axis: tuple = (0.0, 1.0, 0.0),       # 修正为 Y
# ) -> torch.Tensor:
    
#     # 1. 获取对象
#     lid = env.scene[lid_cfg.name]
#     robot = env.scene[robot_cfg.name]
    
#     # 2. 获取姿态
#     lid_quat_w = lid.data.root_quat_w
#     body_ids, _ = robot.find_bodies([gripper_body_name], preserve_order=True)
#     gripper_quat_w = robot.data.body_quat_w[:, body_ids[0], :]

#     # 3. 垂直对齐 (Vertical Alignment)
#     # 目标：夹爪 +Y 指向 World -Z
#     grip_fwd_local = torch.tensor(gripper_forward_axis, device=env.device).repeat(env.num_envs, 1)
#     grip_fwd_world = quat_apply(gripper_quat_w, grip_fwd_local)
#     # print(grip_fwd_world[0])
#     world_down = torch.tensor([0.0, 0.0, -1.0], device=env.device).repeat(env.num_envs, 1)
    
#     dot_vertical = torch.sum(grip_fwd_world * world_down, dim=1)
#     dot_vertical = torch.clamp(dot_vertical, -1.0, 1.0)  #add
#     # 映射 [-1, 1] -> [0, 1]，且全程有梯度
#     # 即使是朝天 (+1)，只要稍微往下转一点，分数就会增加
#     # rew_vertical = (dot_vertical + 1.0) / 2.0
#     val_vertical = torch.clamp((dot_vertical + 1.0) / 2.0, 0.0, 1.0)
#     rew_vertical = torch.pow(val_vertical, 2) + 1e-4
#     # nan -> 0
#     rew_vertical = torch.where(torch.isnan(rew_vertical), torch.zeros_like(rew_vertical), rew_vertical)

#     # 4. 偏航对齐 (Yaw Alignment)
#     # 目标：夹爪 +Z 平行于 Lid +Y
#     grip_width_local = torch.tensor(gripper_width_axis, device=env.device).repeat(env.num_envs, 1)
#     grip_width_world = quat_apply(gripper_quat_w, grip_width_local)
    
#     lid_handle_local = torch.tensor(lid_handle_axis, device=env.device).repeat(env.num_envs, 1)
#     lid_handle_world = quat_apply(lid_quat_w, lid_handle_local)
    
#     dot_yaw = torch.sum(grip_width_world * lid_handle_world, dim=1)
#     rew_yaw = torch.square(torch.abs(dot_yaw)) + 1e-4
#     # nan -> 0
#     rew_yaw = torch.where(torch.isnan(rew_yaw), torch.zeros_like(rew_yaw), rew_yaw)

#     # 5. 组合
#     total_reward = rew_vertical * rew_yaw

#     total_reward=torch.nan_to_num(total_reward, nan=0.0, posinf=0.0, neginf=0.0)

#     # debug
#     if not torch.isfinite(total_reward).all():
#         print(f"[DEBUG] Non-finite reward detected in gripper_lid_orientation_alignment: {total_reward}")
#         print(f"rew_vertical: {rew_vertical}")
#         print(f"rew_yaw: {rew_yaw}")
        
#     return total_reward


# def gripper_handle_orientation_alignment(
#     env,
#     dryingbox_cfg: SceneEntityCfg = SceneEntityCfg("dryingbox"),
#     robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
#     gripper_body_name: str = "right_hand_body",
#     handle_body_name: str = "Equipment_BB_13_C", 
#     max_angle_deg: float = 60.0,
# ) -> torch.Tensor:
#     """
#     特定姿态对齐奖励：
#     1. [X-Y Alignment] 夹爪 X 轴 (1,0,0) 必须平行于 烘箱手柄 Y 轴 (0,1,0)。
#     2. [Z-Z Alignment] 夹爪 Z 轴 (0,0,1) 必须同向于 烘箱本体 -Z 轴 (0,0,-1)。
#     """
    
#     # 1. 获取各个对象
#     dryingbox = env.scene[dryingbox_cfg.name]
#     robot = env.scene[robot_cfg.name]

#     # 2. 获取姿态四元数
#     # A. 夹爪
#     body_ids, _ = robot.find_bodies([gripper_body_name], preserve_order=True)
#     gripper_quat_w = robot.data.body_quat_w[:, body_ids[0], :]

#     # B. 烘箱手柄 (Handle Body)
#     handle_ids, _ = dryingbox.find_bodies([handle_body_name], preserve_order=True)
#     handle_quat_w = dryingbox.data.body_quat_w[:, handle_ids[0], :]

#     # C. 烘箱本体根节点 (Drying Box Root)
#     box_quat_w = dryingbox.data.root_quat_w

#     # ============================================================
#     # 目标 1: 夹爪 X 轴 (1,0,0) 平行于 手柄 Y 轴 (0,1,0)
#     # ============================================================
    
#     # 夹爪局部 X 轴 -> 世界坐标
#     grip_x_local = torch.tensor([1.0, 0.0, 0.0], device=env.device).repeat(env.num_envs, 1)
#     grip_x_world = quat_apply(gripper_quat_w, grip_x_local)

#     # 手柄局部 Y 轴 -> 世界坐标
#     handle_y_local = torch.tensor([0.0, 1.0, 0.0], device=env.device).repeat(env.num_envs, 1)
#     handle_y_world = quat_apply(handle_quat_w, handle_y_local)

#     # 计算点积 (取绝对值，允许正反两个方向平行)
#     dot_xy = torch.abs(torch.sum(grip_x_world * handle_y_world, dim=1))
    
#     # 角度偏差计算
#     dot_xy = torch.clamp(dot_xy, 0.0, 1.0)
#     angle_xy = torch.acos(dot_xy)
    
#     # 归一化奖励
#     max_angle_rad = torch.tensor(max_angle_deg * 3.14159 / 180.0, device=env.device)
#     rew_xy = 1.0 - (angle_xy / max_angle_rad)
#     rew_xy = torch.where(angle_xy <= max_angle_rad, rew_xy, torch.zeros_like(rew_xy))

#     # ============================================================
#     # 目标 2: 夹爪 Z 轴 (0,0,1) 同向于 烘箱 -Z 轴 (0,0,-1)
#     # ============================================================
    
#     # 夹爪局部 Z 轴 -> 世界坐标
#     grip_z_local = torch.tensor([0.0, 0.0, 1.0], device=env.device).repeat(env.num_envs, 1)
#     grip_z_world = quat_apply(gripper_quat_w, grip_z_local)

#     # 烘箱局部 -Z 轴 -> 世界坐标
#     # 注意：这里使用的是 dryingbox 的根节点姿态
#     box_neg_z_local = torch.tensor([0.0, 0.0, -1.0], device=env.device).repeat(env.num_envs, 1)
#     box_neg_z_world = quat_apply(box_quat_w, box_neg_z_local)

#     # 计算点积 (不取绝对值，要求方向严格相同)
#     # 我们希望 dot 接近 1.0
#     dot_zz = torch.sum(grip_z_world * box_neg_z_world, dim=1)
    
#     # 角度偏差计算 (clip -1 到 1 防止 nan)
#     dot_zz = torch.clamp(dot_zz, -1.0, 1.0)
#     angle_zz = torch.acos(dot_zz) # 0 是同向, pi 是反向

#     # 奖励计算
#     rew_zz = 1.0 - (angle_zz / max_angle_rad)
#     rew_zz = torch.where(angle_zz <= max_angle_rad, rew_zz, torch.zeros_like(rew_zz))
#     # 增强 Z 轴方向的敏感度
#     rew_zz = torch.pow(rew_zz, 2)

#     # ============================================================
#     # 组合奖励
#     # ============================================================
#     total_reward = rew_xy * rew_zz
#     # import pdb; pdb.set_trace()
#     return torch.nan_to_num(total_reward, nan=0.0)


def gripper_handle_orientation_alignment(
    env,
    dryingbox_cfg: SceneEntityCfg = SceneEntityCfg("dryingbox"),
    robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
    handle_name: str = "Equipment_BB_13_C",
    gripper_body_name: str = "right_hand_body",
    gripper_forward_axis: tuple = (0.0, 0.0, 1.0),  # 夹爪指尖指向
    gripper_width_axis: tuple = (0.0, 1.0, 0.0),    # 夹爪手指开合轴
    max_angle_deg: float = 30.0,                    # 容忍角度缩小以提高精度
) -> torch.Tensor:
    """
    针对垂直手柄的姿态对齐奖励：
    1. 夹爪前向轴 (Local Z) 需指向烘箱方向 (World -X)
    2. 夹爪宽度轴 (Local Y) 需保持水平 (垂直于 World Z)，以实现水平夹取垂直手柄
    """
    
    # 1. 获取场景对象
    dryingbox = env.scene[dryingbox_cfg.name]
    robot = env.scene[robot_cfg.name]
    
    # 2. 获取夹爪的世界坐标系姿态 (Quaternion)
    body_ids, _ = robot.find_bodies([gripper_body_name], preserve_order=True)
    gripper_quat_w = robot.data.body_quat_w[:, body_ids[0], :]

    # --- 约束 A: 前向指向约束 (Forward Alignment) ---
    # 夹爪的指尖 (Local Z) 应该指向烘箱 (-X 方向)
    grip_fwd_local = torch.tensor(gripper_forward_axis, device=env.device).repeat(env.num_envs, 1)
    grip_fwd_world = quat_apply(gripper_quat_w, grip_fwd_local)
    
    # 目标方向是世界坐标系的 -X
    target_fwd = torch.tensor([1.0, 0.0, 0.0], device=env.device).repeat(env.num_envs, 1)
    import pdb;pdb.set_trace()
    # 计算点积：1.0 代表完全对准 -X，-1.0 代表背对
    dot_fwd = torch.sum(grip_fwd_world * target_fwd, dim=1)
    # 将奖励映射到 [0, 1] 空间
    rew_forward = torch.clamp((dot_fwd - torch.cos(torch.tensor(max_angle_deg * torch.pi / 180.0))) / (1.0 - torch.cos(torch.tensor(max_angle_deg * torch.pi / 180.0))), 0.0, 1.0)

    # --- 约束 B: 水平开合约束 (Roll/Horizontal Alignment) ---
    # 夹爪的宽度轴 (Local Y，即手指张开的方向) 必须是水平的 (垂直于 World Z)
    # 这样才能夹住竖直方向 (Z轴) 的手柄
    grip_width_local = torch.tensor(gripper_width_axis, device=env.device).repeat(env.num_envs, 1)
    grip_width_world = quat_apply(gripper_quat_w, grip_width_local)
    
    world_up = torch.tensor([0.0, 0.0, 1.0], device=env.device).repeat(env.num_envs, 1)
    
    # 我们希望 grip_width_world 与 world_up 的点积为 0 (互相垂直)
    dot_width_z = torch.abs(torch.sum(grip_width_world * world_up, dim=1))
    # 点积越接近 0，奖励越高
    rew_horizontal_fingers = 1.0 - torch.clamp(dot_width_z, 0.0, 1.0)

    # 3. 组合奖励
    # 使用乘法：只有当“指针对了”且“手指水平了”时，才给予高分
    total_reward = rew_forward * rew_horizontal_fingers
    
    # 异常处理
    total_reward = torch.nan_to_num(total_reward, nan=0.0)
        
    return total_reward





# def gripper_handle_orientation_alignment(
#     env,
#     dryingbox_cfg: SceneEntityCfg = SceneEntityCfg("dryingbox"),
#     robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
#     handle_name: str = "Equipment_BB_13_C",  # 修复 1: 补上了逗号
#     gripper_body_name: str = "right_hand_body",
#     gripper_forward_axis: tuple = (0.0, 0.0, 1.0),  # 夹爪自身的 Z 轴（指向前方）
#     gripper_width_axis: tuple = (0.0, 1.0, 0.0),    # 夹爪自身的 Y 轴（指向手指开合方向）
#     handle_axis: tuple = (0.0, 1.0, 0.0),           # 盖子把手轴向
#     max_angle_deg: float = 60.0,                    # 允许的角度偏差
# ) -> torch.Tensor:
    
#     # 1. 获取对象
#     dryingbox = env.scene[dryingbox_cfg.name]
#     robot = env.scene[robot_cfg.name]
    
#     # 2. 获取姿态
#     # 修复 2: 统一使用 handle_name
#     # 注意：如果 handle_name 为空字符串，find_bodies 可能会报错，需确保调用时传入了正确名称
#     if handle_name == "":
#         print("[Warning] handle_name is empty in gripper_handle_orientation_alignment")
        
#     handle_ids, _ = dryingbox.find_bodies([handle_name], preserve_order=True)
#     handle_quat_w = dryingbox.data.body_quat_w[:, handle_ids[0], :]
    
#     body_ids, _ = robot.find_bodies([gripper_body_name], preserve_order=True)
#     gripper_quat_w = robot.data.body_quat_w[:, body_ids[0], :]

#     # 3. 水平对齐 (Horizontal Alignment)
#     # 目标：夹爪前向轴 (Z) 应与世界坐标系 Z 轴垂直 (夹角 90度)
    
#     grip_fwd_local = torch.tensor(gripper_forward_axis, device=env.device).repeat(env.num_envs, 1)
#     grip_fwd_world = quat_apply(gripper_quat_w, grip_fwd_local)
    
#     world_up = torch.tensor([0.0, 0.0, 1.0], device=env.device).repeat(env.num_envs, 1)
    
#     # 计算点积
#     dot_z = torch.sum(grip_fwd_world * world_up, dim=1)
#     dot_z = torch.clamp(dot_z, -1.0, 1.0)
    
#     # 计算当前夹爪与“垂直线”的夹角
#     angle_from_vertical_rad = torch.acos(dot_z) # 范围 [0, pi]
    
#     # 我们希望这个角度接近 90度 (pi/2)
#     target_angle = torch.pi / 2.0
#     angle_error = torch.abs(angle_from_vertical_rad - target_angle)
    
#     # 将 max_angle_deg 转换为弧度
#     max_error_rad = torch.tensor(max_angle_deg * 3.14159265359 / 180.0, device=env.device)
    
#     # 计算奖励
#     angle_normalized = torch.clamp(angle_error / max_error_rad, 0.0, 1.0)
#     rew_horizontal = 1.0 - angle_normalized
    
#     rew_horizontal = torch.where(angle_error <= max_error_rad, rew_horizontal, torch.zeros_like(rew_horizontal))
#     rew_horizontal = torch.where(torch.isnan(rew_horizontal), torch.zeros_like(rew_horizontal), rew_horizontal)

#     # 4. 偏航对齐 (Yaw Alignment)
#     grip_width_local = torch.tensor(gripper_width_axis, device=env.device).repeat(env.num_envs, 1)
#     grip_width_world = quat_apply(gripper_quat_w, grip_width_local)
    
#     handle_local = torch.tensor(handle_axis, device=env.device).repeat(env.num_envs, 1)
#     handle_world = quat_apply(handle_quat_w, handle_local)
    
#     dot_yaw = torch.sum(grip_width_world * handle_world, dim=1)
#     rew_yaw = torch.square(torch.abs(dot_yaw)) + 1e-4
#     rew_yaw = torch.where(torch.isnan(rew_yaw), torch.zeros_like(rew_yaw), rew_yaw)

#     # 5. 组合
#     total_reward = rew_horizontal * rew_yaw
#     total_reward = torch.nan_to_num(total_reward, nan=0.0, posinf=0.0, neginf=0.0)
        
#     return total_reward


# def gripper_lid_orientation_alignment(
#     env,
#     lid_cfg: SceneEntityCfg = SceneEntityCfg("lid"),
#     robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
#     gripper_body_name: str = "right_hand_body",
#     gripper_forward_axis: tuple = (0.0, 0.0, 1.0),  # 修正为 Z
#     gripper_width_axis: tuple = (0.0, 1.0, 0.0),    # 修正为 Y (用于与把手平行)
#     handle_axis: tuple = (0.0, 1.0, 0.0),       # 修正为 Y
#     max_angle_deg: float = 60.0,  # 允许的最大角度偏差（度）
# ) -> torch.Tensor:
    
#     # 1. 获取对象
#     lid = env.scene[lid_cfg.name]
#     robot = env.scene[robot_cfg.name]
    
#     # 2. 获取姿态
#     lid_quat_w = lid.data.root_quat_w
#     body_ids, _ = robot.find_bodies([gripper_body_name], preserve_order=True)
#     gripper_quat_w = robot.data.body_quat_w[:, body_ids[0], :]
#     # 3. 水平对齐 (Level Alignment) - 确保夹爪水平伸出
#     # ============================================================
#     # 目标：夹爪的前向轴 (Z) 应该垂直于 世界的 Up 轴 (Z)。
#     # 也就是夹爪是指向地平线的，不能指天或指地。
    
#     # 将夹爪局部前向向量转到世界坐标系
#     grip_fwd_local = torch.tensor(gripper_forward_axis, device=env.device).repeat(env.num_envs, 1)
#     grip_fwd_world = quat_apply(gripper_quat_w, grip_fwd_local)
    
#     # 世界坐标系的垂直向上向量 (0, 0, 1)
#     world_up = torch.tensor([0.0, 0.0, 1.0], device=env.device).repeat(env.num_envs, 1)
    
#     # 计算点积：
#     # 如果夹爪完全水平，与 Z 轴夹角 90 度，cos(90) = 0
#     # 如果夹爪向下/向上，与 Z 轴夹角 0/180 度，cos = 1/-1
#     dot_level = torch.sum(grip_fwd_world * world_up, dim=1)
    
#     # 我们希望 dot_level 越接近 0 越好
#     # 计算偏差角度 (用于阈值截断)
#     # asin(dot) 得到的是偏离水平面的角度 (因为 dot 是与法线的余弦，即与平面的正弦)
#     angle_from_horizon = torch.asin(torch.clamp(torch.abs(dot_level), 0.0, 1.0))
#     max_angle_rad = torch.tensor(max_angle_deg * 3.14159 / 180.0, device=env.device)
    
#     # 奖励计算：1.0 - |dot|
#     rew_level = 1.0 - torch.abs(dot_level)
    
#     # 阈值截断：如果仰头或低头超过 max_angle_deg，奖励归零
#     rew_level = torch.where(angle_from_horizon < max_angle_rad, rew_level, torch.zeros_like(rew_level))

#     # ============================================================
#     # 4. 滚转对齐 (Roll Alignment) - 确保手指左右夹击
#     # ============================================================
#     # 场景分析：
#     # 把手是竖直的 (Handle Axis 在世界系中接近 Z 轴方向)
#     # 我们希望机器人的两个手指在把手的左右两侧 (Horizontal Grasp)
#     # 这意味着夹爪的宽度方向 (Width Axis/Y轴) 应该是水平的
#     # 因此：夹爪的 Width Axis 应该与 把手的 Handle Axis **垂直**
    
#     # 将夹爪局部宽度向量转到世界坐标系
#     grip_width_local = torch.tensor(gripper_width_axis, device=env.device).repeat(env.num_envs, 1)
#     grip_width_world = quat_apply(gripper_quat_w, grip_width_local)
    
#     # 将把手局部长轴向量转到世界坐标系
#     handle_axis_local = torch.tensor(handle_axis, device=env.device).repeat(env.num_envs, 1)
#     handle_long_world = quat_apply(handle_quat_w, handle_axis_local)
    
#     # 计算点积
#     # 我们希望它们垂直 (90度) -> dot 接近 0
#     dot_cross = torch.sum(grip_width_world * handle_long_world, dim=1)
    
#     # 奖励计算：1.0 - |dot|
#     # 如果平行 (dot=1)，说明手指是一上一下去抓竖直杆子（姿势别扭），得 0 分
#     # 如果垂直 (dot=0)，说明手指是一左一右去抓竖直杆子（姿势正确），得 1 分
#     rew_roll = 1.0 - torch.abs(dot_cross)
    
#     # 同样可以加一个阈值截断，或者使用平方让梯度更陡峭
#     rew_roll = torch.square(rew_roll)

#     # ============================================================
#     # 5. 组合奖励
#     # ============================================================
#     total_reward = rew_level * rew_roll

#     # 处理 NaN
#     total_reward = torch.nan_to_num(total_reward, nan=0.0, posinf=0.0, neginf=0.0)
        
#     return total_reward


# def gripper_open_close_reward(
#     env: ManagerBasedRLEnv,
#     lid_cfg: SceneEntityCfg = SceneEntityCfg("lid"),
#     robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
#     left_gripper_name: str = "left_hand_body", # 用于计算距离的 body
#     # 注意：确保 joint_names 只包含那两个夹爪关节
#     joint_names: list = ["left_hand_joint_left", "left_hand_joint_right"], 
#     close_threshold: float = 0.05, # 5cm 以内视为“近”，需要闭合
#     target_open_pos: float = 0.0,  # 【修正】张开是 0.0
#     target_close_pos: float = 0.03 # 【修正】闭合是 0.03
# ) -> torch.Tensor:
#     """
#     逻辑很简单：
#     1. 如果 距离 < close_threshold：目标关节位置 = 0.03 (闭合)
#     2. 如果 距离 >= close_threshold：目标关节位置 = 0.0 (张开)
#     3. 返回 -abs(当前关节 - 目标关节)，即惩罚误差。
#     """
    
#     # 1. 获取对象
#     lid = env.scene[lid_cfg.name]
#     robot = env.scene[robot_cfg.name]
    
#     # 2. 计算距离 (末端 vs Lid中心)
#     lid_pos_w = lid.data.root_pos_w[:, :3]
    
#     body_ids, _ = robot.find_bodies([left_gripper_name])
#     gripper_pos_w = robot.data.body_pos_w[:, body_ids[0], :3]
    
#     # distance: (num_envs,)
#     distance = torch.norm(gripper_pos_w - lid_pos_w, dim=-1)
    
#     # 3. 动态确定目标关节值
#     # 如果距离小于阈值，目标设为 close_pos，否则设为 open_pos
#     # target_val: (num_envs, 1)
#     # print(distance)
#     target_val = torch.where(
#         distance < close_threshold,
#         torch.tensor(target_close_pos, device=env.device),
#         torch.tensor(target_open_pos, device=env.device)
#     ).unsqueeze(1)
    
#     # 4. 获取当前关节位置
#     joint_ids, _ = robot.find_joints(joint_names)
#     # joint_pos: (num_envs, 2)
#     current_joint_pos = robot.data.joint_pos[:, joint_ids]
    
#     # 取绝对值（防止左指是负数的情况，虽然你的配置里范围看起来都是正的，加上abs更保险）
#     current_joint_pos_abs = torch.abs(current_joint_pos)
    
#     # 5. 计算误差
#     # error: (num_envs,) -> 两个手指误差的平均值
#     error = torch.mean(torch.abs(current_joint_pos_abs - target_val), dim=1)
    
#     # # 6. 返回负奖励 (惩罚)
#     # # 误差越大，惩罚越大。完全符合要求时奖励为 0。
#     # return -error
#     # =====================================================
#     # 🚀 核心修复：安全输出
#     # =====================================================
#     # 1. 使用 tanh 限制数值范围在 [-1, 0]
#     reward = -torch.tanh(error / 0.01)
    
#     # 2. 过滤 NaN！这一步能救命！
#     # 如果物理引擎算出 NaN，这里会把它变成 0.0，防止训练崩溃
#     reward = torch.nan_to_num(reward, nan=0.0, posinf=0.0, neginf=-1.0)
    
#     return reward


# def gripper_open_close_reward(
#     env: ManagerBasedRLEnv,
#     lid_cfg: SceneEntityCfg = SceneEntityCfg("lid"),
#     robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
#     left_gripper_name: str = "left_hand_body", 
#     joint_names: list = ["left_hand_joint_left", "left_hand_joint_right"], 
#     close_threshold: float = 0.05, 
#     target_open_pos: float = 0.0,
#     target_close_pos: float = 0.03,
#     height_offset: float = 0.06,  # 新增：与位置奖励保持一致
# ) -> torch.Tensor:
#     """
#     逻辑：
#     1. 计算夹爪到【目标抓取点】(Lid上方 height_offset 处) 的距离。
#     2. 如果距离 < close_threshold：目标关节位置 = 闭合。
#     3. 否则：目标关节位置 = 张开。
#     """
    
#     # 1. 获取对象
#     lid = env.scene[lid_cfg.name]
#     robot = env.scene[robot_cfg.name]
    
#     # 2. 计算目标抓取点位置 (Lid位置 + Z轴偏移)
#     lid_pos_w = lid.data.root_pos_w[:, :3].clone()
#     lid_pos_w[:, 2] += height_offset # 修正为抓取点位置
    
#     # 获取夹爪位置
#     body_ids, _ = robot.find_bodies([left_gripper_name])
#     gripper_pos_w = robot.data.body_pos_w[:, body_ids[0], :3]
    
#     # 3. 计算距离 (夹爪 vs 目标抓取点)
#     distance = torch.norm(gripper_pos_w - lid_pos_w, dim=-1)
    
#     # 4. 动态确定目标关节值
#     # 靠近抓取点 -> 闭合；远离 -> 张开
#     target_val = torch.where(
#         distance < close_threshold,
#         torch.tensor(target_close_pos, device=env.device),
#         torch.tensor(target_open_pos, device=env.device)
#     ).unsqueeze(1)
    
#     # 5. 获取当前关节位置
#     joint_ids, _ = robot.find_joints(joint_names)
#     current_joint_pos = robot.data.joint_pos[:, joint_ids]
#     current_joint_pos_abs = torch.abs(current_joint_pos)
    
#     # 6. 计算误差并返回奖励
#     error = torch.mean(torch.abs(current_joint_pos_abs - target_val), dim=1)
    
#     # 使用 tanh 限制数值范围
#     reward = 1.0 - torch.tanh(error / 0.01) # 误差越小奖励越大
    
#     return torch.nan_to_num(reward, nan=0.0)


def _is_grasped(
    env: ManagerBasedRLEnv,
    lid_cfg: SceneEntityCfg,
    robot_cfg: SceneEntityCfg,
    left_gripper_name: str,
    right_gripper_name: str,
    joint_names: list,
    height_offset: float,
    grasp_radius: float,
    target_close_pos: float
) -> torch.Tensor:
    """
    内部辅助函数：判定是否成功抓取。
    条件：
    1. 夹爪中心在把手目标点附近 (Sphere Check)
    2. 夹爪处于闭合发力状态 (Joint Effort Check)
    3. (关键) 夹爪Z轴高度合适，不能压在盖子上面 (Z-Axis Check)
    """
    # 1. 获取对象
    lid = env.scene[lid_cfg.name]
    robot = env.scene[robot_cfg.name]
    
    # 2. 目标点位置 (把手中心)
    lid_pos_w = lid.data.root_pos_w[:, :3].clone()
    lid_pos_w[:, 2] += height_offset 
    
    # 3. 夹爪位置
    l_ids, _ = robot.find_bodies([left_gripper_name])
    r_ids, _ = robot.find_bodies([right_gripper_name])
    pos_L = robot.data.body_pos_w[:, l_ids[0], :3]
    pos_R = robot.data.body_pos_w[:, r_ids[0], :3]
    gripper_center = (pos_L + pos_R) / 2.0
    
    # 4. 条件一：距离判定 (在把手球范围内)
    dist_to_handle = torch.norm(gripper_center - lid_pos_w, dim=-1)
    is_near = dist_to_handle < grasp_radius
    
    # 5. 条件二：Z轴高度判定 (防止压在盖子上)
    # 夹爪中心必须在把手目标点附近，允许上下微小浮动，但不能太高
    # 假设 height_offset 是把手几何中心，那么夹爪不应该比它高太多
    z_diff = gripper_center[:, 2] - lid_pos_w[:, 2]
    is_z_valid = torch.abs(z_diff) < 0.03 # 3cm 容差
    
    # 6. 条件三：闭合判定 (正在尝试闭合)
    joint_ids, _ = robot.find_joints(joint_names)
    # 获取关节位置 (绝对值)
    joint_pos = torch.abs(robot.data.joint_pos[:, joint_ids])
    
    # === 修改闭合判定 ===
    
    # 1. 关节位置判定 (Effort Check)
    # 只要用力了就行，去掉上限判定，防止夹太紧被误判
    # 假设 0.05 是你的新 target_close_pos
    is_effort_closing = torch.all(joint_pos > 0.005, dim=1)
    
    # 2. 几何排斥判定 (Geometry Check)
    # 如果真的夹住了东西，左右指尖的距离应该被把手撑开，不会太小
    # 获取指尖位置
    dist_fingertips = torch.norm(pos_L - pos_R, dim=-1)
    
    # 假设把手厚度是 1cm (0.01m)
    # 如果指尖距离 < 0.005，说明两个手指已经碰到一起了（空夹），没夹住东西
    # 如果指尖距离 >= 0.005，说明中间有东西挡着
    is_not_empty = dist_fingertips > 0.005
    
    # 综合判定：
    # 1. 在把手附近 (is_near)
    # 2. 高度合适 (is_z_valid)
    # 3. 在用力闭合 (is_effort_closing)
    # 4. 指尖没贴死 (is_not_empty) -> 这就证明夹住东西了！
    return is_near & is_z_valid & is_effort_closing & is_not_empty


def gripper_close_when_near(
    env: ManagerBasedRLEnv,
    lid_cfg: SceneEntityCfg = SceneEntityCfg("lid"),
    robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
    left_gripper_name: str = "left_hand__l",
    right_gripper_name: str = "left_hand_r",
    joint_names: list = ["left_hand_joint_left", "left_hand_joint_right"], 
    target_close_pos: float = 0.03,
    height_offset: float = 0.02,
    grasp_radius: float = 0.05
) -> torch.Tensor:
    
    # 1. 判定基础抓取条件是否满足
    is_grasped = _is_grasped(
        env, lid_cfg, robot_cfg, left_gripper_name, right_gripper_name,
        joint_names, height_offset, grasp_radius, target_close_pos
    )
    
    # 2. 计算夹紧程度 (Clamping Intensity)
    robot = env.scene[robot_cfg.name]
    joint_ids, _ = robot.find_joints(joint_names)
    # 获取当前关节位置绝对值 (num_envs, num_joints)
    current_joint_pos = torch.abs(robot.data.joint_pos[:, joint_ids])
    # 计算两个指关节的平均闭合深度
    mean_joint_pos = torch.mean(current_joint_pos, dim=1)
    
    # 计算奖励系数：当前位置 / 目标闭合位置
    # 这样当关节越接近 0.03 时，奖励越接近 1.0
    # 强制让 Agent 产生“压满”动作的冲动
    clamping_factor = torch.clamp(mean_joint_pos / target_close_pos, max=1.0)
    
    # 3. 只有在满足抓取判定时，才发放带有权重的夹紧奖励
    # 如果没对准或者没夹到，奖励依然是 0
    return torch.where(is_grasped, clamping_factor, 0.0)

# def gripper_close_when_near(
#     env: ManagerBasedRLEnv,
#     lid_cfg: SceneEntityCfg = SceneEntityCfg("lid"),
#     robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
#     # 注意：这里我们需要具体的指尖 body 名字来做几何判定
#     left_gripper_name: str = "left_hand__l", # 请确认你的USD里左指尖body名
#     right_gripper_name: str = "left_hand_r", # 请确认你的USD里右指尖body名
#     joint_names: list = ["left_hand_joint_left", "left_hand_joint_right"], 
#     target_close_pos: float = 0.03,
#     height_offset: float = 0.03,
#     grasp_radius: float = 0.02 # 球面半径 2cm
# ) -> torch.Tensor:
    
#     # 1. 获取位置
#     lid = env.scene[lid_cfg.name]
#     robot = env.scene[robot_cfg.name]
#     lid_pos_w = lid.data.root_pos_w[:, :3].clone()
#     lid_pos_w[:, 2] += height_offset # 把手中心
    
#     # 获取左右指尖位置
#     l_ids, _ = robot.find_bodies([left_gripper_name])
#     r_ids, _ = robot.find_bodies([right_gripper_name])
#     pos_L = robot.data.body_pos_w[:, l_ids[0], :3]
#     pos_R = robot.data.body_pos_w[:, r_ids[0], :3]
    
#     # 计算夹爪中心
#     pos_center = (pos_L + pos_R) / 2.0
    
#     # 2. 距离判定 (Is Inside Sphere?)
#     dist_center = torch.norm(pos_center - lid_pos_w, dim=-1)
#     is_in_sphere = (dist_center < grasp_radius).float()
    
#     # 3. "在中间"判定 (Is Between?)
#     # 投影逻辑：物体投影在 LR 连线上，且位于 L 和 R 之间
#     vec_LR = pos_R - pos_L # 左指指向右指
#     vec_LO = lid_pos_w - pos_L # 左指指向物体
    
#     # 计算投影比例 t
#     # P_proj = P_L + t * (P_R - P_L)
#     # t = (vec_LO . vec_LR) / |vec_LR|^2
#     # 如果 0 < t < 1，说明投影在两个指尖之间
    
#     len_sq = torch.sum(vec_LR * vec_LR, dim=-1) + 1e-6
#     dot = torch.sum(vec_LO * vec_LR, dim=-1)
#     t = dot / len_sq
    
#     is_between = (t > 0.0) & (t < 1.0)
#     is_between = is_between.float()
    
#     # 4. 闭合判定
#     joint_ids, _ = robot.find_joints(joint_names)
#     # 注意：如果是 Binary Action，可能很难拿到连续的 0~0.03 控制值，如果是仿真状态则没问题
#     current_joint_pos = torch.abs(robot.data.joint_pos[:, joint_ids])
#     close_error = torch.mean(torch.abs(current_joint_pos - target_close_pos), dim=1)
#     # 只有当非常接近闭合目标时才给分
#     is_closing = (close_error < 0.005).float() # 允许 5mm 误差
    
#     # 5. 最终奖励
#     # 只有三者同时满足才给 1.0 分
#     reward = is_in_sphere * is_between * is_closing
    
#     return torch.nan_to_num(reward, nan=0.0)


# def lid_is_lifted(
#     env:ManagerBasedRLEnv,
#     minimal_height:float,
#     lid_cfg:SceneEntityCfg = SceneEntityCfg("lid")
# ) -> torch.Tensor:

#     lid = env.scene[lid_cfg.name]
#     lid_height = lid.data.root_pos_w[:, 2]
#     reward=torch.where(lid_height > minimal_height, 1.0, 0.0)
#     reward=torch.nan_to_num(reward, nan=0.0)
    
#     return reward


def lid_is_lifted(
    env: ManagerBasedRLEnv,
    minimal_height: float = 0.05, # 相对提升阈值
    lid_cfg: SceneEntityCfg = SceneEntityCfg("lid"),
    robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
    left_gripper_name: str = "left_hand__l",
    right_gripper_name: str = "left_hand_r",
    joint_names: list = ["left_hand_joint_left", "left_hand_joint_right"], 
    height_offset: float = 0.03,
    grasp_radius: float = 0.1,
    target_close_pos: float = 0.03,
) -> torch.Tensor:
    
    is_grasped = _is_grasped(
        env, lid_cfg, robot_cfg, left_gripper_name, right_gripper_name,
        joint_names, height_offset, grasp_radius, target_close_pos
    )

    lid = env.scene[lid_cfg.name]
    
    # 1. 获取高度
    current_height = lid.data.root_pos_w[:, 2]
    # 自动获取初始高度
    initial_height = lid.data.default_root_state[:, 2]
    
    # 2. 计算提升量
    lift_height = torch.clamp(current_height - initial_height, min=0.0)
    
    # 3. 速度检查 (防撞飞)
    # 如果速度 > 1.0 m/s，视为被撞飞，不给分
    lid_vel = torch.norm(lid.data.root_lin_vel_w, dim=1)
    is_stable = lid_vel < 1.0
    
    # 4. 计算奖励
    # 基础分：高度越高分越高
    shaping_reward = lift_height * 2.0
    # 成功分：超过阈值
    success_bonus = torch.where(lift_height > minimal_height, 1.0, 0.0)
    
    # 组合
    # 必须 is_grasped AND is_stable
    reward = torch.where(is_stable & is_grasped, shaping_reward + success_bonus, torch.zeros_like(lift_height))
    
    return torch.nan_to_num(reward, nan=0.0)



def lid_lift_success_reward(
    env: ManagerBasedRLEnv,
    lid_cfg: SceneEntityCfg = SceneEntityCfg("lid"),
    robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
    left_gripper_name: str = "left_hand_body",
    joint_names: list = ["left_hand_joint_left", "left_hand_joint_right"],
    # 关键参数
    settled_height: float = 0.75,  # 盖子落在超声仪上的稳定高度 (需根据仿真实际情况微调)
    target_lift_height: float = 0.05, # 目标提升高度 (5cm)
    grasp_dist_threshold: float = 0.05, # 抓取判定距离
    closed_pos: float = 0.03 # 夹爪闭合阈值
) -> torch.Tensor:
    """
    提升奖励：只有在【夹稳】的前提下，将盖子【相对】提升达到目标高度才给高分。
    """
    # 1. 获取数据
    lid = env.scene[lid_cfg.name]
    robot = env.scene[robot_cfg.name]
    
    # 盖子实时高度
    lid_z = lid.data.root_pos_w[:, 2]
    lid_pos = lid.data.root_pos_w[:, :3]
    
    # 夹爪位置
    body_ids, _ = robot.find_bodies([left_gripper_name])
    gripper_pos = robot.data.body_pos_w[:, body_ids[0], :3]
    
    # 2. 【条件一：是否夹稳 (Is Grasped?)】
    # 距离检查
    distance = torch.norm(gripper_pos - lid_pos, dim=-1)
    is_close = distance < grasp_dist_threshold
    
    # 闭合检查
    joint_ids, _ = robot.find_joints(joint_names)
    joint_pos_abs = torch.abs(robot.data.joint_pos[:, joint_ids])
    # 假设 > 80% 的闭合度算抓紧了
    is_closed = torch.all(joint_pos_abs > (closed_pos * 0.8), dim=-1)
    
    is_grasped = is_close & is_closed
    
    # 3. 【条件二：计算相对提升 (Relative Lift)】
    # 当前高度 - 初始稳定高度
    current_lift = lid_z - settled_height
    
    # 4. 计算奖励
    lift_progress = torch.clamp(current_lift / target_lift_height, min=0.0, max=1.0)
    
    # 基础提升分 (Shaping Reward)
    lift_reward = lift_progress
    
    # 成功大奖 (Success Bonus)
    # 如果提升超过目标高度 (比如 > 95% 的 5cm)，给额外大奖
    success_bonus = torch.where(current_lift >= target_lift_height, 2.0, 0.0)
    
    # 组合：只有在 is_grasped 为 True 时才发放提升奖励
    total_reward = torch.where(is_grasped, lift_reward + success_bonus, torch.zeros_like(lift_reward))
    
    # 5. 安全输出
    return torch.nan_to_num(total_reward, nan=0.0)

def lid_lift_progress_reward(
    env: ManagerBasedRLEnv,
    lid_cfg: SceneEntityCfg = SceneEntityCfg("lid"),
    robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
    left_gripper_name: str = "left_hand__l",
    right_gripper_name: str = "left_hand_r",
    joint_names: list = ["left_hand_joint_left", "left_hand_joint_right"], 
    initial_height: float = 0.8, 
    target_height_lift: float = 0.15,
    height_offset: float = 0.02, # 必须与抓取奖励保持一致
    grasp_radius: float = 0.1,   # 提升时可以稍微放宽一点半径，防止抖动丢分
    target_close_pos: float = 0.03,
    std: float = 0.05 # 标准差
) -> torch.Tensor:

    # 1. 判定是否夹住
    is_grasped = _is_grasped(
        env, lid_cfg, robot_cfg, left_gripper_name, right_gripper_name,
        joint_names, height_offset, grasp_radius, target_close_pos
    )
    
    # 2. 计算高度
    lid = env.scene[lid_cfg.name]
    current_height = lid.data.root_pos_w[:, 2]
    lift_height = torch.clamp(current_height - initial_height, min=0.0, max=target_height_lift)
    # print(current_height)
    # print(lift_height)
    # 3. 计算奖励
    # 只有在 is_grasped 为 True 时，才发放高度奖励
    # 这样彻底杜绝了 "砸飞/撞飞" 拿分的情况
    progress = torch.tanh(lift_height / std)
    
    reward = torch.where(is_grasped, progress, 0.0)
    
    return reward

# def lid_grasped_bonus(
#     env: ManagerBasedRLEnv,
#     lid_cfg: SceneEntityCfg,
#     robot_cfg: SceneEntityCfg,
#     left_gripper_name: str,
#     joint_names: list,
#     distance_threshold: float = 0.05,
#     closed_pos: float = 0.03
# ) -> torch.Tensor:
#     # 1. 计算距离
#     lid = env.scene[lid_cfg.name]
#     robot = env.scene[robot_cfg.name]
#     body_ids, _ = robot.find_bodies([left_gripper_name])
#     gripper_pos = robot.data.body_pos_w[:, body_ids[0], :3]
#     lid_pos = lid.data.root_pos_w[:, :3]
#     distance = torch.norm(gripper_pos - lid_pos, dim=-1)
    
#     # 2. 检查距离是否达标
#     is_close = distance < distance_threshold
    
#     # 3. 检查夹爪是否闭合
#     joint_ids, _ = robot.find_joints(joint_names)
#     joint_pos = torch.abs(robot.data.joint_pos[:, joint_ids])
#     # 检查是否接近闭合 (比如 > 0.02)
#     is_closed = torch.all(joint_pos > (closed_pos * 0.8), dim=-1)
    
#     # 4. 给分
#     # 只有同时满足才给 1.0
#     reward = torch.where(is_close & is_closed, 1.0, 0.0)
    
#     return torch.nan_to_num(reward)