init commit
This commit is contained in:
zyhe
172
workflows/simbox/core/skills/dexpick.py
Normal file
172
workflows/simbox/core/skills/dexpick.py
Normal file
@@ -0,0 +1,172 @@
|
||||
import os
|
||||
from copy import deepcopy
|
||||
|
||||
import numpy as np
|
||||
from core.skills.base_skill import BaseSkill, register_skill
|
||||
from omegaconf import DictConfig, OmegaConf
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
# pylint: disable=unused-argument
|
||||
@register_skill
|
||||
class Dexpick(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
|
||||
if kwargs:
|
||||
self.world = kwargs["world"]
|
||||
self.skill_cfg = cfg
|
||||
object_name = self.skill_cfg["objects"][0]
|
||||
self.object = 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)
|
||||
dexpick_pose_path = usd_path.replace("Aligned_obj.usd", "dexpick_pose.yaml")
|
||||
self.pick_poses = []
|
||||
if os.path.exists(dexpick_pose_path):
|
||||
with open(dexpick_pose_path, "r", encoding="utf-8") as f:
|
||||
pick_data = OmegaConf.load(f)
|
||||
pick_poses = pick_data.pick_poses
|
||||
for pick_pose in pick_poses:
|
||||
self.pick_poses.append((np.array(pick_pose[:3]), np.array(pick_pose[3:])))
|
||||
|
||||
self.pick_pose_idx = cfg.get("pick_pose_idx", 0)
|
||||
self.pose_ee2o = self.pick_poses[self.pick_pose_idx]
|
||||
self.manip_list = []
|
||||
if "left" in self.controller.robot_file:
|
||||
self.robot_ee_path = self.robot.fl_ee_path
|
||||
self.robot_base_path = self.robot.fl_base_path
|
||||
lr_arm = "left"
|
||||
elif "right" in self.controller.robot_file:
|
||||
self.robot_ee_path = self.robot.fr_ee_path
|
||||
self.robot_base_path = self.robot.fr_base_path
|
||||
lr_arm = "right"
|
||||
else:
|
||||
raise NotImplementedError
|
||||
self.pickcontact_view = task.pickcontact_views[robot.name][lr_arm][object_name]
|
||||
self.process_valid = True
|
||||
self.obj_init_trans = deepcopy(self.object.get_local_pose()[0])
|
||||
|
||||
def simple_generate_manip_cmds(self):
|
||||
manip_list = []
|
||||
|
||||
# Update
|
||||
p_base_ee_cur, q_base_ee_cur = self.controller.get_ee_pose()
|
||||
ignore_substring = deepcopy(self.controller.ignore_substring + self.skill_cfg.get("ignore_substring", []))
|
||||
ignore_substring.append(self.object.name)
|
||||
ignore_substring += self.task.ignore_objects
|
||||
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)
|
||||
|
||||
T_world_base = get_relative_transform(
|
||||
get_prim_at_path(self.robot_base_path), get_prim_at_path(self.task.root_prim_path)
|
||||
)
|
||||
|
||||
# Reach
|
||||
T_world_obj = tf_matrix_from_pose(*self.object.get_local_pose())
|
||||
T_obj_ee_grasp = tf_matrix_from_pose(*self.pose_ee2o)
|
||||
T_world_ee_grasp = T_world_obj @ T_obj_ee_grasp
|
||||
T_base_ee_grasp = np.linalg.inv(T_world_base) @ T_world_ee_grasp
|
||||
p_base_ee_grasp, q_base_ee_grasp = pose_from_tf_matrix(T_base_ee_grasp)
|
||||
|
||||
# Pre grasp
|
||||
pre_grasp_offset = self.skill_cfg.get("pre_grasp_offset", 0.1)
|
||||
if pre_grasp_offset:
|
||||
T_base_ee_pregrasp = T_base_ee_grasp.copy()
|
||||
if "r5a" in self.controller.robot_file:
|
||||
T_base_ee_pregrasp[0:3, 3] -= T_base_ee_pregrasp[0:3, 0] * pre_grasp_offset
|
||||
else:
|
||||
T_base_ee_pregrasp[0:3, 3] -= T_base_ee_pregrasp[0:3, 2] * pre_grasp_offset
|
||||
|
||||
cmd = (*pose_from_tf_matrix(T_base_ee_pregrasp), "open_gripper", {})
|
||||
manip_list.append(cmd)
|
||||
|
||||
# Grasp
|
||||
cmd = (p_base_ee_grasp, q_base_ee_grasp, "open_gripper", {})
|
||||
manip_list.append(cmd)
|
||||
cmd = (p_base_ee_grasp, q_base_ee_grasp, "close_gripper", {})
|
||||
manip_list.extend(
|
||||
[cmd] * self.skill_cfg.get("gripper_change_steps", 40)
|
||||
) # here 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_grasp,
|
||||
q_base_ee_grasp,
|
||||
"update_specific",
|
||||
{"ignore_substring": ignore_substring, "reference_prim_path": self.controller.reference_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_postgrasp = deepcopy(p_base_ee_grasp)
|
||||
p_base_ee_postgrasp[2] += post_grasp_offset
|
||||
cmd = (p_base_ee_postgrasp, q_base_ee_grasp, "close_gripper", {})
|
||||
manip_list.append(cmd)
|
||||
|
||||
self.manip_list = manip_list
|
||||
|
||||
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=10):
|
||||
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):
|
||||
_, indices = self.get_contact()
|
||||
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.object.get_linear_velocity())) < 5
|
||||
)
|
||||
|
||||
flag = flag and self.process_valid
|
||||
|
||||
if self.skill_cfg.get("lift_th", 0.0) > 0.0:
|
||||
obj_curr_trans = deepcopy(self.object.get_local_pose()[0])
|
||||
flag = flag and ((obj_curr_trans[2] - self.obj_init_trans[2]) > self.skill_cfg.get("lift_th", 0.0))
|
||||
|
||||
return flag
|
||||
Reference in New Issue
Block a user