03d9a5b909
- Fix USD metersPerUnit/upAxis for IS 5.0.0 (no longer auto-compensated) - Batch fix all Aligned_obj.usd, table, and art USD files with backups - Fix DomeLight rotation to Z-axis only (prevent tilted environment map) - Fix scene reuse across episodes (arena_file caching, task clearing, prim guard) - Add migration tools: scan_usd_metadata.py, fix_usd_metadata.py - Add migration guide: migerate/migerate.md - Add nvidia-curobo to .gitignore - Fix sort_the_rubbish config: obj_0 -> obj_1 (obj_0 does not exist) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
392 lines
18 KiB
Python
392 lines
18 KiB
Python
import os
|
|
import random
|
|
from copy import deepcopy
|
|
|
|
import numpy as np
|
|
from core.skills.base_skill import BaseSkill, register_skill
|
|
from core.utils.constants import CUROBO_BATCH_SIZE
|
|
from core.utils.plan_utils import (
|
|
select_index_by_priority_dual,
|
|
select_index_by_priority_single,
|
|
)
|
|
from core.utils.transformation_utils import poses_from_tf_matrices
|
|
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,
|
|
tf_matrix_from_pose,
|
|
)
|
|
|
|
|
|
# pylint: disable=unused-argument
|
|
@register_skill
|
|
class Pick(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.skill_cfg = cfg
|
|
object_name = self.skill_cfg["objects"][0]
|
|
self.pick_obj = task.objects[object_name]
|
|
|
|
# Get grasp annotation
|
|
usd_path = [obj["path"] for obj in task.cfg["objects"] if obj["name"] == object_name][0]
|
|
usd_path = os.path.join(self.task.asset_root, usd_path)
|
|
grasp_pose_path = usd_path.replace(
|
|
"Aligned_obj.usd", self.skill_cfg.get("npy_name", "Aligned_grasp_sparse.npy")
|
|
)
|
|
sparse_grasp_poses = np.load(grasp_pose_path)
|
|
lr_arm = "right" if "right" in self.controller.robot_file else "left"
|
|
self.T_obj_ee, self.scores = self.robot.pose_post_process_fn(
|
|
sparse_grasp_poses,
|
|
lr_arm=lr_arm,
|
|
grasp_scale=self.skill_cfg.get("grasp_scale", 1),
|
|
tcp_offset=self.skill_cfg.get("tcp_offset", self.robot.tcp_offset),
|
|
constraints=self.skill_cfg.get("constraints", None),
|
|
)
|
|
|
|
# Keyposes should be generated after previous skill is done
|
|
self.manip_list = []
|
|
self.pickcontact_view = task.pickcontact_views[robot.name][lr_arm][object_name]
|
|
self.process_valid = True
|
|
self.obj_init_trans = deepcopy(self.pick_obj.get_local_pose()[0])
|
|
final_gripper_state = self.skill_cfg.get("final_gripper_state", -1)
|
|
if final_gripper_state == 1:
|
|
self.gripper_cmd = "open_gripper"
|
|
elif final_gripper_state == -1:
|
|
self.gripper_cmd = "close_gripper"
|
|
else:
|
|
raise ValueError(f"final_gripper_state must be 1 or -1, got {final_gripper_state}")
|
|
self.fixed_orientation = self.skill_cfg.get("fixed_orientation", None)
|
|
if self.fixed_orientation is not None:
|
|
self.fixed_orientation = np.array(self.fixed_orientation)
|
|
|
|
def simple_generate_manip_cmds(self):
|
|
manip_list = []
|
|
|
|
# Update
|
|
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)
|
|
|
|
ignore_substring = deepcopy(self.controller.ignore_substring + self.skill_cfg.get("ignore_substring", []))
|
|
ignore_substring.append(self.pick_obj.name)
|
|
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)
|
|
|
|
# Pre grasp
|
|
T_base_ee_grasps = self.sample_ee_pose() # (N, 4, 4)
|
|
T_base_ee_pregrasps = deepcopy(T_base_ee_grasps)
|
|
self.controller.update_specific(
|
|
ignore_substring=ignore_substring, reference_prim_path=self.controller.reference_prim_path
|
|
)
|
|
|
|
if "r5a" in self.controller.robot_file:
|
|
T_base_ee_pregrasps[:, :3, 3] -= T_base_ee_pregrasps[:, :3, 0] * self.skill_cfg.get("pre_grasp_offset", 0.1)
|
|
else:
|
|
T_base_ee_pregrasps[:, :3, 3] -= T_base_ee_pregrasps[:, :3, 2] * self.skill_cfg.get("pre_grasp_offset", 0.1)
|
|
|
|
p_base_ee_pregrasps, q_base_ee_pregrasps = poses_from_tf_matrices(T_base_ee_pregrasps)
|
|
p_base_ee_grasps, q_base_ee_grasps = poses_from_tf_matrices(T_base_ee_grasps)
|
|
|
|
if self.controller.use_batch:
|
|
# Check if the input arrays are exactly the same
|
|
if np.array_equal(p_base_ee_pregrasps, p_base_ee_grasps) and np.array_equal(
|
|
q_base_ee_pregrasps, q_base_ee_grasps
|
|
):
|
|
# Inputs are identical, compute only once to avoid redundant computation
|
|
result = self.controller.test_batch_forward(p_base_ee_grasps, q_base_ee_grasps)
|
|
index = select_index_by_priority_single(result)
|
|
else:
|
|
# Inputs are different, compute separately
|
|
pre_result = self.controller.test_batch_forward(p_base_ee_pregrasps, q_base_ee_pregrasps)
|
|
result = self.controller.test_batch_forward(p_base_ee_grasps, q_base_ee_grasps)
|
|
index = select_index_by_priority_dual(pre_result, result)
|
|
else:
|
|
for index in range(T_base_ee_grasps.shape[0]):
|
|
p_base_ee_pregrasp, q_base_ee_pregrasp = p_base_ee_pregrasps[index], q_base_ee_pregrasps[index]
|
|
p_base_ee_grasp, q_base_ee_grasp = p_base_ee_grasps[index], q_base_ee_grasps[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("pick plan success")
|
|
break
|
|
else:
|
|
if result_pre == 1:
|
|
print("pick plan success")
|
|
break
|
|
|
|
if self.fixed_orientation is not None:
|
|
q_base_ee_pregrasps[index] = self.fixed_orientation
|
|
q_base_ee_grasps[index] = self.fixed_orientation
|
|
|
|
# Pre-grasp
|
|
cmd = (p_base_ee_pregrasps[index], q_base_ee_pregrasps[index], "open_gripper", {})
|
|
manip_list.append(cmd)
|
|
if self.skill_cfg.get("pre_grasp_hold_vec_weight", None) is not None:
|
|
cmd = (
|
|
p_base_ee_pregrasps[index],
|
|
q_base_ee_pregrasps[index],
|
|
"update_pose_cost_metric",
|
|
{"hold_vec_weight": self.skill_cfg.get("pre_grasp_hold_vec_weight", None)},
|
|
)
|
|
manip_list.append(cmd)
|
|
|
|
# Grasp
|
|
cmd = (p_base_ee_grasps[index], q_base_ee_grasps[index], "open_gripper", {})
|
|
manip_list.append(cmd)
|
|
cmd = (p_base_ee_grasps[index], q_base_ee_grasps[index], self.gripper_cmd, {})
|
|
manip_list.extend(
|
|
[cmd] * self.skill_cfg.get("gripper_change_steps", 40)
|
|
) # Default we use 40 steps to make sure the gripper is fully closed
|
|
ignore_substring = deepcopy(self.controller.ignore_substring + self.skill_cfg.get("ignore_substring", []))
|
|
cmd = (
|
|
p_base_ee_grasps[index],
|
|
q_base_ee_grasps[index],
|
|
"update_specific",
|
|
{"ignore_substring": ignore_substring, "reference_prim_path": self.controller.reference_prim_path},
|
|
)
|
|
manip_list.append(cmd)
|
|
cmd = (
|
|
p_base_ee_grasps[index],
|
|
q_base_ee_grasps[index],
|
|
"attach_obj",
|
|
{"obj_prim_path": self.pick_obj.mesh_prim_path},
|
|
)
|
|
manip_list.append(cmd)
|
|
|
|
# Post-grasp
|
|
post_grasp_offset = np.random.uniform(
|
|
self.skill_cfg.get("post_grasp_offset_min", 0.05), self.skill_cfg.get("post_grasp_offset_max", 0.05)
|
|
)
|
|
if post_grasp_offset:
|
|
p_base_ee_postgrasps = deepcopy(p_base_ee_grasps)
|
|
p_base_ee_postgrasps[index][2] += post_grasp_offset
|
|
cmd = (p_base_ee_postgrasps[index], q_base_ee_grasps[index], self.gripper_cmd, {})
|
|
manip_list.append(cmd)
|
|
|
|
# Whether return to pre-grasp
|
|
if self.skill_cfg.get("return_to_pregrasp", False):
|
|
cmd = (p_base_ee_pregrasps[index], q_base_ee_pregrasps[index], self.gripper_cmd, {})
|
|
manip_list.append(cmd)
|
|
|
|
self.manip_list = manip_list
|
|
|
|
def sample_ee_pose(self, max_length=CUROBO_BATCH_SIZE):
|
|
T_base_ee = self.get_ee_poses("armbase")
|
|
|
|
num_pose = T_base_ee.shape[0]
|
|
flags = {
|
|
"x": np.ones(num_pose, dtype=bool),
|
|
"y": np.ones(num_pose, dtype=bool),
|
|
"z": np.ones(num_pose, dtype=bool),
|
|
"direction_to_obj": np.ones(num_pose, dtype=bool),
|
|
}
|
|
filter_conditions = {
|
|
"x": {
|
|
"forward": (0, 0, 1), # (row, col, direction)
|
|
"backward": (0, 0, -1),
|
|
"upward": (2, 0, 1),
|
|
"downward": (2, 0, -1),
|
|
},
|
|
"y": {"forward": (0, 1, 1), "backward": (0, 1, -1), "downward": (2, 1, -1), "upward": (2, 1, 1)},
|
|
"z": {"forward": (0, 2, 1), "backward": (0, 2, -1), "downward": (2, 2, -1), "upward": (2, 2, 1)},
|
|
}
|
|
for axis in ["x", "y", "z"]:
|
|
filter_list = self.skill_cfg.get(f"filter_{axis}_dir", None)
|
|
if filter_list is not None:
|
|
# direction, value = filter_list
|
|
direction = filter_list[0]
|
|
row, col, sign = filter_conditions[axis][direction]
|
|
if len(filter_list) == 2:
|
|
value = filter_list[1]
|
|
cos_val = np.cos(np.deg2rad(value))
|
|
flags[axis] = T_base_ee[:, row, col] >= cos_val if sign > 0 else T_base_ee[:, row, col] <= 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:
|
|
flags[axis] = np.logical_and(
|
|
T_base_ee[:, row, col] >= cos_val1, T_base_ee[:, row, col] <= cos_val2
|
|
)
|
|
else:
|
|
flags[axis] = np.logical_and(
|
|
T_base_ee[:, row, col] <= cos_val1, T_base_ee[:, row, col] >= cos_val2
|
|
)
|
|
if self.skill_cfg.get("direction_to_obj", None) is not None:
|
|
direction_to_obj = self.skill_cfg["direction_to_obj"]
|
|
T_world_obj = tf_matrix_from_pose(*self.pick_obj.get_local_pose())
|
|
T_base_world = get_relative_transform(
|
|
get_prim_at_path(self.task.root_prim_path), get_prim_at_path(self.controller.reference_prim_path)
|
|
)
|
|
T_base_obj = T_base_world @ T_world_obj
|
|
if direction_to_obj == "right":
|
|
flags["direction_to_obj"] = T_base_ee[:, 1, 3] <= T_base_obj[1, 3]
|
|
elif direction_to_obj == "left":
|
|
flags["direction_to_obj"] = T_base_ee[:, 1, 3] > T_base_obj[1, 3]
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
combined_flag = np.logical_and.reduce(list(flags.values()))
|
|
if sum(combined_flag) == 0:
|
|
# idx_list = [i for i in range(max_length)]
|
|
idx_list = list(range(max_length))
|
|
else:
|
|
tmp_scores = self.scores[combined_flag]
|
|
tmp_idxs = np.arange(num_pose)[combined_flag]
|
|
combined = list(zip(tmp_scores, tmp_idxs))
|
|
combined.sort()
|
|
idx_list = [idx for (score, idx) in combined[:max_length]]
|
|
score_list = self.scores[idx_list]
|
|
weights = 1.0 / (score_list + 1e-8)
|
|
weights = weights / weights.sum()
|
|
|
|
sampled_idx = random.choices(idx_list, weights=weights, k=max_length)
|
|
sampled_scores = self.scores[sampled_idx]
|
|
|
|
# Sort indices by their scores (ascending)
|
|
sorted_pairs = sorted(zip(sampled_scores, sampled_idx))
|
|
idx_list = [idx for _, idx in sorted_pairs]
|
|
|
|
print(self.scores[idx_list])
|
|
# print((T_base_ee[idx_list])[:, 0, 1])
|
|
return T_base_ee[idx_list]
|
|
|
|
def get_ee_poses(self, frame: str = "world"):
|
|
# get grasp poses at specific frame
|
|
if frame not in ["world", "body", "armbase"]:
|
|
raise ValueError(
|
|
f"poses in {frame} frame is not supported: accepted values are [world, body, armbase] only"
|
|
)
|
|
|
|
if frame == "body":
|
|
return self.T_obj_ee
|
|
|
|
local_pose = self.pick_obj.get_local_pose()
|
|
world_pose = self.pick_obj.get_world_pose()
|
|
print(f"[PICK_DBG] prim_path={self.pick_obj.prim_path}", flush=True)
|
|
print(f"[PICK_DBG] local_pose trans={local_pose[0]}", flush=True)
|
|
print(f"[PICK_DBG] world_pose trans={world_pose[0]}", flush=True)
|
|
print(f"[PICK_DBG] local_pose ori={local_pose[1]}", flush=True)
|
|
print(f"[PICK_DBG] world_pose ori={world_pose[1]}", flush=True)
|
|
parent_prim = get_prim_at_path(self.pick_obj.prim_path).GetParent()
|
|
print(f"[PICK_DBG] parent_prim_path={parent_prim.GetPrimPath()}", flush=True)
|
|
grandparent_prim = parent_prim.GetParent()
|
|
print(f"[PICK_DBG] grandparent_prim_path={grandparent_prim.GetPrimPath()}", flush=True)
|
|
try:
|
|
from pxr import UsdGeom
|
|
obj_prim = get_prim_at_path(self.pick_obj.prim_path)
|
|
obj_xf = UsdGeom.Xformable(obj_prim)
|
|
xform_ops = obj_xf.GetOrderedXformOps()
|
|
print(f"[PICK_DBG] obj_xformOps_count={len(xform_ops)}", flush=True)
|
|
for op in xform_ops:
|
|
print(f"[PICK_DBG] obj_xformOp: {op.GetName()} val={op.Get()}", flush=True)
|
|
obj_l2w = obj_xf.ComputeLocalToWorldTransform(0)
|
|
print(f"[PICK_DBG] obj_USD_l2w=({obj_l2w[3][0]:.6f}, {obj_l2w[3][1]:.6f}, {obj_l2w[3][2]:.6f})", flush=True)
|
|
parent_xf = UsdGeom.Xformable(parent_prim)
|
|
parent_l2w = parent_xf.ComputeLocalToWorldTransform(0)
|
|
print(f"[PICK_DBG] parent_l2w_translate=({parent_l2w[3][0]}, {parent_l2w[3][1]}, {parent_l2w[3][2]})", flush=True)
|
|
gp_xf = UsdGeom.Xformable(grandparent_prim)
|
|
gp_l2w = gp_xf.ComputeLocalToWorldTransform(0)
|
|
print(f"[PICK_DBG] grandparent_l2w_translate=({gp_l2w[3][0]}, {gp_l2w[3][1]}, {gp_l2w[3][2]})", flush=True)
|
|
stage = parent_prim.GetStage()
|
|
mpu = UsdGeom.GetStageMetersPerUnit(stage)
|
|
print(f"[PICK_DBG] stage_metersPerUnit={mpu}", flush=True)
|
|
# Check get_local_pose source
|
|
glp_method = type(self.pick_obj).get_local_pose
|
|
print(f"[PICK_DBG] get_local_pose_from={glp_method.__qualname__}", flush=True)
|
|
except Exception as e:
|
|
import traceback
|
|
print(f"[PICK_DBG] UsdGeom error: {e}", flush=True)
|
|
traceback.print_exc()
|
|
print(f"[PICK_DBG] root_prim_path={self.task.root_prim_path}", flush=True)
|
|
print(f"[PICK_DBG] reference_prim_path={self.controller.reference_prim_path}", flush=True)
|
|
T_world_obj = tf_matrix_from_pose(*local_pose)
|
|
print(f"[PICK_DBG] T_world_obj translation={T_world_obj[:3,3]}", flush=True)
|
|
T_world_ee = T_world_obj[None] @ self.T_obj_ee
|
|
print(f"[PICK_DBG] T_world_ee[0] translation={T_world_ee[0,:3,3]}", flush=True)
|
|
|
|
if frame == "world":
|
|
return T_world_ee
|
|
|
|
if frame == "armbase": # arm base frame
|
|
T_world_base = get_relative_transform(
|
|
get_prim_at_path(self.controller.reference_prim_path), get_prim_at_path(self.task.root_prim_path)
|
|
)
|
|
print(f"[PICK_DBG] T_world_base translation={T_world_base[:3,3]}", flush=True)
|
|
T_base_world = np.linalg.inv(T_world_base)
|
|
print(f"[PICK_DBG] T_base_world translation={T_base_world[:3,3]}", flush=True)
|
|
T_base_ee = T_base_world[None] @ T_world_ee
|
|
print(f"[PICK_DBG] T_base_ee[0] translation={T_base_ee[0,:3,3]}", flush=True)
|
|
return T_base_ee
|
|
|
|
def get_contact(self, contact_threshold=0.0):
|
|
contact = np.abs(self.pickcontact_view.get_contact_force_matrix()).squeeze()
|
|
contact = np.sum(contact, axis=-1)
|
|
indices = np.where(contact > contact_threshold)[0]
|
|
return contact, 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)
|
|
return len(self.manip_list) == 0
|
|
|
|
def is_success(self):
|
|
flag = True
|
|
|
|
_, indices = self.get_contact()
|
|
if self.gripper_cmd == "close_gripper":
|
|
flag = len(indices) >= 1
|
|
|
|
if self.skill_cfg.get("process_valid", True):
|
|
self.process_valid = np.max(np.abs(self.robot.get_joints_state().velocities)) < 5 and (
|
|
np.max(np.abs(self.pick_obj.get_linear_velocity())) < 5
|
|
)
|
|
flag = flag and self.process_valid
|
|
|
|
if self.skill_cfg.get("lift_th", 0.0) > 0.0:
|
|
p_world_obj = deepcopy(self.pick_obj.get_local_pose()[0])
|
|
flag = flag and ((p_world_obj[2] - self.obj_init_trans[2]) > self.skill_cfg.get("lift_th", 0.0))
|
|
|
|
return flag
|