98 lines
4.3 KiB
Python
98 lines
4.3 KiB
Python
from __future__ import annotations
|
||
import torch
|
||
from typing import TYPE_CHECKING
|
||
from isaaclab.assets import Articulation, RigidObject
|
||
from isaaclab.managers import SceneEntityCfg
|
||
|
||
if TYPE_CHECKING:
|
||
from isaaclab.envs import ManagerBasedRLEnv
|
||
|
||
def lid_dropped(env: ManagerBasedRLEnv,
|
||
minimum_height: float = -0.05,
|
||
lid_cfg: SceneEntityCfg = SceneEntityCfg("DexCube")) -> torch.Tensor:
|
||
lid: RigidObject = env.scene[lid_cfg.name]
|
||
return lid.data.root_pos_w[:, 2] < minimum_height
|
||
|
||
def lid_successfully_grasped(env: ManagerBasedRLEnv,
|
||
distance_threshold: float = 0.03,
|
||
height_threshold: float = 0.2,
|
||
lid_cfg: SceneEntityCfg = SceneEntityCfg("DexCube"),
|
||
robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
|
||
gripper_body_name: str = "left_hand_body") -> torch.Tensor:
|
||
lid: RigidObject = env.scene[lid_cfg.name]
|
||
robot: Articulation = env.scene[robot_cfg.name]
|
||
body_idx = robot.find_bodies(gripper_body_name)[0]
|
||
gripper_pos_w = robot.data.body_pos_w[:, body_idx, :]
|
||
distance = torch.norm(lid.data.root_pos_w[:, :3] - gripper_pos_w, dim=1)
|
||
is_close = distance < distance_threshold
|
||
is_lifted = lid.data.root_pos_w[:, 2] > height_threshold
|
||
return is_close & is_lifted
|
||
|
||
def gripper_at_lid_side(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",
|
||
side_distance: float = 0.05,
|
||
side_tolerance: float = 0.01,
|
||
alignment_tolerance: float = 0.02,
|
||
height_offset: float = 0.1,
|
||
height_tolerance: float = 0.02) -> torch.Tensor:
|
||
"""Terminate when gripper center is positioned on the side of the lid at specified height.
|
||
|
||
坐标系说明：
|
||
- X 方向：两个夹爪朝中心合并的方向
|
||
- Y 方向：夹爪间空隙的方向，和 lid 的把手方向一致
|
||
- Z 方向：高度方向
|
||
"""
|
||
lid: RigidObject = env.scene[lid_cfg.name]
|
||
robot: Articulation = env.scene[robot_cfg.name]
|
||
|
||
lid_pos_w = lid.data.root_pos_w[:, :3]
|
||
|
||
# 获取两个夹爪的位置
|
||
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], :]
|
||
right_gripper_pos = robot.data.body_pos_w[:, right_body_ids[0], :]
|
||
|
||
# 计算夹爪中心位置
|
||
gripper_center = (left_gripper_pos + right_gripper_pos) / 2.0
|
||
rel_pos = gripper_center - lid_pos_w
|
||
|
||
# Y 方向：应该在 lid 的 Y 方向两侧（距离 side_distance）
|
||
y_dist = torch.abs(rel_pos[:, 1])
|
||
y_ok = (y_dist >= side_distance - side_tolerance) & (y_dist <= side_distance + side_tolerance)
|
||
|
||
# X 方向：应该对齐（接近 0）
|
||
x_dist = torch.abs(rel_pos[:, 0])
|
||
x_ok = x_dist < alignment_tolerance
|
||
|
||
# Z 方向：应该在 lid 上方 height_offset 处
|
||
z_error = torch.abs(rel_pos[:, 2] - height_offset)
|
||
z_ok = z_error < height_tolerance
|
||
|
||
# 所有条件都要满足
|
||
return x_ok & y_ok & z_ok
|
||
|
||
|
||
def base_height_failure(env: ManagerBasedRLEnv,
|
||
minimum_height: float | None = None,
|
||
maximum_height: float | None = None,
|
||
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
|
||
"""Terminate when the robot's base height is outside the specified range."""
|
||
# extract the used quantities (to enable type-hinting)
|
||
asset: Articulation = env.scene[asset_cfg.name]
|
||
root_pos_z = asset.data.root_pos_w[:, 2]
|
||
|
||
# check if height is outside the range
|
||
out_of_bounds = torch.zeros_like(root_pos_z, dtype=torch.bool)
|
||
if minimum_height is not None:
|
||
out_of_bounds |= root_pos_z < minimum_height
|
||
if maximum_height is not None:
|
||
out_of_bounds |= root_pos_z > maximum_height
|
||
|
||
return out_of_bounds
|
||
|