tangger/issacdataengine
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

init commit

Browse Source
This commit is contained in:
zyhe
2026-03-16 11:44:10 +00:00
commit 94384a93c9
552 changed files with 363038 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

448
workflows/simbox/core/skills/place.py Normal file
View File

@@ -0,0 +1,448 @@
from copy import deepcopy
import numpy as np
from core.skills.base_skill import BaseSkill, register_skill
from core.utils.box import Box, get_bbox_center_and_corners
from core.utils.constants import CUROBO_BATCH_SIZE
from core.utils.iou import IoU
from core.utils.plan_utils import (
select_index_by_priority_dual,
select_index_by_priority_single,
)
from core.utils.transformation_utils import create_pose_matrices, poses_from_tf_matrices
from core.utils.usd_geom_utils import compute_bbox
from omegaconf import DictConfig
from omni.isaac.core.controllers import BaseController
from omni.isaac.core.robots.robot import Robot
from omni.isaac.core.tasks import BaseTask
from omni.isaac.core.utils.prims import get_prim_at_path
from omni.isaac.core.utils.transformations import (
get_relative_transform,
pose_from_tf_matrix,
tf_matrix_from_pose,
)
from scipy.spatial.transform import Rotation as R
# pylint: disable=consider-using-generator,too-many-public-methods,unused-argument
@register_skill
class Place(BaseSkill):
def __init__(self, robot: Robot, controller: BaseController, task: BaseTask, cfg: DictConfig, *args, **kwargs):
super().__init__()
self.robot = robot
self.controller = controller
self.task = task
self.name = cfg["name"]
self.pick_obj = task._task_objects[cfg["objects"][0]]
self.place_obj = task._task_objects[cfg["objects"][1]]
self.place_align_axis = cfg.get("place_align_axis", None)
self.pick_align_axis = cfg.get("pick_align_axis", None)
self.constraint_gripper_x = cfg.get("constraint_gripper_x", False)
self.place_part_prim_path = cfg.get("place_part_prim_path", None)
if self.place_part_prim_path:
self.place_prim_path = f"{self.place_obj.prim_path}/{self.place_part_prim_path}"
else:
self.place_prim_path = self.place_obj.prim_path
self.manip_list = []
# if "Franka" in self.controller.robot_file or "Franka" in self.robot.cfg["name"]:
# self.robot_ee_path = self.robot.ee_path
# self.robot_base_path = self.robot.base_path
# else:
# self.robot_ee_path = self.controller.robot_ee_path
# self.robot_base_path = self.controller.robot_base_path
self.robot_ee_path = self.controller.robot_ee_path
self.robot_base_path = self.controller.robot_base_path
self.skill_cfg = cfg
self.align_pick_obj_axis = self.skill_cfg.get("align_pick_obj_axis", None)
self.align_place_obj_axis = self.skill_cfg.get("align_place_obj_axis", None)
self.align_plane_x_axis = self.skill_cfg.get("align_plane_x_axis", None)
self.align_plane_y_axis = self.skill_cfg.get("align_plane_y_axis", None)
self.align_obj_tol = self.skill_cfg.get("align_obj_tol", None)
def simple_generate_manip_cmds(self):
manip_list = []
p_base_ee_cur, q_base_ee_cur = self.controller.get_ee_pose()
cmd = (p_base_ee_cur, q_base_ee_cur, "update_pose_cost_metric", {"hold_vec_weight": None})
manip_list.append(cmd)
if self.skill_cfg.get("ignore_substring", []):
ignore_substring = deepcopy(self.controller.ignore_substring + self.skill_cfg.get("ignore_substring", []))
cmd = (
p_base_ee_cur,
q_base_ee_cur,
"update_specific",
{"ignore_substring": ignore_substring, "reference_prim_path": self.controller.reference_prim_path},
)
manip_list.append(cmd)
if self.skill_cfg.get("pre_place_hold_vec_weight", None) is not None:
cmd = (
p_base_ee_cur,
q_base_ee_cur,
"update_pose_cost_metric",
{"hold_vec_weight": self.skill_cfg.get("pre_place_hold_vec_weight", None)},
)
manip_list.append(cmd)
result = self.sample_gripper_place_traj()
cmd = (result[0][0], result[0][1], "close_gripper", {})
manip_list.append(cmd)
if self.skill_cfg.get("post_place_hold_vec_weight", None) is not None:
cmd = (
result[0][0],
result[0][1],
"update_pose_cost_metric",
{"hold_vec_weight": self.skill_cfg.get("post_place_hold_vec_weight", None)},
)
manip_list.append(cmd)
p_base_ee_place, q_base_ee_place = result[1][0], result[1][1]
cmd = (p_base_ee_place, q_base_ee_place, "close_gripper", {})
manip_list.append(cmd)
manip_list.extend([cmd] * self.skill_cfg.get("hesitate_steps", 0))
cmd = (p_base_ee_place, q_base_ee_place, "open_gripper", {})
manip_list.extend([cmd] * self.skill_cfg.get("gripper_change_steps", 10))
cmd = (p_base_ee_place, q_base_ee_place, "detach_obj", {})
manip_list.append(cmd)
ignore_substring = deepcopy(self.controller.ignore_substring)
cmd = (
p_base_ee_place,
q_base_ee_place,
"update_specific",
{"ignore_substring": ignore_substring, "reference_prim_path": self.controller.reference_prim_path},
)
manip_list.append(cmd)
# Postplace
if self.skill_cfg.get("post_place_vector", None):
ignore_substring = deepcopy(self.controller.ignore_substring)
ignore_substring.append(self.pick_obj.name)
cmd = (
p_base_ee_place,
q_base_ee_place,
"update_specific",
{"ignore_substring": ignore_substring, "reference_prim_path": self.controller.reference_prim_path},
)
manip_list.append(cmd)
cmd = (result[2][0], result[2][1], "open_gripper", {})
manip_list.append(cmd)
self.manip_list = manip_list
self.place_ee_trans = p_base_ee_place
def sample_gripper_place_traj(self):
self.T_world_obj = tf_matrix_from_pose(*self.pick_obj.get_local_pose())
self.T_world_ee = get_relative_transform(
get_prim_at_path(self.robot_ee_path), get_prim_at_path(self.task.root_prim_path)
)
self.T_base_world = get_relative_transform(
get_prim_at_path(self.task.root_prim_path), get_prim_at_path(self.robot_base_path)
)
self.T_obj_ee = np.linalg.inv(self.T_world_obj) @ self.T_world_ee
ee2o_distance = np.linalg.norm(self.T_obj_ee[0:3, 3])
self.T_world_container = tf_matrix_from_pose(*self.place_obj.get_local_pose())
bbox_place_obj = compute_bbox(get_prim_at_path(self.place_prim_path))
b_min, b_max = bbox_place_obj.min, bbox_place_obj.max
def get_range(key, default):
r = self.skill_cfg.get(key, default)
return np.random.uniform(r[0], r[1], size=CUROBO_BATCH_SIZE)
x_ratios = get_range("x_ratio_range", (0.4, 0.6))
y_ratios = get_range("y_ratio_range", (0.4, 0.6))
z_ratios = get_range("z_ratio_range", (0.4, 0.6))
direction = self.skill_cfg.get("place_direction", "vertical")
pos_constraint = self.skill_cfg.get("position_constraint", "gripper")
pre_z_off = self.skill_cfg.get("pre_place_z_offset", 0.2)
plt_z_off = self.skill_cfg.get("place_z_offset", 0.1)
if direction == "vertical":
x = b_min[0] + x_ratios * (b_max[0] - b_min[0])
y = b_min[1] + y_ratios * (b_max[1] - b_min[1])
pre_place_pos_w = np.stack([x, y, np.full(CUROBO_BATCH_SIZE, b_max[2] + pre_z_off)], axis=-1)
place_pos_w = np.stack([x, y, np.full(CUROBO_BATCH_SIZE, b_max[2] + plt_z_off)], axis=-1)
elif direction == "horizontal":
align_axis = self.T_world_container[:3, :3] @ np.array(self.skill_cfg["align_place_obj_axis"])
offset_axis = self.T_world_container[:3, :3] @ np.array(self.skill_cfg["offset_place_obj_axis"])
tmp_pos_w = b_min + np.stack([x_ratios, y_ratios, z_ratios], axis=-1) * (b_max - b_min)
if pos_constraint == "object":
pre_align = self.skill_cfg.get("pre_place_align", 0.2)
pre_offset = self.skill_cfg.get("pre_place_offset", 0.2)
plt_align = self.skill_cfg.get("place_align", 0.1)
plt_offset = self.skill_cfg.get("place_offset", 0.1)
pre_place_pos_w = tmp_pos_w + align_axis * pre_align + offset_axis * pre_offset
place_pos_w = tmp_pos_w + align_axis * plt_align + offset_axis * plt_offset
else:
pre_place_pos_w = tmp_pos_w - align_axis * (pre_z_off + ee2o_distance)
place_pos_w = tmp_pos_w - align_axis * (plt_z_off + ee2o_distance)
else:
raise NotImplementedError
pre_place_pos_b = (self.T_base_world[:3, :3] @ pre_place_pos_w.T).T + self.T_base_world[:3, 3]
place_pos_b = (self.T_base_world[:3, :3] @ place_pos_w.T).T + self.T_base_world[:3, 3]
R_tgts = self.generate_constrained_rotation_batch()
if pos_constraint == "object":
R_ee_obj = np.linalg.inv(self.T_obj_ee)[:3, :3]
R_base_obj_tgts = R_tgts @ R_ee_obj
T_base_obj_pre_tgt = create_pose_matrices(pre_place_pos_b, R_base_obj_tgts)
T_base_obj_plt_tgt = create_pose_matrices(place_pos_b, R_base_obj_tgts)
T_base_ee_preplaces = T_base_obj_pre_tgt @ self.T_obj_ee
T_base_ee_places = T_base_obj_plt_tgt @ self.T_obj_ee
elif pos_constraint == "gripper":
T_base_ee_preplaces = create_pose_matrices(pre_place_pos_b, R_tgts)
T_base_ee_places = create_pose_matrices(place_pos_b, R_tgts)
else:
raise NotImplementedError
self.controller.update_specific(
ignore_substring=self.controller.ignore_substring + self.skill_cfg.get("ignore_substring", []),
reference_prim_path=self.controller.reference_prim_path,
)
p_base_ee_preplaces, q_base_ee_preplaces = poses_from_tf_matrices(T_base_ee_preplaces)
p_base_ee_places, q_base_ee_places = poses_from_tf_matrices(T_base_ee_places)
if self.controller.use_batch:
# Check if the input arrays are exactly the same
if np.array_equal(p_base_ee_preplaces, p_base_ee_places) and np.array_equal(
q_base_ee_preplaces, q_base_ee_places
):
# Inputs are identical, compute only once to avoid redundant computation
result = self.controller.test_batch_forward(p_base_ee_places, q_base_ee_places)
index = select_index_by_priority_single(result)
else:
# Inputs are different, compute separately
pre_result = self.controller.test_batch_forward(p_base_ee_preplaces, q_base_ee_preplaces)
result = self.controller.test_batch_forward(p_base_ee_places, q_base_ee_places)
index = select_index_by_priority_dual(pre_result, result)
else:
for index in range(T_base_ee_places.shape[0]):
p_base_ee_pregrasp, q_base_ee_pregrasp = p_base_ee_preplaces[index], q_base_ee_preplaces[index]
p_base_ee_grasp, q_base_ee_grasp = p_base_ee_places[index], q_base_ee_places[index]
test_mode = self.skill_cfg.get("test_mode", "forward")
if test_mode == "forward":
result_pre = self.controller.test_single_forward(p_base_ee_pregrasp, q_base_ee_pregrasp)
elif test_mode == "ik":
result_pre = self.controller.test_single_ik(p_base_ee_pregrasp, q_base_ee_pregrasp)
else:
raise NotImplementedError
if self.skill_cfg.get("pre_grasp_offset", 0.1) > 0:
if test_mode == "forward":
result = self.controller.test_single_forward(p_base_ee_grasp, q_base_ee_grasp)
elif test_mode == "ik":
result = self.controller.test_single_ik(p_base_ee_grasp, q_base_ee_grasp)
else:
raise NotImplementedError
if result == 1 and result_pre == 1:
print("place plan success")
break
else:
if result_pre == 1:
print("place plan success")
break
res_pre = list(pose_from_tf_matrix(T_base_ee_preplaces[index]))
res_plt = list(pose_from_tf_matrix(T_base_ee_places[index]))
if self.skill_cfg.get("post_place_vector", None):
post_vec = np.array(self.skill_cfg["post_place_vector"])
T_base_ee_postplace = deepcopy(T_base_ee_places[index])
T_base_ee_postplace[:3, 3] = T_base_ee_places[index][:3, :3] @ post_vec + T_base_ee_places[index][:3, 3]
res_post = list(pose_from_tf_matrix(T_base_ee_postplace))
return [res_pre, res_plt, res_post]
return [res_pre, res_plt]
def generate_constrained_rotation_batch(self, batch_size=3000):
filter_conditions = {
"x": {
"forward": (0, 0, 1), # (row, col, direction)
"backward": (0, 0, -1),
"leftward": (1, 0, 1),
"rightward": (1, 0, -1),
"upward": (2, 0, 1),
"downward": (2, 0, -1),
},
"y": {
"forward": (0, 1, 1),
"backward": (0, 1, -1),
"leftward": (1, 1, 1),
"rightward": (1, 1, -1),
"upward": (2, 1, 1),
"downward": (2, 1, -1),
},
"z": {
"forward": (0, 2, 1),
"backward": (0, 2, -1),
"leftward": (1, 2, 1),
"rightward": (1, 2, -1),
"upward": (2, 2, 1),
"downward": (2, 2, -1),
},
}
rot_mats = R.random(batch_size).as_matrix()
valid_mask = np.ones(batch_size, dtype=bool)
for axis in ["x", "y", "z"]:
filter_list = self.skill_cfg.get(f"filter_{axis}_dir", None)
if filter_list is not None:
direction = filter_list[0]
row, col, sign = filter_conditions[axis][direction]
elements = rot_mats[:, row, col]
if len(filter_list) == 2:
value = filter_list[1]
cos_val = np.cos(np.deg2rad(value))
if sign > 0:
valid_mask &= elements >= cos_val
else:
valid_mask &= elements <= cos_val
elif len(filter_list) == 3:
value1, value2 = filter_list[1:]
cos_val1 = np.cos(np.deg2rad(value1))
cos_val2 = np.cos(np.deg2rad(value2))
if sign > 0:
valid_mask &= (elements >= cos_val1) & (elements <= cos_val2)
else:
valid_mask &= (elements <= cos_val1) & (elements >= cos_val2)
if (
self.align_pick_obj_axis is not None
and self.align_place_obj_axis is not None
and self.align_obj_tol is not None
):
align_pick_obj_axis = np.array(self.align_pick_obj_axis)
align_place_obj_axis = np.array(self.align_place_obj_axis)
R_ee_obj = np.linalg.inv(self.T_obj_ee)[:3, :3]
R_base_container_tgt = (self.T_base_world @ self.T_world_container)[:3, :3]
R_base_obj_tgts = rot_mats @ R_ee_obj # (N, 3, 3)
pick_vecs_tgt = np.einsum("ijk,k->ij", R_base_obj_tgts, align_pick_obj_axis) # (N, 3)
place_vec_tgt = R_base_container_tgt @ align_place_obj_axis
dot_products = np.dot(pick_vecs_tgt, place_vec_tgt)
norms = np.linalg.norm(pick_vecs_tgt, axis=1) * np.linalg.norm(place_vec_tgt)
radians = np.arccos(np.clip(dot_products / norms, -1.0, 1.0))
valid_mask &= radians < np.deg2rad(self.align_obj_tol)
valid_rot_mats = rot_mats[valid_mask]
print("length of valid place rots :", len(valid_rot_mats))
if len(valid_rot_mats) == 0:
print("Warning: No matrix satisfies constraints")
return rot_mats[:CUROBO_BATCH_SIZE]
else:
indices = np.random.choice(len(valid_rot_mats), CUROBO_BATCH_SIZE)
return valid_rot_mats[indices]
def is_feasible(self, th=5):
return self.controller.num_plan_failed <= th
def is_subtask_done(self, t_eps=1e-3, o_eps=5e-3):
assert len(self.manip_list) != 0
p_base_ee_cur, q_base_ee_cur = self.controller.get_ee_pose()
p_base_ee, q_base_ee, *_ = self.manip_list[0]
diff_trans = np.linalg.norm(p_base_ee_cur - p_base_ee)
diff_ori = 2 * np.arccos(min(abs(np.dot(q_base_ee_cur, q_base_ee)), 1.0))
pose_flag = np.logical_and(
diff_trans < t_eps,
diff_ori < o_eps,
)
self.plan_flag = self.controller.num_last_cmd > 10
return np.logical_or(pose_flag, self.plan_flag)
def is_done(self):
if len(self.manip_list) == 0:
return True
if self.is_subtask_done(t_eps=self.skill_cfg.get("t_eps", 1e-3), o_eps=self.skill_cfg.get("o_eps", 5e-3)):
self.manip_list.pop(0)
# if self.is_success():
# self.manip_list.clear()
# print("Rotate Done")
return len(self.manip_list) == 0
def is_success(self, th=0.0):
if self.skill_cfg.get("success_mode", "3diou") == "3diou":
bbox_pick_obj = compute_bbox(self.pick_obj.prim)
bbox_place_obj = compute_bbox(get_prim_at_path(self.place_prim_path))
iou = IoU(
Box(get_bbox_center_and_corners(bbox_pick_obj)), Box(get_bbox_center_and_corners(bbox_place_obj))
).iou()
print(iou)
return iou > th
elif self.skill_cfg.get("success_mode", "3diou") == "height":
T_o2r = get_relative_transform(
get_prim_at_path(self.pick_obj.prim_path), get_prim_at_path(self.robot_base_path)
)
T_o2r_trans = T_o2r[:3, 3]
return T_o2r_trans[2] < self.place_ee_trans[2] - 0.4
elif self.skill_cfg.get("success_mode", "3diou") == "xybbox":
bbox_place_obj = compute_bbox(get_prim_at_path(self.place_prim_path))
pick_x, pick_y = self.pick_obj.get_local_pose()[0][:2]
place_xy_min = bbox_place_obj.min[:2]
place_xy_max = bbox_place_obj.max[:2]
return ((place_xy_min[0] + 0.015) < pick_x < (place_xy_max[0] - 0.015)) and (
(place_xy_min[1] + 0.015) < pick_y < (place_xy_max[1] - 0.015)
)
elif self.skill_cfg.get("success_mode", "3diou") == "left":
bbox_place_obj = compute_bbox(get_prim_at_path(self.place_prim_path))
pick_x, pick_y = self.pick_obj.get_local_pose()[0][:2]
place_xy_min = bbox_place_obj.min[:2]
place_xy_max = bbox_place_obj.max[:2]
return pick_x < place_xy_min[0] - self.skill_cfg.get("threshold", 0.03)
elif self.skill_cfg.get("success_mode", "3diou") == "right":
bbox_place_obj = compute_bbox(get_prim_at_path(self.place_prim_path))
pick_x, pick_y = self.pick_obj.get_local_pose()[0][:2]
place_xy_min = bbox_place_obj.min[:2]
place_xy_max = bbox_place_obj.max[:2]
return pick_x > place_xy_max[0] + self.skill_cfg.get("threshold", 0.03)
elif self.skill_cfg.get("success_mode", "3diou") == "flower":
bbox_pick_obj = compute_bbox(self.pick_obj.prim)
bbox_place_obj = compute_bbox(get_prim_at_path(self.place_prim_path))
iou = IoU(
Box(get_bbox_center_and_corners(bbox_pick_obj)), Box(get_bbox_center_and_corners(bbox_place_obj))
).iou()
print("iou", iou)
th = self.skill_cfg.get("success_th", 0.0)
middle = self.pick_obj.get_local_pose()[0]
x_min, y_min, _ = bbox_place_obj.min
x_max, y_max, _ = bbox_place_obj.max
x_middle, y_middle, _ = middle[0], middle[1], middle[2]
if x_min < x_middle < x_max and y_min < y_middle < y_max:
return iou > th
else:
return False
elif self.skill_cfg.get("success_mode", "3diou") == "cup":
bbox_pick_obj = compute_bbox(self.pick_obj.prim)
bbox_place_obj = compute_bbox(get_prim_at_path(self.place_prim_path))
iou = IoU(
Box(get_bbox_center_and_corners(bbox_pick_obj)), Box(get_bbox_center_and_corners(bbox_place_obj))
).iou()
x_cup, y_cup, z_cup = bbox_pick_obj.min
x_shelf, y_shelf, z_shelf = bbox_place_obj.min
print("iou", iou)
print("x_cup, y_cup, z_cup", x_cup, y_cup, z_cup)
print("x_shelf, y_shelf, z_shelf", x_shelf, y_shelf, z_shelf)
return (z_cup > z_shelf + 0.05) and (iou > self.skill_cfg.get("success_th", 0.0))