init commit

workflows/simbox/core/skills/__init__.py

@@ -0,0 +1,72 @@
+"""Skills module initialization."""
+
+from core.skills.base_skill import SKILL_DICT
+
+from .approach_rotate import Approach_Rotate
+from .artpreplan import Artpreplan
+from .close import Close
+from .dexpick import Dexpick
+from .dexplace import Dexplace
+from .dynamicpick import Dynamicpick
+from .failpick import FailPick
+from .flip import Flip
+from .goto_pose import Goto_Pose
+from .gripper_action import Gripper_Action
+from .heuristic_skill import Heuristic_Skill
+from .home import Home
+from .joint_ctrl import Joint_Ctrl
+from .manualpick import Manualpick
+from .mobile_rotate import Mobile_Rotate
+from .mobile_translate import Mobile_Translate
+from .move import Move
+from .open import Open
+from .pick import Pick
+from .place import Place
+from .pour_water_succ import Pour_Water_Succ
+from .rotate import Rotate
+from .rotate_obj import Rotate_Obj
+from .scan import Scan
+from .track import Track
+from .wait import Wait
+
+# Explicitly declare the public interface
+__all__ = [
+    "Approach_Rotate",
+    "Artpreplan",
+    "Close",
+    "Dexpick",
+    "Dexplace",
+    "Dynamicpick",
+    "FailPick",
+    "Flip",
+    "Goto_Pose",
+    "Gripper_Action",
+    "Heuristic_Skill",
+    "Home",
+    "Joint_Ctrl",
+    "Mobile_Rotate",
+    "Mobile_Translate",
+    "Move",
+    "Open",
+    "Pick",
+    "Manualpick",
+    "Place",
+    "Pour_Water_Succ",
+    "Rotate",
+    "Rotate_Obj",
+    "Scan",
+    "Track",
+    "Wait",
+    "get_skill_cls",
+    "get_skill_dict",
+]
+
+
+def get_skill_cls(category_name):
+    """Get skill class by category name."""
+    return SKILL_DICT[category_name]
+
+
+def get_skill_dict():
+    """Get skill dictionary."""
+    return SKILL_DICT

workflows/simbox/core/skills/approach_rotate.py

@@ -0,0 +1,201 @@
+# pylint: skip-file
+import numpy as np
+from core.skills.base_skill import BaseSkill, register_skill
+from core.utils.interpolate_utils import linear_interpolation
+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=unused-argument
+@register_skill
+class Approach_Rotate(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.move_obj = task.objects[cfg["objects"][0]]
+        self.approach_obj = task.objects[cfg["objects"][1]]
+        self.skill_cfg = cfg
+        self.rotate_cfg = self.skill_cfg.get("rotate", None)
+
+        self.T_world_base = get_relative_transform(
+            get_prim_at_path(self.controller.robot_base_path), get_prim_at_path(self.task.root_prim_path)
+        )
+
+        self.manip_list = []
+        self.success_threshold_move = self.skill_cfg["success_threshold"]
+        if self.rotate_cfg:
+            self.success_threshold_rotate = self.rotate_cfg["success_threshold"]
+
+    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": self.skill_cfg.get("hold_vec_weight", None)},
+        )
+        manip_list.append(cmd)
+        ignore_substring = 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)
+
+        self.dummy_forward_cfg = self.skill_cfg.get("dummy_forward", None)
+        if self.dummy_forward_cfg:
+            curr_js, tgt_js = self.get_tgt_js()
+            interp_js_list = linear_interpolation(curr_js, tgt_js, self.dummy_forward_cfg.get("num_steps", 10))
+            for js in interp_js_list:
+                p_base_ee, q_base_ee = self.controller.forward_kinematic(js)
+                cmd = (
+                    p_base_ee,
+                    q_base_ee,
+                    "dummy_forward",
+                    {
+                        "arm_action": js,
+                        "gripper_state": self.dummy_forward_cfg.get("gripper_state", 1.0),
+                    },
+                )
+                manip_list.append(cmd)
+
+        self.p_base_ee_tgt, self.q_base_ee_tgt = self.getTgtPose()
+        cmd = (
+            self.p_base_ee_tgt,
+            self.q_base_ee_tgt,
+            "close_gripper",
+            {},
+        )
+        manip_list.append(cmd)
+
+        self.manip_list = manip_list
+
+    def get_tgt_js(self):
+        raise NotImplementedError
+
+    def getTgtPose(self):
+
+        T_world_ee_cur = self.T_world_base @ tf_matrix_from_pose(*self.controller.get_ee_pose())
+        T_world_move_obj_cur = tf_matrix_from_pose(*self.move_obj.get_world_pose())
+
+        if self.skill_cfg.get("obj_axis_offset", None):
+            obj_axis_offset = self.skill_cfg.get("obj_axis_offset", None)
+            axis_map = {"x": 0, "y": 1, "z": 2}
+            dT = np.eye(4)
+            for axis, offset in obj_axis_offset:
+                if axis in axis_map:
+                    dT[axis_map[axis], 3] = offset
+            T_world_move_obj_cur = T_world_move_obj_cur @ dT
+
+        T_world_move_obj_tgt = T_world_move_obj_cur.copy()
+        if self.rotate_cfg:
+            q_world_move_obj_tgt = self.getTgtMoveObjOrientation()
+            T_world_move_obj_tgt[:3, :3] = R.from_quat(q_world_move_obj_tgt, scalar_first=True).as_matrix()
+
+        p_base_ee_tgt, q_base_ee_tgt = self.calTgtEEPose(
+            T_world_ee_cur,
+            T_world_move_obj_cur,
+            T_world_move_obj_tgt,
+        )
+        p_base_ee_tgt[2] += self.skill_cfg.get("z_offset", 0.0)
+
+        return p_base_ee_tgt, q_base_ee_tgt
+
+    def calTgtEEPose(self, T_world_ee_cur=None, T_world_move_obj_cur=None, T_world_move_obj_tgt=None):
+        # Choose object local axis as approach direction (default: x-axis)
+        approach_axis = self.skill_cfg.get("approach_axis", "+x")  # "x", "y", or "z"
+        axis_map = {"+x": [0, 1], "+y": [1, 1], "+z": [2, 1], "-x": [0, -1], "-y": [1, -1], "-z": [2, -1]}
+        local_axis = axis_map[approach_axis][0]
+        local_flag = axis_map[approach_axis][1]
+        approach_dir_world = T_world_move_obj_tgt[:3, local_axis] * (local_flag)  # object axis expressed in world frame
+
+        # Apply additional yaw rotation around world Z axis (obj_yaw_offset)
+        yaw_offset = self.skill_cfg.get("obj_yaw_offset", 0.0) / 180.0 * np.pi
+        if abs(yaw_offset) > 1e-6:
+            R_z = R.from_euler("z", yaw_offset).as_matrix()
+            approach_dir_world = R_z @ approach_dir_world
+
+        # Construct approach point in front of tgt object (world frame)
+        p_world_apr_obj_cur, _ = self.approach_obj.get_world_pose()
+        distance = self.skill_cfg.get("distance", 0.1)
+        p_world_apr_obj_tgt = p_world_apr_obj_cur - distance * approach_dir_world
+        T_world_move_obj_tgt[:2, 3] = p_world_apr_obj_tgt[:2]
+
+        T_base_ee_tgt = (
+            np.linalg.inv(self.T_world_base)
+            @ T_world_move_obj_tgt
+            @ np.linalg.inv(T_world_move_obj_cur)
+            @ T_world_ee_cur
+        )
+
+        return pose_from_tf_matrix(T_base_ee_tgt)
+
+    def getTgtMoveObjOrientation(self):
+        rotate_type = self.rotate_cfg["type"]
+        if rotate_type == "random":
+            q_world_move_obj_tgt = np.random.uniform(*self.rotate_cfg.get("rotate_obj_euler", [[0, 0, 0], [0, 0, 0]]))
+            q_world_move_obj_tgt = (R.from_euler("xyz", q_world_move_obj_tgt, degrees=True)).as_quat(scalar_first=True)
+
+        elif rotate_type == "towards":
+            tgt_obj_2 = self.task.objects[self.rotate_cfg["objects"][1]]
+            curr_obj_2_trans, _ = tgt_obj_2.get_world_pose()
+            x_tgt, y_tgt = curr_obj_2_trans[1] - self.T_world_base[1, 3], -(
+                curr_obj_2_trans[0] - self.T_world_base[0, 3]
+            )
+
+            tgt_obj_yaw = np.arctan2(y_tgt, x_tgt)
+
+            tgt_obj_yaw = np.mod(tgt_obj_yaw + np.pi, 2 * np.pi) - np.pi
+            q_world_move_obj_tgt = np.array([90, 0, tgt_obj_yaw / np.pi * 180])  # z-up
+            R_obj_tgt = R.from_euler("xyz", q_world_move_obj_tgt, degrees=True)
+            q_world_move_obj_tgt = R_obj_tgt.as_quat(scalar_first=True)
+
+        return q_world_move_obj_tgt
+
+    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_pos = 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_pos < 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():
+            self.manip_list.pop(0)
+        return len(self.manip_list) == 0
+
+    def is_success(self):
+        p_base_ee_cur, q_base_ee_cur = self.controller.get_ee_pose()
+        distance = np.linalg.norm(p_base_ee_cur - self.p_base_ee_tgt)
+        flag = (distance < self.success_threshold_move) and (len(self.manip_list) == 0)
+        if self.rotate_cfg:
+            dot = np.clip(np.dot(q_base_ee_cur, self.q_base_ee_tgt), -1.0, 1.0)
+            angle_diff = 2 * np.arccos(np.abs(dot))
+            flag = (angle_diff < self.success_threshold_rotate) and flag
+        return flag

workflows/simbox/core/skills/artpreplan.py

@@ -0,0 +1,165 @@
+from copy import deepcopy
+
+import numpy as np
+from core.skills.base_skill import BaseSkill, register_skill
+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
+from solver.planner import KPAMPlanner
+
+
+# pylint: disable=unused-argument
+@register_skill
+class Artpreplan(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.stage = task.stage
+        self.name = cfg["name"]
+        art_obj_name = cfg["objects"][0]
+        self.skill_cfg = cfg
+        self.art_obj = task.objects[art_obj_name]
+        self.planner_setting = cfg["planner_setting"]
+        self.contact_pose_index = self.planner_setting["contact_pose_index"]
+        self.success_threshold = self.planner_setting["success_threshold"]
+        self.update_art_joint = self.planner_setting.get("update_art_joint", False)
+        if kwargs:
+            self.world = kwargs["world"]
+            self.draw = kwargs["draw"]
+        self.manip_list = []
+
+        if self.skill_cfg.get("obj_info_path", None):
+            self.art_obj.update_articulated_info(self.skill_cfg["obj_info_path"])
+
+        lr_arm = "left" if "left" in self.controller.robot_file else "right"
+        self.fingers_link_contact_view = task.artcontact_views[robot.name][lr_arm][art_obj_name + "_fingers_link"]
+        self.fingers_base_contact_view = task.artcontact_views[robot.name][lr_arm][art_obj_name + "_fingers_base"]
+        self.forbid_collision_contact_view = task.artcontact_views[robot.name][lr_arm][
+            art_obj_name + "_forbid_collision"
+        ]
+        self.collision_valid = True
+        self.process_valid = True
+
+    def setup_kpam(self):
+        self.planner = KPAMPlanner(
+            env=self.world,
+            robot=self.robot,
+            object=self.art_obj,
+            cfg_path=self.planner_setting,
+            controller=self.controller,
+            draw_points=self.draw,
+            stage=self.stage,
+        )
+
+    def simple_generate_manip_cmds(self):
+        if self.skill_cfg.get("obj_info_path", None):
+            self.art_obj.update_articulated_info(self.skill_cfg["obj_info_path"])
+
+        self.setup_kpam()
+        traj_keyframes, sample_times = self.planner.get_keypose()
+        if len(traj_keyframes) == 0 and len(sample_times) == 0:
+            print("No keyframes found, return empty manip_list")
+            self.manip_list = []
+            return
+
+        T_world_base = get_relative_transform(
+            get_prim_at_path(self.robot.base_path), get_prim_at_path(self.task.root_prim_path)
+        )
+        self.traj_keyframes = traj_keyframes
+        self.sample_times = sample_times
+        if self.draw:
+            for keypose in traj_keyframes:
+                self.draw.draw_points([(T_world_base @ np.append(keypose[:3, 3], 1))[:3]], [(0, 0, 0, 1)], [7])
+        manip_list = []
+
+        # Update
+        p_base_ee, q_base_ee = self.controller.get_ee_pose()
+        ignore_substring = deepcopy(self.controller.ignore_substring)
+        cmd = (
+            p_base_ee,
+            q_base_ee,
+            "update_specific",
+            {"ignore_substring": ignore_substring, "reference_prim_path": self.controller.reference_prim_path},
+        )
+        manip_list.append(cmd)
+
+        self.p_base_ee_tgt = p_base_ee
+        self.q_base_ee_tgt = q_base_ee
+        self.manip_list = manip_list
+
+    def update(self):
+        curr_joint_p = self.art_obj._articulation_view.get_joint_positions()[:, self.art_obj.object_joint_index]
+        if self.update_art_joint and self.is_success():
+            self.art_obj._articulation_view.set_joint_position_targets(
+                positions=curr_joint_p, joint_indices=self.art_obj.object_joint_index
+            )
+
+    def get_contact(self, contact_threshold=0.0):
+        contact = {}
+        fingers_link_contact = np.abs(self.fingers_link_contact_view.get_contact_force_matrix()).squeeze()
+        fingers_link_contact = np.sum(fingers_link_contact, axis=-1)
+        fingers_link_contact_indices = np.where(fingers_link_contact > contact_threshold)[0]
+        contact["fingers_link"] = {
+            "fingers_link_contact": fingers_link_contact,
+            "fingers_link_contact_indices": fingers_link_contact_indices,
+        }
+
+        fingers_base_contact = np.abs(self.fingers_base_contact_view.get_contact_force_matrix()).squeeze()
+        fingers_base_contact = np.sum(fingers_base_contact, axis=-1)
+        fingers_base_contact_indices = np.where(fingers_base_contact > contact_threshold)[0]
+        contact["fingers_base"] = {
+            "fingers_base_contact": fingers_base_contact,
+            "fingers_base_contact_indices": fingers_base_contact_indices,
+        }
+
+        forbid_collision_contact = np.abs(self.forbid_collision_contact_view.get_contact_force_matrix()).squeeze()
+        forbid_collision_contact = np.sum(forbid_collision_contact, axis=-1)
+        forbid_collision_contact_indices = np.where(forbid_collision_contact > contact_threshold)[0]
+        contact["forbid_collision"] = {
+            "forbid_collision_contact": forbid_collision_contact,
+            "forbid_collision_contact_indices": forbid_collision_contact_indices,
+        }
+
+        return contact
+
+    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():
+            self.manip_list.pop(0)
+            print("POP one manip cmd")
+        if self.is_success():
+            self.manip_list.clear()
+            print("Close Pre Plan Done")
+        return len(self.manip_list) == 0
+
+    def is_success(self, t_eps=5e-3, o_eps=0.087):
+        p_base_ee_cur, q_base_ee_cur = self.controller.get_ee_pose()
+        diff_pos = np.linalg.norm(p_base_ee_cur - self.p_base_ee_tgt)
+        diff_ori = 2 * np.arccos(min(abs(np.dot(q_base_ee_cur, self.q_base_ee_tgt)), 1.0))
+        pose_flag = np.logical_or(
+            diff_pos < t_eps,
+            diff_ori < o_eps,
+        )
+        return pose_flag

workflows/simbox/core/skills/base_skill.py

@@ -0,0 +1,35 @@
+import re
+from abc import ABC
+
+SKILL_DICT = {}
+
+
+def register_skill(target_class):
+    key = "_".join(re.sub(r"([A-Z0-9])", r" \1", target_class.__name__).split()).lower()
+    # key = target_class.__name__
+    # assert key not in SKILL_DICT
+    SKILL_DICT[key] = target_class
+    return target_class
+
+
+class BaseSkill(ABC):
+    def __init__(self):
+        self.plan_flag = False
+
+    def is_ready(self):
+        return True
+
+    def is_done(self):
+        raise NotImplementedError
+
+    def is_success(self):
+        raise NotImplementedError
+
+    def update(self):
+        pass
+
+    def is_feasible(self):
+        return True
+
+    def is_record(self):
+        return True

workflows/simbox/core/skills/close.py

@@ -0,0 +1,194 @@
+from copy import deepcopy
+
+import numpy as np
+from core.skills.base_skill import BaseSkill, register_skill
+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
+from scipy.spatial.transform import Rotation as R
+from solver.planner import KPAMPlanner
+
+
+# pylint: disable=unused-argument
+@register_skill
+class Close(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.stage = task.stage
+        self.name = cfg["name"]
+        art_obj_name = cfg["objects"][0]
+        self.skill_cfg = cfg
+        self.art_obj = task.objects[art_obj_name]
+        self.planner_setting = cfg["planner_setting"]
+        self.contact_pose_index = self.planner_setting["contact_pose_index"]
+        self.success_threshold = self.planner_setting["success_threshold"]
+        self.update_art_joint = self.planner_setting.get("update_art_joint", False)
+        if kwargs:
+            self.world = kwargs["world"]
+            self.draw = kwargs["draw"]
+        self.manip_list = []
+        # self.draw = True
+
+        if self.skill_cfg.get("obj_info_path", None):
+            self.art_obj.update_articulated_info(self.skill_cfg["obj_info_path"])
+
+        lr_arm = "left" if "left" in self.controller.robot_file else "right"
+        self.fingers_link_contact_view = task.artcontact_views[robot.name][lr_arm][art_obj_name + "_fingers_link"]
+        self.fingers_base_contact_view = task.artcontact_views[robot.name][lr_arm][art_obj_name + "_fingers_base"]
+        self.forbid_collision_contact_view = task.artcontact_views[robot.name][lr_arm][
+            art_obj_name + "_forbid_collision"
+        ]
+        self.collision_valid = True
+        self.process_valid = True
+
+    def setup_kpam(self):
+        self.planner = KPAMPlanner(
+            env=self.world,
+            robot=self.robot,
+            object=self.art_obj,
+            cfg_path=self.planner_setting,
+            controller=self.controller,
+            draw_points=self.draw,
+            stage=self.stage,
+        )
+
+    def simple_generate_manip_cmds(self):
+        if self.skill_cfg.get("obj_info_path", None):
+            self.art_obj.update_articulated_info(self.skill_cfg["obj_info_path"])
+
+        self.setup_kpam()
+        traj_keyframes, sample_times = self.planner.get_keypose()
+        if len(traj_keyframes) == 0 and len(sample_times) == 0:
+            print("No keyframes found, return empty manip_list")
+            self.manip_list = []
+            return
+
+        T_world_base = get_relative_transform(
+            get_prim_at_path(self.robot.base_path), get_prim_at_path(self.task.root_prim_path)
+        )
+        self.traj_keyframes = traj_keyframes
+        self.sample_times = sample_times
+        if self.draw:
+            for keypose in traj_keyframes:
+                self.draw.draw_points([(T_world_base @ np.append(keypose[:3, 3], 1))[:3]], [(0, 0, 0, 1)], [7])
+        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", []))
+        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)
+
+        for i in range(len(self.traj_keyframes)):
+            p_base_ee_tgt = self.traj_keyframes[i][:3, 3]
+            q_base_ee_tgt = R.from_matrix(self.traj_keyframes[i][:3, :3]).as_quat(scalar_first=True)
+            cmd = (p_base_ee_tgt, q_base_ee_tgt, "close_gripper", {})
+            manip_list.append(cmd)
+
+            if i == self.contact_pose_index - 1:
+                p_base_ee = self.traj_keyframes[i][:3, 3]
+                q_base_ee = R.from_matrix(self.traj_keyframes[i][:3, :3]).as_quat(scalar_first=True)
+                ignore_substring = deepcopy(
+                    self.controller.ignore_substring + self.skill_cfg.get("ignore_substring", [])
+                )
+                parent_name = self.art_obj.prim_path.split("/")[-2]
+                ignore_substring.append(parent_name)
+                cmd = (
+                    p_base_ee,
+                    q_base_ee,
+                    "update_specific",
+                    {"ignore_substring": ignore_substring, "reference_prim_path": self.controller.reference_prim_path},
+                )
+                manip_list.append(cmd)
+        self.manip_list = manip_list
+
+    def update(self):
+        curr_joint_p = self.art_obj._articulation_view.get_joint_positions()[:, self.art_obj.object_joint_index]
+        if self.update_art_joint and self.is_success():
+            self.art_obj._articulation_view.set_joint_position_targets(
+                positions=curr_joint_p, joint_indices=self.art_obj.object_joint_index
+            )
+
+    def get_contact(self, contact_threshold=0.0):
+        contact = {}
+        fingers_link_contact = np.abs(self.fingers_link_contact_view.get_contact_force_matrix()).squeeze()
+        fingers_link_contact = np.sum(fingers_link_contact, axis=-1)
+        fingers_link_contact_indices = np.where(fingers_link_contact > contact_threshold)[0]
+        contact["fingers_link"] = {
+            "fingers_link_contact": fingers_link_contact,
+            "fingers_link_contact_indices": fingers_link_contact_indices,
+        }
+
+        fingers_base_contact = np.abs(self.fingers_base_contact_view.get_contact_force_matrix()).squeeze()
+        fingers_base_contact = np.sum(fingers_base_contact, axis=-1)
+        fingers_base_contact_indices = np.where(fingers_base_contact > contact_threshold)[0]
+        contact["fingers_base"] = {
+            "fingers_base_contact": fingers_base_contact,
+            "fingers_base_contact_indices": fingers_base_contact_indices,
+        }
+
+        forbid_collision_contact = np.abs(self.forbid_collision_contact_view.get_contact_force_matrix()).squeeze()
+        forbid_collision_contact = np.sum(forbid_collision_contact, axis=-1)
+        forbid_collision_contact_indices = np.where(forbid_collision_contact > contact_threshold)[0]
+        contact["forbid_collision"] = {
+            "forbid_collision_contact": forbid_collision_contact,
+            "forbid_collision_contact_indices": forbid_collision_contact_indices,
+        }
+
+        return contact
+
+    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():
+            self.manip_list.pop(0)
+            print("POP one manip cmd")
+        if self.is_success():
+            self.manip_list.clear()
+            print("Close Done")
+        return len(self.manip_list) == 0
+
+    def is_success(self):
+        contact = self.get_contact()
+
+        if self.skill_cfg.get("collision_valid", True):
+            self.collision_valid = (
+                self.collision_valid
+                and len(contact["forbid_collision"]["forbid_collision_contact_indices"]) == 0
+                and len(contact["fingers_base"]["fingers_base_contact_indices"]) == 0
+            )
+        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.art_obj.get_linear_velocity())) < 5
+            )
+
+        curr_joint_p = self.art_obj._articulation_view.get_joint_positions()[:, self.art_obj.object_joint_index]
+
+        return np.abs(curr_joint_p) <= self.success_threshold and self.collision_valid and self.process_valid

workflows/simbox/core/skills/dexpick.py

@@ -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

workflows/simbox/core/skills/dexplace.py

@@ -0,0 +1,232 @@
+import os
+from copy import deepcopy
+
+import numpy as np
+from core.skills.base_skill import BaseSkill, register_skill
+from core.utils.usd_geom_utils import compute_bbox
+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,
+)
+from scipy.spatial.transform import Rotation as R
+from scipy.spatial.transform import Slerp
+
+
+# pylint: disable=unused-argument
+@register_skill
+class Dexplace(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
+        self.name = cfg["name"]
+        self.pick_obj = task._task_objects[cfg["objects"][0]]
+        self.place_obj = task._task_objects[cfg["objects"][1]]
+        self.gripper_axis = cfg.get("gripper_axis", None)
+        self.camera_axis_filter = cfg.get("camera_axis_filter", None)
+        self.place_part_prim_path = cfg.get("place_part_prim_path", None)
+        # Get place annotation
+        usd_path = [obj["path"] for obj in task.cfg["objects"] if obj["name"] == self.skill_cfg["objects"][0]][0]
+        usd_path = os.path.join(self.task.asset_root, usd_path)
+        place_range_path = usd_path.replace("Aligned_obj.usd", "place_range.yaml")
+        if os.path.exists(place_range_path):
+            with open(place_range_path, "r", encoding="utf-8") as f:
+                place_data = OmegaConf.load(f)
+                self.x_range = place_data.x_range
+                self.y_range = place_data.y_range
+        else:
+            self.x_range = [0.4, 0.6]
+            self.y_range = [0.4, 0.6]
+        # Get place_prim
+        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
+        # Get left or right
+        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
+        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
+        if kwargs:
+            self.draw = kwargs["draw"]
+        self.manip_list = []
+
+    def simple_generate_manip_cmds(self):
+        manip_list = []
+        place_traj, post_place_level = self.sample_gripper_place_traj()
+        if len(place_traj) > 1:
+            # Having waypoints
+            for waypoint in place_traj[:-1]:
+                p_base_ee_mid, q_base_ee_mid = waypoint[:3], waypoint[3:]
+                cmd = (p_base_ee_mid, q_base_ee_mid, "close_gripper", {})
+                manip_list.append(cmd)
+
+        # The last waypoint
+        p_base_ee_place, q_base_ee_place = place_traj[-1][:3], place_traj[-1][3:]
+        cmd = (p_base_ee_place, q_base_ee_place, "close_gripper", {})
+        manip_list.append(cmd)
+        cmd = (p_base_ee_place, q_base_ee_place, "open_gripper", {})
+        manip_list.extend([cmd] * self.skill_cfg.get("gripper_change_steps", 10))
+
+        # Adding a pose place pose to avoid collision when combining place skill and close skill
+        T_base_ee_place = tf_matrix_from_pose(p_base_ee_place, q_base_ee_place)
+        # Post place
+        T_base_ee_postplace = deepcopy(T_base_ee_place)
+        # Retreat for a bit along gripper axis
+        if "r5a" in self.controller.robot_file:
+            T_base_ee_postplace[0:3, 3] = T_base_ee_postplace[0:3, 3] - T_base_ee_postplace[0:3, 0] * post_place_level
+        else:
+            T_base_ee_postplace[0:3, 3] = T_base_ee_postplace[0:3, 3] - T_base_ee_postplace[0:3, 2] * post_place_level
+        cmd = (*pose_from_tf_matrix(T_base_ee_postplace), "open_gripper", {})
+        manip_list.append(cmd)
+        self.manip_list = manip_list
+
+    def sample_gripper_place_traj(self):
+        place_traj = []
+        T_base_ee = get_relative_transform(get_prim_at_path(self.robot_ee_path), get_prim_at_path(self.robot_base_path))
+        T_world_base = get_relative_transform(
+            get_prim_at_path(self.robot_base_path), get_prim_at_path(self.task.root_prim_path)
+        )
+        T_world_ee = T_world_base @ T_base_ee
+        p_world_ee_start, q_world_ee_start = pose_from_tf_matrix(T_world_ee)
+        # Getting the object pose
+        T_world_obj = tf_matrix_from_pose(*self.pick_obj.get_local_pose())
+        # Calculate the pose of the end-effector in the object's coordinate frame
+        T_obj_world = np.linalg.inv(T_world_obj)
+        # Getting the relation pose and distance of ee to object (after picking, before placing)
+        T_obj_ee = T_obj_world @ T_world_ee
+        ee2o_distance = np.linalg.norm(T_obj_ee[0:3, 3])
+        place_part_prim = get_prim_at_path(self.place_prim_path)
+        bbox_place_obj = compute_bbox(place_part_prim)
+        x_min, y_min, z_min = bbox_place_obj.min
+        x_max, y_max, z_max = bbox_place_obj.max
+        self.place_boundary = [[x_min, y_min, z_min], [x_max, y_max, z_max]]
+        # Calculate the bounding box vertices
+        vertices = [
+            [x_min, y_min, z_min],
+            [x_min, y_max, z_min],
+            [x_max, y_min, z_min],
+            [x_max, y_max, z_min],
+            [x_min, y_min, z_max],
+            [x_min, y_max, z_max],
+            [x_max, y_min, z_max],
+            [x_max, y_max, z_max],
+        ]
+        # Draw the bounding box vertices
+        if self.draw:
+            for vertex in vertices:
+                self.draw.draw_points([vertex], [(0, 0, 0, 1)], [7])  # black
+
+        # 1. Obtaining ee_ori
+        p_world_ee_init = self.controller.T_world_ee_init[0:3, 3]  # getting initial ee position
+        container_position = self.place_obj.get_local_pose()[0]  # getting container position
+        container_position[1] += 0.0
+        gripper_axis = container_position - p_world_ee_init  # gripper_axis is aligned with the container direction
+        gripper_axis = gripper_axis / np.linalg.norm(gripper_axis)  # Normalize the target vector
+        q_world_ee = self.get_ee_ori(gripper_axis, T_world_ee, self.camera_axis_filter)
+        # 2. Obtaining p_world_ee
+        x = x_min + np.random.uniform(self.x_range[0], self.x_range[1]) * (x_max - x_min)
+        y = y_min + np.random.uniform(self.y_range[0], self.y_range[1]) * (y_max - y_min)
+        z = z_min + 0.15
+        obj_place_position = np.array([x, y, z])
+        if self.draw:
+            self.draw.draw_points([obj_place_position.tolist()], [(1, 0, 0, 1)], [7])  # red
+        p_world_ee = obj_place_position - gripper_axis * ee2o_distance
+        # 3. Adding Waypoint
+        # Pre place
+        p_world_ee_mid = (p_world_ee_start + p_world_ee) / 2.0
+        p_world_ee_mid[2] += 0.05
+        slerp = Slerp([0, 1], R.from_quat([q_world_ee_start, q_world_ee]))
+        q_world_ee_mid = slerp([0.5]).as_quat()[0]
+        if self.draw:
+            self.draw.draw_points([p_world_ee_mid.tolist()], [(0, 1, 0, 1)], [7])  # green
+        place_traj.append(self.adding_waypoint(p_world_ee_mid, q_world_ee_mid, T_world_base))
+        # Place
+        if self.draw:
+            self.draw.draw_points([p_world_ee.tolist()], [(0, 1, 0, 1)], [7])  # green
+        place_traj.append(self.adding_waypoint(p_world_ee, q_world_ee, T_world_base))
+        post_place_level = 0.1
+
+        return place_traj, post_place_level
+
+    def get_ee_ori(self, gripper_axis, T_world_ee, camera_axis_filter=None):
+        gripper_x = gripper_axis
+        if camera_axis_filter is not None:
+            direction = camera_axis_filter[0]["direction"]
+            degree = camera_axis_filter[1]["degree"]
+            direction = np.array(direction) / np.linalg.norm(direction)  # Normalize the direction vector
+            angle = np.radians(np.random.uniform(degree[0], degree[1]))
+            gripper_z = direction - np.dot(direction, gripper_x) * gripper_x
+            gripper_z = gripper_z / np.linalg.norm(gripper_z)
+            rotation_axis = np.cross(gripper_z, gripper_x)
+            rotation_axis = rotation_axis / np.linalg.norm(rotation_axis)
+            gripper_z = R.from_rotvec(angle * rotation_axis).apply(gripper_z)
+
+        else:
+            current_z = T_world_ee[0:3, 1]
+            gripper_z = current_z - np.dot(current_z, gripper_x) * gripper_x
+
+        gripper_z = gripper_z / np.linalg.norm(gripper_z)
+        gripper_y = np.cross(gripper_z, gripper_x)
+        gripper_y = gripper_y / np.linalg.norm(gripper_y)
+        gripper_z = np.cross(gripper_x, gripper_y)
+        R_world_ee = np.column_stack((gripper_x, gripper_y, gripper_z))
+        q_world_ee = R.from_matrix(R_world_ee).as_quat(scalar_first=True)
+        return q_world_ee
+
+    def adding_waypoint(self, p_world_ee, q_world_ee, T_world_base):
+        """
+        Adding a waypoint, also transform from wolrd frame to robot frame
+        """
+        T_world_ee = tf_matrix_from_pose(p_world_ee, q_world_ee)
+        T_base_ee = np.linalg.inv(T_world_base) @ T_world_ee
+        p_base_ee, q_base_ee = pose_from_tf_matrix(T_base_ee)
+        waypoint = np.concatenate((p_base_ee, q_base_ee))
+        return waypoint
+
+    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):
+        x, y, z = self.pick_obj.get_local_pose()[0]  # pick_obj position
+        within_boundary = (
+            self.place_boundary[0][0] <= x <= self.place_boundary[1][0]
+            and self.place_boundary[0][1] <= y <= self.place_boundary[1][1]
+            and self.place_boundary[0][2] <= z  # <= self.place_boundary[1][2]
+        )
+
+        print("pos :", self.pick_obj.get_local_pose()[0])
+        print("boundary :", self.place_boundary)
+        print("within_boundary :", within_boundary)
+
+        return within_boundary

workflows/simbox/core/skills/dynamicpick.py

@@ -0,0 +1,390 @@
+# pylint: skip-file
+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,
+)
+from omni.timeline import get_timeline_interface
+from scipy.spatial.transform import Rotation as R
+
+
+# pylint: disable=unused-argument
+@register_skill
+class Dynamicpick(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]
+        self.predict_pick = False
+        self.meet_pose_o2w = None
+        self.grasp_scale = self.skill_cfg.get("grasp_scale", 1)
+        self.tcp_offset = self.skill_cfg.get("tcp_offset", 0.125)
+
+        # 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", "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.grasp_scale, tcp_offset=self.tcp_offset
+        )
+        self.robot_name = self.controller.robot_file.split("/")[-1].split(".yml")[0]
+        self.object_name = object_name
+
+        # !!! 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.cmd_time = 0
+        self.delta_x = np.random.uniform(self.skill_cfg["pick_range"][0], self.skill_cfg["pick_range"][1])
+        self.time_bias = self.skill_cfg.get("time_bias", 0)
+        self.pick_bias = self.skill_cfg.get("pick_bias", 0)
+        self.process_valid = True
+        self.obj_init_trans = deepcopy(self.pick_obj.get_local_pose()[0])
+
+    def simple_generate_manip_cmds(self):
+        pass
+
+    def predict_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)
+        ignore_substring += self.task.ignore_objects
+        self.controller.update_specific(ignore_substring, self.controller.reference_prim_path)
+        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)
+
+        cmd_time, expected_js = self.controller.pre_forward(p_base_ee_cur, q_base_ee_cur, ds_ratio=2)
+        self.cmd_time += cmd_time
+
+        # Pre grasp
+        T_base_ee_grasps = self.sample_ee_pose()  # (N, 4, 4)
+        # Batch grasp pose adjustment if needed (operate on all T_base_ee_grasps at once)
+        if self.skill_cfg.get("pivot_angle_z", None) is not None:
+            num_grasps = T_base_ee_grasps.shape[0]
+            # sample per-grasp pivot angles
+            pivot_angles_z = np.random.uniform(
+                self.skill_cfg["pivot_angle_z"][0],
+                self.skill_cfg["pivot_angle_z"][1],
+                size=num_grasps,
+            )
+            # compute batch rotation matrices R_z(-pivot_angle_z)
+            pivot_rotations = R.from_euler("z", -pivot_angles_z, degrees=True).as_matrix()  # (N, 3, 3)
+            # apply rotations to all rotation blocks
+            T_base_ee_grasps[:, :3, :3] = np.einsum("nij,njk->nik", T_base_ee_grasps[:, :3, :3], pivot_rotations)
+            # sample per-grasp z translation adjustments
+            pos_adjust_z = np.random.uniform(
+                self.skill_cfg["pos_adjust_z"][0],
+                self.skill_cfg["pos_adjust_z"][1],
+                size=num_grasps,
+            )
+            T_base_ee_grasps[:, 2, 3] += pos_adjust_z
+        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
+
+        # Pre-grasp
+        cmd = (p_base_ee_pregrasps[index], q_base_ee_pregrasps[index], "open_gripper", {})
+        manip_list.append(cmd)
+        cmd_time, expected_js = self.controller.pre_forward(
+            p_base_ee_pregrasps[index], q_base_ee_pregrasps[index], expected_js, ds_ratio=2
+        )
+        self.cmd_time += cmd_time
+
+        # Grasp
+        cmd = (p_base_ee_grasps[index], q_base_ee_grasps[index], "open_gripper", {})
+        manip_list.append(cmd)
+        cmd_time, expected_js = self.controller.pre_forward(
+            p_base_ee_grasps[index], q_base_ee_grasps[index], expected_js, ds_ratio=2
+        )
+        self.cmd_time += cmd_time
+        cmd = (p_base_ee_grasps[index], q_base_ee_grasps[index], "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
+
+        # 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], "close_gripper", {})
+            manip_list.append(cmd)
+
+        self.manip_list = manip_list
+        self.cmd_time += self.time_bias
+
+    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 = 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])
+        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
+
+        if self.meet_pose_o2w is not None:
+            T_world_obj = tf_matrix_from_pose(*self.meet_pose_o2w)
+        else:
+            T_world_obj = tf_matrix_from_pose(*self.pick_obj.get_local_pose())
+        T_world_ee = T_world_obj[None] @ self.T_obj_ee
+
+        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)
+            )
+            T_base_world = np.linalg.inv(T_world_base)
+            T_base_ee = T_base_world[None] @ T_world_ee
+            return T_base_ee
+
+    def is_ready(self):
+        object_position = self.pick_obj.get_local_pose()[0]
+        ee_init_position = deepcopy(self.controller.T_world_ee_init[0:3, 3])
+        x = object_position[0] - ee_init_position[0]
+        self.obj_velocity = self.task.conveyor_velocity
+        if (self.obj_velocity < 0 and x < 0.5) or (self.obj_velocity > 0 and x > -0.5):
+            if not self.predict_pick:
+                print(f"###{self.robot_name} PREDICTING {self.object_name}###")
+                position = deepcopy(object_position)
+                delta_x = self.delta_x
+                position[0] = ee_init_position[0] + delta_x
+                orientation = self.pick_obj.get_local_pose()[1]
+                self.meet_pose_o2w = (position, orientation)
+                self.predict_manip_cmds()
+                self.epsilon = delta_x - (self.cmd_time * self.obj_velocity) + self.pick_bias
+                self.predict_pick = True
+            if (self.obj_velocity < 0 and x < self.epsilon) or (
+                self.obj_velocity > 0 and x > self.epsilon
+            ):  # start real pick
+                return True
+            else:
+                return False
+        else:
+            return False
+
+    def get_obj_velocity(self, x):
+        timeline = get_timeline_interface()
+        current_time = timeline.get_current_time()
+        previous_time = getattr(self, "_previous_time", None)
+        previous_x = getattr(self, "_previous_x", None)
+        if previous_time is not None and previous_x is not None:
+            time_delta = current_time - previous_time
+            if time_delta > 0:
+                x_velocity = (x - previous_x) / time_delta
+            else:
+                x_velocity = 0
+        else:
+            x_velocity = 0
+
+        self._previous_time = current_time
+        self._previous_x = x
+
+        return x_velocity
+
+    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 not self.is_ready():
+            return False
+        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.pick_obj.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.pick_obj.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

workflows/simbox/core/skills/failpick.py

@@ -0,0 +1,231 @@
+# pylint: skip-file
+import os
+from copy import deepcopy
+
+import numpy as np
+from core.skills.base_skill import BaseSkill, register_skill
+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,
+)
+
+
+# pylint: disable=unused-argument
+@register_skill
+class FailPick(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.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)
+        grasp_pose_path = usd_path.replace("Aligned_obj.usd", "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_ee2o, self.scores = self.robot.pose_post_process_fn(sparse_grasp_poses, lr_arm=lr_arm)
+
+        # !!! keyposes should be generated after previous skill is done
+        self.manip_list = []
+        self.process_valid = True
+        self.obj_init_trans = deepcopy(self.object.get_local_pose()[0])
+
+    def simple_generate_manip_cmds(self):
+        manip_list = []
+
+        # Update
+        ignore_substring = deepcopy(self.controller.ignore_substring + self.skill_cfg.get("ignore_substring", []))
+        ignore_substring.append(self.object.name)
+
+        # Pre grasp
+        poses = self.sample_ee_pose()
+        self.controller.update_specific(
+            ignore_substring=ignore_substring, reference_prim_path=self.controller.reference_prim_path
+        )
+        for pose in poses:
+            grasp_trans, grasp_ori = pose_from_tf_matrix(pose)
+            x_offset = np.random.choice(
+                [
+                    np.random.uniform(
+                        self.skill_cfg.get("grasp_x_offset_min", 0.05), self.skill_cfg.get("grasp_x_offset_max", 0.1)
+                    ),
+                    np.random.uniform(
+                        -self.skill_cfg.get("grasp_x_offset_max", 0.1), -self.skill_cfg.get("grasp_x_offset_min", 0.05)
+                    ),
+                ],
+            )
+            y_offset = np.random.choice(
+                [
+                    np.random.uniform(
+                        self.skill_cfg.get("grasp_y_offset_min", 0.05), self.skill_cfg.get("grasp_y_offset_max", 0.1)
+                    ),
+                    np.random.uniform(
+                        -self.skill_cfg.get("grasp_y_offset_max", 0.1), -self.skill_cfg.get("grasp_y_offset_min", 0.05)
+                    ),
+                ],
+            )
+            grasp_trans[0] += x_offset
+            grasp_trans[1] += y_offset
+
+            test_mode = self.skill_cfg.get("test_mode", "forward")
+            if test_mode == "forward":
+                result = self.controller.test_single_forward(grasp_trans, grasp_ori)
+            elif test_mode == "ik":
+                result = self.controller.test_single_ik(grasp_trans, grasp_ori)
+            else:
+                raise NotImplementedError
+            if result == 1:
+                print("pick plan success")
+                break
+
+        cmd = (grasp_trans, grasp_ori, "close_gripper", {})
+        manip_list.append(cmd)
+
+        # Grasp
+        manip_list.extend([cmd] * self.skill_cfg.get("gripper_change_steps", 10))
+
+        # Post grasp
+        post_grasp_pose = pose.copy()
+        post_grasp_pose[2, 3] += np.random.uniform(
+            self.skill_cfg.get("post_grasp_offset_min", 0.05), self.skill_cfg.get("post_grasp_offset_max", 0.05)
+        )
+        cmd = (*pose_from_tf_matrix(post_grasp_pose), "close_gripper", {})
+        manip_list.append(cmd)
+
+        # cmd = (*pose_from_tf_matrix(post_grasp_pose), 'open_gripper', {})
+        # manip_list.extend([cmd] * self.skill_cfg.get("gripper_change_steps", 10))
+
+        self.manip_list = manip_list
+
+    def sample_ee_pose(self, max_length=30):
+        T_ee2r = self.get_ee_poses("robot")
+
+        num_pose = T_ee2r.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_ee2r[:, row, col] >= cos_val if sign > 0 else T_ee2r[:, 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] = T_ee2r[:, row, col] >= cos_val1 and T_ee2r[:, row, col] <= cos_val2
+                        flags[axis] = np.logical_and(T_ee2r[:, row, col] >= cos_val1, T_ee2r[:, row, col] <= cos_val2)
+                    else:
+                        # flags[axis] = T_ee2r[:, row, col] <= cos_val1 and T_ee2r[:, row, col] >= cos_val2
+                        flags[axis] = np.logical_and(T_ee2r[:, row, col] <= cos_val1, T_ee2r[:, 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_o2w = tf_matrix_from_pose(*self.object.get_local_pose())
+            T_w2r = get_relative_transform(
+                get_prim_at_path(self.task.root_prim_path), get_prim_at_path(self.controller.reference_prim_path)
+            )
+            T_o2r = T_w2r @ T_o2w
+            if direction_to_obj == "right":
+                flags["direction_to_obj"] = T_ee2r[:, 1, 3] <= T_o2r[1, 3]
+            elif direction_to_obj == "left":
+                flags["direction_to_obj"] = T_ee2r[:, 1, 3] > T_o2r[1, 3]
+            else:
+                raise NotImplementedError
+
+        combined_flag = np.logical_and.reduce(list(flags.values()))
+        if sum(combined_flag) == 0:
+            # idx = np.random.choice(np.arange(num_pose), size=max_length, replace=True)
+            idx = [0]
+        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 = [idx for (score, idx) in combined[:max_length]]
+            # idx = np.random.choice(np.arange(num_pose)[combined_flag], size=max_length, replace=True)
+        return T_ee2r[idx]
+
+    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_ee2o
+
+        T_o2w = tf_matrix_from_pose(*self.object.get_local_pose())
+        T_ee2w = T_o2w[None] @ self.T_ee2o
+
+        if frame == "world":
+            return T_ee2w
+
+        if frame == "robot":  # robot base frame
+            T_r2w = get_relative_transform(
+                get_prim_at_path(self.controller.reference_prim_path), get_prim_at_path(self.task.root_prim_path)
+            )
+            T_w2r = np.linalg.inv(T_r2w)
+            T_ee2r = T_w2r[None] @ T_ee2w
+            return T_ee2r
+
+    def is_subtask_done(self, t_eps=1e-3, o_eps=5e-3):
+        assert len(self.manip_list) != 0
+        curr_ee_trans, curr_ee_ori = self.controller.get_ee_pose()
+        ee_trans, ee_ori, *_ = self.manip_list[0]
+        diff_trans = np.linalg.norm(curr_ee_trans - ee_trans)
+        diff_ori = 2 * np.arccos(min(abs(np.dot(curr_ee_ori, ee_ori)), 1.0))
+        pose_flag = np.logical_and(
+            diff_trans < t_eps,
+            diff_ori < o_eps,
+        )
+        self.plan_flag = self.controller.num_last_cmd > 10
+        # self.plan_flag = False
+        return np.logical_or(pose_flag, self.plan_flag)
+
+    def is_record(self):
+        return len(self.manip_list) < (1 * self.skill_cfg.get("gripper_change_steps", 10) + 2)
+
+    def is_success(self):
+        return True
+
+    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)):
+            # set_trace()
+            self.manip_list.pop(0)
+        return len(self.manip_list) == 0
+
+    def is_feasible(self, th=5):
+        return self.controller.num_plan_failed <= th

workflows/simbox/core/skills/flip.py

@@ -0,0 +1,169 @@
+from copy import deepcopy
+
+import numpy as np
+from core.skills.base_skill import BaseSkill, register_skill
+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=unused-argument
+@register_skill
+class Flip(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.objects[cfg["objects"][0]]
+        self.skill_cfg = cfg
+        self.gripper_axis = cfg.get("gripper_axis", False)
+        self.manip_list = []
+        if kwargs:
+            self.draw = kwargs["draw"]
+        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
+        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
+
+    def simple_generate_manip_cmds(self):
+        manip_list = []
+        place_traj, post_place_level = self.sample_place_traj()
+        if len(place_traj) > 1:
+            # Having waypoints
+            for waypoint in place_traj[:-1]:
+                p_base_ee, q_base_ee = waypoint[:3], waypoint[3:]
+                cmd = (p_base_ee, q_base_ee, "close_gripper", {})
+                manip_list.append(cmd)
+        # The last waypoint
+        p_base_ee_place, q_base_ee_place = place_traj[-1][:3], place_traj[-1][3:]
+        cmd = (p_base_ee_place, q_base_ee_place, "close_gripper", {})
+        manip_list.append(cmd)
+        cmd = (p_base_ee_place, q_base_ee_place, "open_gripper", {})
+        manip_list.extend(
+            [cmd] * self.skill_cfg.get("open_wait_steps", 20)
+        )  # here we use 20 steps to make sure the gripper is fully open
+
+        # Adding a pose place pose to avoid collision when combining place skill and close skill
+        T_base_ee_place = tf_matrix_from_pose(p_base_ee_place, q_base_ee_place)
+        # Post place
+        T_base_ee_postplace = deepcopy(T_base_ee_place)
+        # Retreat for a bit along gripper axis
+        T_base_ee_postplace[0:3, 3] = T_base_ee_postplace[0:3, 3] - T_base_ee_postplace[0:3, 0] * post_place_level
+        cmd = (*pose_from_tf_matrix(T_base_ee_postplace), "open_gripper", {})
+        manip_list.append(cmd)
+        self.manip_list = manip_list
+
+    def sample_place_traj(self):
+        place_traj = []
+        T_base_ee = get_relative_transform(get_prim_at_path(self.robot_ee_path), get_prim_at_path(self.robot_base_path))
+        T_world_base = get_relative_transform(get_prim_at_path(self.robot_base_path), get_prim_at_path("/World"))
+        T_world_ee = T_world_base @ T_base_ee
+
+        # 1. Obtaining ee_ori
+        gripper_axis = np.array(self.gripper_axis)
+        gripper_axis = gripper_axis / np.linalg.norm(gripper_axis)  # Normalize the vector
+        camera_axis = np.array([0, 1, 0])
+        q_world_ee = self.get_ee_ori(gripper_axis, T_world_ee, camera_axis)
+        # 2. Obtaining ee_trans
+        p_world_ee_init = self.controller.T_world_ee_init[0:3, 3]  # getting initial ee position
+        p_world_ee = p_world_ee_init.copy()
+        p_world_ee[0] += np.random.uniform(0.19, 0.21)  # 0.2
+        p_world_ee[1] += np.random.uniform(0.23, 0.27)  # 0.25
+        p_world_ee[2] += 0
+        if self.draw:
+            self.draw.draw_points([p_world_ee.tolist()], [(1, 0, 0, 1)], [7])  # red
+        # 3. Adding waypoint
+        # Pre place
+        p_world_ee_mid = p_world_ee_init.copy()
+        p_world_ee_mid[1] += np.random.uniform(0.23, 0.27)  # 0.25
+        p_world_ee_mid[2] += np.random.uniform(0.14, 0.16)  # 0.15
+        gripper_axis_mid = np.array([0, 1, -1])
+        camera_axis_mid = np.array([0, 1, 1])
+        gripper_axis_mid = gripper_axis_mid / np.linalg.norm(gripper_axis_mid)
+        camera_axis_mid = camera_axis_mid / np.linalg.norm(camera_axis_mid)
+        q_world_ee_mid = self.get_ee_ori(gripper_axis_mid, T_world_ee, camera_axis_mid)
+
+        if self.draw:
+            self.draw.draw_points([p_world_ee_mid.tolist()], [(1, 0, 0, 1)], [7])  # red
+        place_traj.append(self.adding_waypoint(p_world_ee_mid, q_world_ee_mid, T_world_base))
+        # Place
+        place_traj.append(self.adding_waypoint(p_world_ee, q_world_ee, T_world_base))
+        post_place_level = 0.1
+
+        return place_traj, post_place_level
+
+    def get_ee_ori(self, gripper_axis, T_world_ee, camera_axis=None):
+        gripper_x = gripper_axis
+        if camera_axis is not None:
+            gripper_z = camera_axis
+        else:
+            current_z = T_world_ee[0:3, 1]
+            gripper_z = current_z - np.dot(current_z, gripper_x) * gripper_x
+
+        gripper_z = gripper_z / np.linalg.norm(gripper_z)
+        gripper_y = np.cross(gripper_z, gripper_x)
+        gripper_y = gripper_y / np.linalg.norm(gripper_y)
+        gripper_z = np.cross(gripper_x, gripper_y)
+        R_world_ee = np.column_stack((gripper_x, gripper_y, gripper_z))
+        q_world_ee = R.from_matrix(R_world_ee).as_quat(scalar_first=True)
+        return q_world_ee
+
+    def adding_waypoint(self, p_world_ee, q_world_ee, T_world_base):
+        """
+        Adding a waypoint, also transform from wolrd frame to robot frame
+        """
+        T_world_ee = tf_matrix_from_pose(p_world_ee, q_world_ee)
+        T_base_ee = np.linalg.inv(T_world_base) @ T_world_ee
+        p_base_ee, q_base_ee = pose_from_tf_matrix(T_base_ee)
+        waypoint = np.concatenate((p_base_ee, q_base_ee))
+        return waypoint
+
+    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):
+        # Calculate the angle between the object's local y-axis and the world's z-axis
+        T_world_obj = tf_matrix_from_pose(*self.pick_obj.get_local_pose())
+        obj_y_axis = T_world_obj[0:3, 1]  # Extract the object's y-axis in world coordinates
+        world_z_axis = np.array([0, 0, 1])  # World z-axis
+        # Compute the angle between the two vectors
+        dot_product = np.dot(obj_y_axis, world_z_axis)
+        angle = np.arccos(np.clip(dot_product, -1.0, 1.0))  # Clip to handle numerical errors
+        angle_degrees = np.degrees(angle)
+        # Position
+        ee_init_position = self.controller.T_world_ee_init[0:3, 3]  # getting initial ee position
+        obj_position = T_world_obj[0:3, 3]
+        delta_y = obj_position[1] - ee_init_position[1]
+        return angle_degrees < 90 and 0 < delta_y < 0.7

workflows/simbox/core/skills/goto_pose.py

@@ -0,0 +1,272 @@
+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_single
+from core.utils.transformation_utils import (
+    create_pose_matrices,
+    get_orientation,
+    perturb_orientation,
+    perturb_position,
+    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,
+)
+from scipy.spatial.transform import Rotation as R
+from scipy.spatial.transform import Slerp
+
+
+# pylint: disable=line-too-long,unused-argument
+@register_skill
+class Goto_Pose(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
+        self.frame = self.skill_cfg.get("frame", "robot")
+        gripper_state = self.skill_cfg.get("gripper_state", 1)
+        self.gripper_fn = "open_gripper" if gripper_state == 1 else "close_gripper"
+        position = self.skill_cfg.get("position", None)
+        if position is not None:
+            self.p_base_ee_tgt = perturb_position(position, self.skill_cfg.get("max_noise_m", 0.0))
+        else:
+            raise KeyError(
+                f"Required config 'position' not found in skill_cfg. Available keys: {list(self.skill_cfg.keys())}"
+            )
+
+        if self.skill_cfg.get("quaternion", None) or self.skill_cfg.get("euler", None):
+            self.q_base_ee_tgt = perturb_orientation(
+                np.array(get_orientation(self.skill_cfg.get("euler"), self.skill_cfg.get("quaternion"))),
+                self.skill_cfg.get("max_noise_deg", 0),
+            )
+        else:
+            self.q_base_ee_tgt = None
+
+        # !!! keyposes should be generated after previous skill is done
+        self.manip_list = []
+
+        # self.robot_ee_path = self.controller.robot_ee_path
+        # self.robot_base_path = self.controller.robot_base_path
+        # T_base_world = get_relative_transform(
+        #     get_prim_at_path(self.task.root_prim_path),
+        #     get_prim_at_path(self.robot_base_path)
+        # )
+        # self.target_trans = np.array([-0.05, -0.2, 1.1])
+        # place_pos = T_base_world[:3, :3] @ self.target_trans + T_base_world[:3, 3]
+
+        # from pdb import set_trace
+        # set_trace()
+
+    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.q_base_ee_tgt is None:
+            # Start Filter according to constraints
+            obj = self.task.objects[self.skill_cfg["objects"][0]]
+            T_world_obj = tf_matrix_from_pose(*obj.get_local_pose())
+            T_world_ee = get_relative_transform(
+                get_prim_at_path(self.controller.robot_ee_path), get_prim_at_path(self.task.root_prim_path)
+            )
+            self.T_obj_ee = np.linalg.inv(T_world_obj) @ T_world_ee
+
+            self.align_obj_axis = self.skill_cfg["align_obj_axis"]
+            self.align_ref_axis = self.skill_cfg["align_ref_axis"]
+            self.align_obj_tol = self.skill_cfg["align_obj_tol"]
+            R_tgts = self.generate_constrained_rotation_batch()
+            p_tgts = perturb_position(
+                self.skill_cfg.get("position", None), self.skill_cfg.get("max_noise_m", 0.0), CUROBO_BATCH_SIZE
+            )
+
+            pos_constraint = self.skill_cfg.get("position_constraint", "gripper")
+            if pos_constraint == "gripper":
+                T_base_ee_tgts = create_pose_matrices(p_tgts, R_tgts)
+            elif 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_tgts = create_pose_matrices(p_tgts, R_base_obj_tgts)
+                T_base_ee_tgts = T_base_obj_tgts @ self.T_obj_ee
+            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_tgts, q_base_ee_tgts = poses_from_tf_matrices(T_base_ee_tgts)
+            if self.controller.use_batch:
+                result = self.controller.test_batch_forward(p_base_ee_tgts, q_base_ee_tgts)
+                index = select_index_by_priority_single(result)
+            else:
+                for index in range(T_base_ee_tgts.shape[0]):
+                    p_base_ee_tgt, q_base_ee_tgt = p_base_ee_tgts[index], q_base_ee_tgts[index]
+                    test_mode = self.skill_cfg.get("test_mode", "forward")
+                    if test_mode == "forward":
+                        result_pre = self.controller.test_single_forward(p_base_ee_tgt, q_base_ee_tgt)
+                    elif test_mode == "ik":
+                        result_pre = self.controller.test_single_ik(p_base_ee_tgt, q_base_ee_tgt)
+                    else:
+                        raise NotImplementedError
+                    if result_pre == 1:
+                        print("goto pose plan success")
+                        break
+            self.p_base_ee_tgt, self.q_base_ee_tgt = p_base_ee_tgts[index], q_base_ee_tgts[index]
+
+        ignore_substring = self.controller.ignore_substring + self.skill_cfg.get("ignore_substring", [])
+        interp_nums = self.skill_cfg.get("interp_nums", 1)
+        if interp_nums >= 2:
+            interp_trans_list, interp_ori_list = self.interp(
+                p_base_ee_cur, q_base_ee_cur, self.p_base_ee_tgt, self.q_base_ee_tgt, interp_nums
+            )
+            for interp_trans, interp_ori in zip(interp_trans_list, interp_ori_list):
+                # if self.controller.test_single_forward(interp_trans, interp_ori):
+                cmd = (
+                    interp_trans,
+                    interp_ori,
+                    "update_specific",
+                    {
+                        "ignore_substring": ignore_substring,
+                        "reference_prim_path": self.controller.reference_prim_path,
+                    },
+                )
+                manip_list.append(cmd)
+        else:
+            cmd = (self.p_base_ee_tgt, self.q_base_ee_tgt, self.gripper_fn, {})
+            manip_list.append(cmd)
+
+        self.manip_list = manip_list
+
+    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_obj_axis is not None and self.align_ref_axis is not None and self.align_obj_tol is not None:
+
+            align_obj_axis = np.array(self.align_obj_axis)
+            align_ref_axis = np.array(self.align_ref_axis)
+            R_ee_obj = np.linalg.inv(self.T_obj_ee)[:3, :3]
+            Rs_base_obj_tgt = rot_mats @ R_ee_obj  # (N, 3, 3)
+            align_obj_vecs = np.einsum("ijk,k->ij", Rs_base_obj_tgt, align_obj_axis)  # (N, 3)
+            align_ref_vec = np.eye(3) @ align_ref_axis
+            dot_products = np.dot(align_obj_vecs, align_ref_vec)
+            norms = np.linalg.norm(align_obj_vecs, axis=1) * np.linalg.norm(align_ref_vec)
+            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 interp(self, curr_trans, curr_ori, target_trans, target_ori, interp_num, normalize_quaternions=True):
+        # position interpolation (linear)
+        interp_trans = np.linspace(curr_trans, target_trans, interp_num, axis=0)
+
+        # Rotation interpolation (spherical linear interpolation - SLERP)
+        if normalize_quaternions:
+            curr_ori = curr_ori / np.linalg.norm(curr_ori)
+            target_ori = target_ori / np.linalg.norm(target_ori)
+
+        # SciPy uses [x, y, z, w] format
+        rotations = R.from_quat([curr_ori, target_ori], scalar_first=True)  # Current  # Target
+
+        slerp = Slerp([0, 1], rotations)
+        times = np.linspace(0, 1, interp_num)
+        interp_ori = slerp(times).as_quat(scalar_first=True)
+
+        return interp_trans, interp_ori
+
+    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_pos = 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_pos < 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():
+            self.manip_list.pop(0)
+        return len(self.manip_list) == 0
+
+    def is_success(self, t_eps=5e-3, o_eps=0.087):
+        p_base_ee_cur, q_base_ee_cur = self.controller.get_ee_pose()
+        diff_pos = np.linalg.norm(p_base_ee_cur - self.p_base_ee_tgt)
+        diff_ori = 2 * np.arccos(min(abs(np.dot(q_base_ee_cur, self.q_base_ee_tgt)), 1.0))
+        pose_flag = np.logical_or(
+            diff_pos < t_eps,
+            diff_ori < o_eps,
+        )
+        return pose_flag

workflows/simbox/core/skills/gripper_action.py

@@ -0,0 +1,61 @@
+import numpy as np
+from core.skills.base_skill import BaseSkill, register_skill
+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
+
+
+# pylint: disable=unused-argument
+@register_skill
+class Gripper_Action(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
+        self._gripper_state = self.skill_cfg["gripper_state"]
+
+        # !!! keyposes should be generated after previous skill is done
+        self.manip_list = []
+
+    def simple_generate_manip_cmds(self):
+        manip_list = []
+
+        p_base_ee_cur, q_base_ee_cur = self.controller.get_ee_pose()
+        if self._gripper_state == 1:  # Open
+            cmd = (p_base_ee_cur, q_base_ee_cur, "open_gripper", {}, self.skill_cfg.get("vel", None))
+        elif self._gripper_state == -1:  # Close
+            cmd = (p_base_ee_cur, q_base_ee_cur, "close_gripper", {}, self.skill_cfg.get("vel", None))
+        else:
+            raise NotImplementedError
+
+        manip_list.extend([cmd] * self.skill_cfg.get("wait_steps", 10))
+        self.manip_list = manip_list
+
+    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():
+            self.manip_list.pop(0)
+        return len(self.manip_list) == 0
+
+    def is_success(self):
+        return True

workflows/simbox/core/skills/heuristic_skill.py

@@ -0,0 +1,171 @@
+import numpy as np
+from core.skills.base_skill import BaseSkill, register_skill
+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.transformations import (
+    pose_from_tf_matrix,
+    tf_matrix_from_pose,
+)
+
+
+# pylint: disable=unused-argument
+@register_skill
+class Heuristic_Skill(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
+
+        self.lr_hand = "right" if "right" in self.controller.robot_file else "left"
+        if self.lr_hand == "left":
+            self._joint_indices = self.robot.left_joint_indices
+            self._joint_home = self.robot.left_joint_home
+            self._joint_home = np.array(self._joint_home)
+            if self.skill_cfg.get("gripper_state", None):
+                self._gripper_state = self.skill_cfg["gripper_state"]
+            else:
+                self._gripper_state = self.robot.left_gripper_state
+        else:
+            self._joint_indices = self.robot.right_joint_indices
+            self._joint_home = self.robot.right_joint_home
+            self._joint_home = np.array(self._joint_home)
+            if self.skill_cfg.get("gripper_state", None):
+                self._gripper_state = self.skill_cfg["gripper_state"]
+            else:
+                self._gripper_state = self.robot.right_gripper_state
+
+        ALLOWED_MODES = {"abs_qpos", "rel_qpos", "rel_ee", "home"}
+
+        self.mode = self.skill_cfg.get("mode", "home")
+        if self.mode not in ALLOWED_MODES:
+            raise ValueError(
+                f"Unsupported mode '{self.mode}' for JointMove. Allowed modes are: {sorted(ALLOWED_MODES)}"
+            )
+        self.move_steps = self.skill_cfg.get("move_steps", 50)
+        self.t_eps = self.skill_cfg.get("t_eps", 0.088)
+
+        # Keyposes should be generated after previous skill is done
+        self.manip_list = []
+        self._goal_joints = None
+
+    def _compute_ee_goal(self, p_base_ee_cur, q_base_ee_cur, rel_ee):
+        """
+        rel_ee: (4,4) transformation matrix
+        """
+        T_base_ee = tf_matrix_from_pose(p_base_ee_cur, q_base_ee_cur)
+        if isinstance(rel_ee, (list, tuple)):
+            rel_ee = np.array(rel_ee)
+        T_base_ee_tgt = rel_ee @ T_base_ee
+        p_base_ee_tgt, q_base_ee_tgt = pose_from_tf_matrix(T_base_ee_tgt)
+        return p_base_ee_tgt, q_base_ee_tgt
+
+    def _solve_goal_joints_via_plan(self, ee_trans_goal, ee_ori_goal):
+        """
+        Use controller.plan to get a collision-free joint path,
+        and take the last waypoint as goal arm joints.
+        """
+        if self.controller.use_batch:
+            raise NotImplementedError
+
+        sim_js = self.robot.get_joints_state()
+        js_names = self.robot.dof_names
+        result = self.controller.plan(ee_trans_goal, ee_ori_goal, sim_js, js_names)
+        succ = result.success.item()
+        if succ:
+            cmd_plan = result.get_interpolated_plan()
+            goal_arm_joints = cmd_plan[-1].position.cpu().numpy()  # replace by ik
+            return goal_arm_joints
+        else:
+            return None
+
+    def _build_joint_traj(self, curr_joints, goal_joints, p_base_ee_cur, q_base_ee_cur):
+        """Build a list of dummy_forward commands interpolating in joint space."""
+        manip_list = []
+        for k in range(self.move_steps):
+            alpha = float(k + 1) / float(self.move_steps) * 1.25
+            arm_action = goal_joints * alpha + curr_joints * (1.0 - alpha)
+            cmd = (
+                p_base_ee_cur,
+                q_base_ee_cur,
+                "dummy_forward",
+                {
+                    "arm_action": arm_action,
+                    "gripper_state": self._gripper_state,
+                },
+            )
+            manip_list.append(cmd)
+        return manip_list
+
+    def simple_generate_manip_cmds(self):
+        self.manip_list = []
+        p_base_ee_cur, q_base_ee_cur = self.controller.get_ee_pose()
+        curr_joints = self.robot.get_joint_positions()[self._joint_indices]
+
+        if self.mode == "home":
+            self._goal_joints = self._joint_home.copy()
+        else:
+            if self.mode == "abs_qpos":
+                self._goal_joints = self.skill_cfg.get("value", self._joint_home)
+            elif self.mode == "rel_qpos":
+                self._goal_joints = self.skill_cfg.get("value", np.zeros(self._joint_home.shape))
+            elif self.mode == "rel_ee":
+                p_base_ee_tgt, q_base_ee_tgt = self._compute_ee_goal(
+                    p_base_ee_cur, q_base_ee_cur, self.skill_cfg.get("value", np.eye(4))
+                )
+                self._goal_joints = self._solve_goal_joints_via_plan(p_base_ee_tgt, q_base_ee_tgt)
+            else:
+                raise NotImplementedError
+
+        if self._goal_joints is None:
+            self.manip_list = []
+            cmd = (
+                p_base_ee_cur,
+                q_base_ee_cur,
+                "update_specific",
+                {
+                    "ignore_substring": self.controller.ignore_substring,
+                    "reference_prim_path": self.controller.reference_prim_path,
+                },
+            )
+            self.manip_list.append(cmd)
+            return
+
+        self.manip_list = self._build_joint_traj(curr_joints, self._goal_joints, p_base_ee_cur, q_base_ee_cur)
+
+    def is_feasible(self, th=5):
+        return self.controller.num_plan_failed <= th
+
+    def is_subtask_done(self, t_eps=0.088):
+        if len(self.manip_list) == 0:
+            return True
+        if self._goal_joints is None:
+            return True
+        curr_joints = self.robot.get_joint_positions()[self._joint_indices]
+        target_joints = self.manip_list[0][3]["arm_action"]
+        diff_trans = np.linalg.norm(curr_joints - target_joints)
+        pose_flag = diff_trans < t_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.t_eps):
+            self.manip_list.pop(0)
+        if self.is_success(t_eps=self.t_eps):
+            self.manip_list.clear()
+            print("Heuristic Skill Done")
+        return len(self.manip_list) == 0
+
+    def is_success(self, t_eps=0.088):
+        if self._goal_joints is None:
+            print("cannot compute goal joints, skill failure")
+            return False
+
+        curr_joints = self.robot.get_joint_positions()[self._joint_indices]
+        diff_trans = np.linalg.norm(curr_joints - self._goal_joints)
+        return diff_trans < t_eps

workflows/simbox/core/skills/home.py

@@ -0,0 +1,82 @@
+import numpy as np
+from core.skills.base_skill import BaseSkill, register_skill
+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
+
+
+# pylint: disable=unused-argument
+@register_skill
+class Home(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
+
+        self.lr_hand = "right" if "right" in self.controller.robot_file else "left"
+        if self.lr_hand == "left":
+            self._joint_indices = self.robot.left_joint_indices
+            self._joint_home = self.robot.left_joint_home
+            if self.skill_cfg.get("gripper_state", None):
+                self._gripper_state = self.skill_cfg["gripper_state"]
+            else:
+                self._gripper_state = self.robot.left_gripper_state
+        elif self.lr_hand == "right":
+            self._joint_indices = self.robot.right_joint_indices
+            self._joint_home = self.robot.right_joint_home
+            if self.skill_cfg.get("gripper_state", None):
+                self._gripper_state = self.skill_cfg["gripper_state"]
+            else:
+                self._gripper_state = self.robot.right_gripper_state
+
+        # !!! keyposes should be generated after previous skill is done
+        self.manip_list = []
+
+    def simple_generate_manip_cmds(self):
+        manip_list = []
+        curr_ee_trans, curr_ee_ori = self.controller.get_ee_pose()
+        curr_joints = self.robot.get_joint_positions()[self._joint_indices]
+        home_joints = self._joint_home
+
+        for k in range(0, 50):
+            arm_action = np.array(home_joints) * ((k + 1) / 40) + np.array(curr_joints) * (1 - (k + 1) / 40)
+            cmd = (
+                curr_ee_trans,
+                curr_ee_ori,
+                "dummy_forward",
+                {"arm_action": arm_action, "gripper_state": self._gripper_state},
+            )
+            manip_list.append(cmd)
+
+        self.manip_list = manip_list
+
+    def is_feasible(self, th=5):
+        return self.controller.num_plan_failed <= th
+
+    def is_subtask_done(self, t_eps=0.088):
+        assert len(self.manip_list) != 0
+        curr_joints = self.robot.get_joint_positions()[self._joint_indices]
+        target_joints = self.manip_list[0][3]["arm_action"]
+        diff_trans = np.linalg.norm(curr_joints - target_joints)
+        pose_flag = (diff_trans < t_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():
+            self.manip_list.pop(0)
+        if self.is_success():
+            self.manip_list.clear()
+            print("Home Done")
+        return len(self.manip_list) == 0
+
+    def is_success(self, t_eps=0.088):
+        curr_joints = self.robot.get_joint_positions()[self._joint_indices]
+        diff_trans = np.linalg.norm(curr_joints - self._joint_home)
+        # print(diff_trans)
+        return diff_trans < t_eps

workflows/simbox/core/skills/joint_ctrl.py

@@ -0,0 +1,149 @@
+# pylint: skip-file
+import numpy as np
+import torch
+from core.skills.base_skill import BaseSkill, register_skill
+from core.utils.interpolate_utils import linear_interpolation
+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
+
+
+# pylint: disable=unused-argument
+@register_skill
+class Joint_Ctrl(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.skill_cfg = cfg
+        self.robot_base_path = self.controller.robot_base_path
+        if "left" in controller.robot_file:
+            self.robot_lr = "left"
+        elif "right" in controller.robot_file:
+            self.robot_lr = "right"
+        self.manip_list = []
+        self.success_threshold_js = self.skill_cfg.get("success_threshold_js", 5e-3)
+
+    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': [0,0,0,0,0,0]})
+        cmd = (p_base_ee_cur, q_base_ee_cur, "update_pose_cost_metric", {"hold_vec_weight": None})
+        manip_list.append(cmd)
+
+        curr_js, target_js = self.get_target_js()
+        interp_js_list = linear_interpolation(curr_js, target_js, self.skill_cfg.get("num_steps", 10))
+        for js in interp_js_list:
+            p_base_ee, q_base_ee = self.controller.forward_kinematic(js)
+            cmd = (
+                p_base_ee,
+                q_base_ee,
+                "dummy_forward",
+                {
+                    "arm_action": js,
+                    "gripper_state": self.skill_cfg.get("gripper_state", 1.0),
+                },
+            )
+            manip_list.append(cmd)
+
+        self.target_js = js
+        self.manip_list = manip_list
+
+    def get_target_js(self):
+        """
+        Compute target joint configuration based on current joint state and
+        joint control commands defined in skill configuration.
+
+        Returns:
+            curr_js (np.ndarray): Current joint positions of the controlled arm.
+            target_js (np.ndarray): Target joint positions after applying commands.
+        """
+
+        # --- Get current joint positions ---
+        joint_positions = self.robot.get_joints_state().positions
+
+        if isinstance(joint_positions, torch.Tensor):
+            curr_js = joint_positions.detach().cpu().numpy()[self.controller.arm_indices]
+        elif isinstance(joint_positions, np.ndarray):
+            curr_js = joint_positions[self.controller.arm_indices]
+        else:
+            raise TypeError(f"Unsupported joint state type: {type(joint_positions)}")
+
+        target_js = curr_js.copy()
+
+        # --- Apply joint control commands ---
+        # ctrl_list: list of (joint_index, angle_in_deg, mode), mode in {"abs", "delta"}
+        ctrl_list = self.skill_cfg.get("ctrl_list", [])
+        for joint_idx, angle_deg, mode in ctrl_list:
+            angle_rad = angle_deg * np.pi / 180.0
+            if mode == "abs":
+                target_js[joint_idx] = angle_rad
+            elif mode == "delta":
+                target_js[joint_idx] += angle_rad
+            else:
+                raise ValueError(f"Unknown control mode: {mode}")
+
+        # --- Apply robot-specific joint limits / safety clamps ---
+        robot_file = self.controller.robot_file.lower()
+
+        if "piper" in robot_file:
+            # Example: clamp elbow and wrist joints for Piper robot
+            target_js[2] = min(target_js[2], 0.0)
+            target_js[4] = np.clip(target_js[4], -1.22, 1.22)
+
+        elif "r5a" in robot_file:
+            # Reserved for R5A-specific constraints
+            pass
+
+        return curr_js, target_js
+
+    def is_feasible(self, th=5):
+        return self.controller.num_plan_failed <= th
+
+    def is_subtask_done(self, js_eps=5e-3, t_eps=1e-3, o_eps=5e-3):
+        assert len(self.manip_list) != 0
+        manip_cmd = self.manip_list[0]
+        if manip_cmd[2] == "joint_ctrl":
+            joint_positions = self.robot.get_joints_state().positions
+            if isinstance(joint_positions, torch.Tensor):
+                curr_js = joint_positions.numpy()[self.controller.arm_indices]  # JointState
+            elif isinstance(joint_positions, np.ndarray):
+                curr_js = joint_positions[self.controller.arm_indices]  # JointState
+            target_js = self.manip_list[0][3]["target"]
+            diff_js = np.linalg.norm(curr_js - target_js)
+            js_flag = diff_js < js_eps
+            self.plan_flag = self.controller.num_last_cmd > 10
+            return np.logical_or(js_flag, self.plan_flag)
+        else:
+            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(self.success_threshold_js):
+            self.manip_list.pop(0)
+
+        return len(self.manip_list) == 0
+
+    def is_success(self):
+        joint_positions = self.robot.get_joints_state().positions
+        if isinstance(joint_positions, torch.Tensor):
+            curr_js = joint_positions.numpy()[self.controller.arm_indices]  # JointState
+        elif isinstance(joint_positions, np.ndarray):
+            curr_js = joint_positions[self.controller.arm_indices]  # JointState
+        distance_js = np.linalg.norm(curr_js - self.target_js)
+        flag = (distance_js < self.success_threshold_js) and (len(self.manip_list) == 0)
+
+        return flag

workflows/simbox/core/skills/manualpick.py

@@ -0,0 +1,422 @@
+# pylint: skip-file
+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,
+)
+
+
+@register_skill
+class Manualpick(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)
+        grasp_scale = self.skill_cfg.get("grasp_scale", 1)
+        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=grasp_scale
+        )
+
+        # !!! keyposes should be generated after previous skill is done
+        self.manip_list = []
+        self.pickcontact_view = task.pickcontact_views[robot.name][lr_arm][object_name]
+
+    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)
+        ignore_substring.append(self.pick_obj.name)
+
+        cmd = (
+            p_base_ee_cur,
+            q_base_ee_cur,
+            "update_pose_cost_metric",
+            {"hold_vec_weight": self.skill_cfg.get("hold_vec_weight", None)},
+        )
+        manip_list.append(cmd)
+
+        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("start_lr_skill", False):
+            cmd = (p_base_ee_cur, q_base_ee_cur, "update_pose_cost_metric", {"hold_vec_weight": None})
+            manip_list.append(cmd)
+
+        # Pre grasp
+        T_base_ee_grasps = self.sample_ee_pose()  # (N, 4, 4)
+        adjust_ori = self.skill_cfg.get("adjust_ori", None)
+        if adjust_ori:
+            pose_axis = adjust_ori[0]
+            base_axis = adjust_ori[1]
+            judge_flag = adjust_ori[2]
+            axis_index = {"x": 0, "y": 1, "z": 2}
+            rotate_axis = self.skill_cfg.get("adjust_rotate_axis", "x")
+
+            if "piper" in self.controller.robot_file or "r5a" in self.controller.robot_file:
+                num_poses = T_base_ee_grasps.shape[0]
+                adjust_angle_list_cfg = self.skill_cfg.get("adjust_angle_list_cfg", [-15, 15, 7])
+                adjust_angle_list = np.linspace(
+                    adjust_angle_list_cfg[0], adjust_angle_list_cfg[1], adjust_angle_list_cfg[2]
+                )  # (K,)
+
+                # build batch rotation matrices of shape (K, 4, 4)
+                thetas = np.radians(adjust_angle_list)
+                rots = []
+                for theta in thetas:
+                    if rotate_axis == "x":
+                        rot = np.array(
+                            [
+                                [1, 0, 0, 0],
+                                [0, np.cos(theta), -np.sin(theta), 0],
+                                [0, np.sin(theta), np.cos(theta), 0],
+                                [0, 0, 0, 1],
+                            ]
+                        )
+                    elif rotate_axis == "y":
+                        rot = np.array(
+                            [
+                                [np.cos(theta), 0, np.sin(theta), 0],
+                                [0, 1, 0, 0],
+                                [-np.sin(theta), 0, np.cos(theta), 0],
+                                [0, 0, 0, 1],
+                            ]
+                        )
+                    elif rotate_axis == "z":
+                        rot = np.array(
+                            [
+                                [np.cos(theta), -np.sin(theta), 0, 0],
+                                [np.sin(theta), np.cos(theta), 0, 0],
+                                [0, 0, 1, 0],
+                                [0, 0, 0, 1],
+                            ]
+                        )
+                    else:
+                        rot = np.eye(4)
+                    rots.append(rot)
+                rots = np.stack(rots, axis=0)  # (K, 4, 4)
+
+                # original poses: (N, 4, 4), broadcast with rotations: (K, 4, 4)
+                original_poses = T_base_ee_grasps.copy()
+                rotated_poses = original_poses[:, None, :, :] @ rots[None, :, :, :]  # (N, K, 4, 4)
+
+                # compute metric for each (pose, angle) candidate
+                base_idx = axis_index[base_axis]
+                pose_idx = axis_index[pose_axis]
+                current_values = rotated_poses[:, :, base_idx, pose_idx]  # (N, K)
+
+                if judge_flag == "min":
+                    best_indices = np.argmin(current_values, axis=1)
+                else:
+                    best_indices = np.argmax(current_values, axis=1)
+
+                # gather best poses per grasp
+                idx_rows = np.arange(num_poses)
+                best_poses = rotated_poses[idx_rows, best_indices]  # (N, 4, 4)
+                T_base_ee_grasps = best_poses
+
+        manual_adjust_ori = self.skill_cfg.get("manual_adjust_ori", None)
+        if manual_adjust_ori:
+            for adjust_ori in manual_adjust_ori:
+                rotate_axis = adjust_ori[0]
+                angle = adjust_ori[1]
+                theta = np.radians(angle)
+                if rotate_axis == "x":
+                    rot = np.array(
+                        [
+                            [1, 0, 0, 0],
+                            [0, np.cos(theta), -np.sin(theta), 0],
+                            [0, np.sin(theta), np.cos(theta), 0],
+                            [0, 0, 0, 1],
+                        ]
+                    )
+                elif rotate_axis == "y":
+                    rot = np.array(
+                        [
+                            [np.cos(theta), 0, np.sin(theta), 0],
+                            [0, 1, 0, 0],
+                            [-np.sin(theta), 0, np.cos(theta), 0],
+                            [0, 0, 0, 1],
+                        ]
+                    )
+                elif rotate_axis == "z":
+                    rot = np.array(
+                        [
+                            [np.cos(theta), -np.sin(theta), 0, 0],
+                            [np.sin(theta), np.cos(theta), 0, 0],
+                            [0, 0, 1, 0],
+                            [0, 0, 0, 1],
+                        ]
+                    )
+                else:
+                    rot = np.eye(4)
+                # apply the same rotation to all grasps in batch
+                T_base_ee_grasps = T_base_ee_grasps @ rot
+
+        adjust_trans_offset = self.skill_cfg.get("adjust_trans_offset", [0, 0, 0])
+        T_base_ee_grasps[:, :3, 3] += adjust_trans_offset
+        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)
+        pre_grasp_offset_manual = self.skill_cfg.get("pre_grasp_offset_manual", None)
+        if pre_grasp_offset_manual:
+            T_base_ee_pregrasps[:, :3, 3] += np.array(pre_grasp_offset_manual)
+
+        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
+
+        cmd = (p_base_ee_pregrasps[index], q_base_ee_pregrasps[index], "open_gripper", {})
+        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], "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 = 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)
+
+        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 Exception(
+                "poses in {} frame is not supported: accepted values are [world, body, armbase] only".format(frame)
+            )
+
+        if frame == "body":
+            return self.T_obj_ee
+
+        T_world_obj = tf_matrix_from_pose(*self.pick_obj.get_local_pose())
+        T_world_ee = T_world_obj[None] @ self.T_obj_ee
+
+        if frame == "world":
+            return T_world_ee
+
+        if frame == "armbase":  # robot base frame
+            T_r2w = get_relative_transform(
+                get_prim_at_path(self.controller.reference_prim_path), get_prim_at_path(self.task.root_prim_path)
+            )
+            T_base_world = np.linalg.inv(T_r2w)
+            T_base_ee = T_base_world[None] @ T_world_ee
+            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():
+            self.manip_list.pop(0)
+        return len(self.manip_list) == 0
+
+    def is_success(self):
+        contact, indices = self.get_contact()
+        return len(indices) >= 1

workflows/simbox/core/skills/mobile_rotate.py

@@ -0,0 +1,66 @@
+import numpy as np
+from core.skills.base_skill import BaseSkill, register_skill
+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
+
+
+# pylint: disable=unused-argument
+@register_skill
+class Mobile_Rotate(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
+        # !!! keyposes should be generated after previous skill is done
+        self.manip_list = []
+        self.target = self.skill_cfg.get("target", 0.785)  # 45 degrees
+        self.mobile_rotate_indices = [self.robot._articulation_view.dof_names.index("mobile_rotate")]
+
+    def simple_generate_manip_cmds(self):
+        manip_list = []
+        self.mobile_rotate_initial = self.robot._articulation_view.get_joint_positions()[0, self.mobile_rotate_indices]
+        curr_ee_trans, curr_ee_ori = self.controller.get_ee_pose()
+        for k in range(100):
+            cmd = (
+                curr_ee_trans,
+                curr_ee_ori,
+                "mobile_move",
+                {
+                    "target": self.target * k / 50,
+                    "joint_indices": self.mobile_rotate_indices,
+                    "initial_position": self.mobile_rotate_initial,
+                },
+            )
+            manip_list.append(cmd)
+        self.manip_list = manip_list
+
+    def is_subtask_done(self, o_eps=5e-3):
+        assert len(self.manip_list) != 0
+        manip_cmd = self.manip_list[0]
+        target = manip_cmd[3]["target"]
+        curr_joint_p = self.robot._articulation_view.get_joint_positions()[0, self.mobile_rotate_indices]
+        distance = np.abs(curr_joint_p - self.mobile_rotate_initial)
+        print(distance)
+        return abs(distance - abs(target)) < o_eps
+
+    def is_success(self):
+        return True
+
+    def mobile_rotate_done(self):
+        curr_joint_p = self.robot._articulation_view.get_joint_positions()[0, self.mobile_rotate_indices]
+        distance = np.abs(curr_joint_p - self.mobile_rotate_initial)
+        return distance >= np.abs(self.target)
+
+    def is_done(self):
+        if len(self.manip_list) == 0:
+            return True
+        if self.is_subtask_done():
+            self.manip_list.pop(0)
+        if self.mobile_rotate_done():
+            self.manip_list.clear()
+            print("Mobile Rotate Done")
+        return len(self.manip_list) == 0

workflows/simbox/core/skills/mobile_translate.py

@@ -0,0 +1,76 @@
+import numpy as np
+from core.skills.base_skill import BaseSkill, register_skill
+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
+
+
+# pylint: disable=unused-argument
+@register_skill
+class Mobile_Translate(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
+        # !!! keyposes should be generated after previous skill is done
+        self.manip_list = []
+
+        self.mobile_translate_indices = [
+            self.robot._articulation_view.dof_names.index(name) for name in ["mobile_translate_x", "mobile_translate_y"]
+        ]
+        self.target = np.array(self.skill_cfg.get("target", [0.0, 0.0]))
+
+    def simple_generate_manip_cmds(self):
+        manip_list = []
+        curr_ee_trans, curr_ee_ori = self.controller.get_ee_pose()
+        self.mobile_translate_initial = self.robot._articulation_view.get_joint_positions()[
+            0, self.mobile_translate_indices
+        ]
+        for k in range(100):
+            cmd = (
+                curr_ee_trans,
+                curr_ee_ori,
+                "mobile_move",
+                {
+                    "target": self.target * k / 50,
+                    "joint_indices": self.mobile_translate_indices,
+                    "initial_position": self.mobile_translate_initial,
+                },
+            )
+            manip_list.append(cmd)
+
+        self.manip_list = manip_list
+
+    def is_subtask_done(self, t_eps=1e-3):
+        assert len(self.manip_list) != 0
+        manip_cmd = self.manip_list[0]
+        target = manip_cmd[3]["target"]
+        curr_joint_p = self.robot._articulation_view.get_joint_positions()[0, self.mobile_translate_indices]
+        distance = np.abs(curr_joint_p - self.mobile_translate_initial)
+
+        return np.linalg.norm(distance - np.abs(target)) < t_eps
+
+    def mobile_translate_done(self):
+        curr_joint_p = self.robot._articulation_view.get_joint_positions()[0, self.mobile_translate_indices]
+        distance = np.abs(curr_joint_p - self.mobile_translate_initial)
+
+        return np.all(distance > np.abs(self.target))
+
+    def is_success(self):
+        return True
+
+    def is_done(self):
+        if len(self.manip_list) == 0:
+            return True
+        if self.is_subtask_done():
+            self.manip_list.pop(0)
+        if self.mobile_translate_done():
+            self.manip_list.clear()
+            print("Mobile Translate Done")
+        return len(self.manip_list) == 0
+
+    def is_feasible(self, th=5):
+        return self.controller.num_plan_failed <= th

workflows/simbox/core/skills/move.py

@@ -0,0 +1,116 @@
+import numpy as np
+from core.skills.base_skill import BaseSkill, register_skill
+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 scipy.spatial.transform import Rotation as R
+
+
+# pylint: disable=unused-argument
+@register_skill
+class Move(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.move_obj = task.objects[cfg["objects"][0]]
+        self.tgt_obj = task.objects[cfg["objects"][1]]
+
+        if "invisible_object" in cfg:
+            self.ivs_obj = task.objects[cfg["invisible_object"][0]]
+            self.ivs_obj.set_visibility(True)
+        else:
+            self.ivs_obj = None
+        self.skill_cfg = cfg
+        self.success_threshold = cfg["success_threshold"]
+        self.delta_trans = np.array(cfg.get("delta_trans", [[0, 0, 0]]))
+        self.manip_list = []
+
+    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": self.skill_cfg.get("hold_vec_weight", [0, 0, 0, 0, 0, 0])},
+        )
+        manip_list.append(cmd)
+
+        ignore_substring = 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)
+
+        p_base_ee_tgt = self.getTgtTranslation()
+        for delta_trans in self.delta_trans:
+            cmd = (p_base_ee_tgt + delta_trans, q_base_ee_cur, "close_gripper", {})
+            manip_list.append(cmd)
+
+        self.manip_list = manip_list
+        self.p_base_ee_tgt = p_base_ee_tgt + self.delta_trans[-1]
+
+    def getTgtTranslation(self):
+        p_world_move_obj = self.move_obj.get_world_pose()[0]
+        p_world_tgt_obj = self.tgt_obj.get_world_pose()[0]
+        global_move = p_world_tgt_obj - p_world_move_obj
+        _, q_world_base_cur = self.controller.get_armbase_pose()
+        p_base_ee_cur, _ = self.controller.get_ee_pose()
+        R_we = R.from_quat(q_world_base_cur, scalar_first=True).as_matrix()  # EE -> World
+        R_ew = R_we.T  # World -> EE
+        ee_move = R_ew @ global_move
+        ee_move[2] *= 0.0
+
+        p_base_ee_tgt = p_base_ee_cur + ee_move
+        return p_base_ee_tgt
+
+    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
+        # print(self.controller.num_last_cmd)
+        if self.plan_flag:
+            print(f"move_only plan_flag: {self.plan_flag}, num_last_cmd: {self.controller.num_last_cmd}")
+        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():
+            self.manip_list.pop(0)
+        if self.ivs_obj is not None:
+            if self.is_success():
+                self.ivs_obj.set_visibility(False)
+        return len(self.manip_list) == 0
+
+    def is_success(self):
+        p_base_ee_cur, _ = self.controller.get_ee_pose()
+        distance = np.linalg.norm(p_base_ee_cur - self.p_base_ee_tgt)
+        flag = (distance < self.success_threshold) and (len(self.manip_list) == 0)
+
+        # print("distance ee :", distance)
+
+        p_world_move_obj = self.move_obj.get_world_pose()[0]
+        p_world_tgt_obj = self.tgt_obj.get_world_pose()[0]
+        distance = np.linalg.norm(p_world_tgt_obj - p_world_move_obj)
+        flag = (distance < self.success_threshold) and flag
+
+        # print("distance move :", distance)
+
+        return flag

workflows/simbox/core/skills/open.py

@@ -0,0 +1,214 @@
+from copy import deepcopy
+
+import numpy as np
+from core.skills.base_skill import BaseSkill, register_skill
+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
+from scipy.spatial.transform import Rotation as R
+from solver.planner import KPAMPlanner
+
+
+# pylint: disable=unused-argument
+@register_skill
+class Open(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.stage = task.stage
+        self.name = cfg["name"]
+        self.skill_cfg = cfg
+        art_obj_name = cfg["objects"][0]
+        self.art_obj = task.objects[art_obj_name]
+        self.planner_setting = cfg["planner_setting"]
+        self.contact_pose_index = self.planner_setting["contact_pose_index"]
+        self.success_threshold = self.planner_setting["success_threshold"]
+        self.success_mode = self.planner_setting.get("success_mode", "abs")
+        self.update_art_joint = self.planner_setting.get("update_art_joint", False)
+        if kwargs:
+            self.world = kwargs["world"]
+            self.draw = kwargs["draw"]
+        self.manip_list = []
+
+        lr_arm = "left" if "left" in self.controller.robot_file else "right"
+        self.fingers_link_contact_view = task.artcontact_views[robot.name][lr_arm][art_obj_name + "_fingers_link"]
+        self.fingers_base_contact_view = task.artcontact_views[robot.name][lr_arm][art_obj_name + "_fingers_base"]
+        self.forbid_collision_contact_view = task.artcontact_views[robot.name][lr_arm][
+            art_obj_name + "_forbid_collision"
+        ]
+        self.collision_valid = True
+        self.process_valid = True
+
+    def setup_kpam(self):
+        self.planner = KPAMPlanner(
+            env=self.world,
+            robot=self.robot,
+            object=self.art_obj,
+            cfg_path=self.planner_setting,
+            controller=self.controller,
+            draw_points=self.draw,
+            stage=self.stage,
+        )
+
+    def simple_generate_manip_cmds(self):
+        if self.skill_cfg.get("obj_info_path", None):
+            self.art_obj.update_articulated_info(self.skill_cfg["obj_info_path"])
+
+        self.setup_kpam()
+
+        traj_keyframes, sample_times = self.planner.get_keypose()
+        if len(traj_keyframes) == 0 and len(sample_times) == 0:
+            print("No keyframes found, return empty manip_list")
+            self.manip_list = []
+            return
+
+        T_world_base = get_relative_transform(
+            get_prim_at_path(self.robot.base_path), get_prim_at_path(self.task.root_prim_path)
+        )
+        self.traj_keyframes = traj_keyframes
+        self.sample_times = sample_times
+        if self.draw:
+            for keypose in traj_keyframes:
+                self.draw.draw_points([(T_world_base @ np.append(keypose[:3, 3], 1))[:3]], [(0, 0, 0, 1)], [7])
+        manip_list = []
+
+        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", []))
+        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)
+
+        for i in range(len(self.traj_keyframes)):
+            p_base_ee_tgt = self.traj_keyframes[i][:3, 3]
+            q_base_ee_tgt = R.from_matrix(self.traj_keyframes[i][:3, :3]).as_quat(scalar_first=True)
+            if i <= self.contact_pose_index:
+                cmd = (p_base_ee_tgt, q_base_ee_tgt, "open_gripper", {})
+            else:
+                cmd = (p_base_ee_tgt, q_base_ee_tgt, "close_gripper", {})
+            manip_list.append(cmd)
+
+            if i == self.contact_pose_index:
+                cmd = (p_base_ee_tgt, q_base_ee_tgt, "close_gripper", {})
+                manip_list.extend([cmd] * 40)
+
+            if i == self.contact_pose_index - 1:
+                p_base_ee = self.traj_keyframes[i][:3, 3]
+                q_base_ee = R.from_matrix(self.traj_keyframes[i][:3, :3]).as_quat(scalar_first=True)
+                ignore_substring = deepcopy(self.controller.ignore_substring)
+                parent_name = self.art_obj.prim_path.split("/")[-2]
+                ignore_substring.append(parent_name)
+                cmd = (
+                    p_base_ee,
+                    q_base_ee,
+                    "update_specific",
+                    {"ignore_substring": ignore_substring, "reference_prim_path": self.controller.reference_prim_path},
+                )
+                manip_list.append(cmd)
+        self.manip_list = manip_list
+
+    def update(self):
+        curr_joint_p = self.art_obj._articulation_view.get_joint_positions()[:, self.art_obj.object_joint_index]
+        if self.update_art_joint:
+            self.art_obj._articulation_view.set_joint_position_targets(
+                positions=curr_joint_p, joint_indices=self.art_obj.object_joint_index
+            )
+
+    def get_contact(self, contact_threshold=0.0):
+        contact = {}
+        fingers_link_contact = np.abs(self.fingers_link_contact_view.get_contact_force_matrix()).squeeze()
+        fingers_link_contact = np.sum(fingers_link_contact, axis=-1)
+        fingers_link_contact_indices = np.where(fingers_link_contact > contact_threshold)[0]
+        contact["fingers_link"] = {
+            "fingers_link_contact": fingers_link_contact,
+            "fingers_link_contact_indices": fingers_link_contact_indices,
+        }
+
+        fingers_base_contact = np.abs(self.fingers_base_contact_view.get_contact_force_matrix()).squeeze()
+        fingers_base_contact = np.sum(fingers_base_contact, axis=-1)
+        fingers_base_contact_indices = np.where(fingers_base_contact > contact_threshold)[0]
+        contact["fingers_base"] = {
+            "fingers_base_contact": fingers_base_contact,
+            "fingers_base_contact_indices": fingers_base_contact_indices,
+        }
+
+        forbid_collision_contact = np.abs(self.forbid_collision_contact_view.get_contact_force_matrix()).squeeze()
+        forbid_collision_contact = np.sum(forbid_collision_contact, axis=-1)
+        forbid_collision_contact_indices = np.where(forbid_collision_contact > contact_threshold)[0]
+        contact["forbid_collision"] = {
+            "forbid_collision_contact": forbid_collision_contact,
+            "forbid_collision_contact_indices": forbid_collision_contact_indices,
+        }
+
+        return contact
+
+    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():
+            self.manip_list.pop(0)
+            print("POP one manip cmd")
+        if self.is_success():
+            self.manip_list.clear()
+            print("Open Done")
+        return len(self.manip_list) == 0
+
+    def is_success(self):
+        contact = self.get_contact()
+
+        if self.skill_cfg.get("collision_valid", True):
+            self.collision_valid = (
+                self.collision_valid
+                and len(contact["forbid_collision"]["forbid_collision_contact_indices"]) == 0
+                and len(contact["fingers_base"]["fingers_base_contact_indices"]) == 0
+            )
+        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.art_obj.get_joints_state().velocities)) < 5
+            )
+
+        curr_joint_p = self.art_obj._articulation_view.get_joint_positions()[:, self.art_obj.object_joint_index]
+        init_joint_p = self.art_obj.articulation_initial_joint_position
+        print(
+            curr_joint_p - init_joint_p,
+            "collision_valid :",
+            self.collision_valid,
+            "process_valid :",
+            self.process_valid,
+        )
+        if self.success_mode == "normal":
+            return (
+                (curr_joint_p - init_joint_p) >= np.abs(self.success_threshold)
+                and self.collision_valid
+                and self.process_valid
+            )
+        elif self.success_mode == "abs":
+            return (
+                np.abs(curr_joint_p - init_joint_p) >= np.abs(self.success_threshold)
+                and self.collision_valid
+                and self.process_valid
+            )

workflows/simbox/core/skills/pick.py

@@ -0,0 +1,347 @@
+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
+
+        T_world_obj = tf_matrix_from_pose(*self.pick_obj.get_local_pose())
+        T_world_ee = T_world_obj[None] @ self.T_obj_ee
+
+        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)
+            )
+            T_base_world = np.linalg.inv(T_world_base)
+            T_base_ee = T_base_world[None] @ T_world_ee
+            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

workflows/simbox/core/skills/place.py

@@ -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))

workflows/simbox/core/skills/pour_water_succ.py

@@ -0,0 +1,152 @@
+import numpy as np
+import torch
+from core.skills.base_skill import BaseSkill, register_skill
+from core.utils.transformation_utils import (
+    get_orientation,
+    perturb_orientation,
+    perturb_position,
+)
+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 scipy.spatial.transform import Rotation as R
+
+
+# pylint: disable=unused-argument
+@register_skill
+class Pour_Water_Succ(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
+        self.frame = self.skill_cfg.get("frame", "robot")
+        if self.skill_cfg.get("translation", None):
+            self.p_base_ee_tgt = np.array(self.skill_cfg.get("translation", None))
+            self.p_base_ee_tgt = perturb_position(self.p_base_ee_tgt, self.skill_cfg.get("max_noise_m", 0.05))
+        else:
+            self.p_base_ee_tgt = None
+        if self.skill_cfg.get("quaternion", None) or self.skill_cfg.get("euler", None):
+            self.q_base_ee_tgt = np.array(
+                get_orientation(self.skill_cfg.get("euler"), self.skill_cfg.get("quaternion"))
+            )
+            self.q_base_ee_tgt = perturb_orientation(self.q_base_ee_tgt, self.skill_cfg.get("max_noise_deg", 5))
+        else:
+            self.q_base_ee_tgt = None
+
+        # !!! keyposes should be generated after previous skill is done
+        self.manip_list = []
+
+    def simple_generate_manip_cmds(self):
+        manip_list = []
+        p_base_ee_cur, q_base_ee_cur = self.controller.get_ee_pose()
+        if self.p_base_ee_tgt is None:
+            self.p_base_ee_tgt = p_base_ee_cur
+        if self.q_base_ee_tgt is None:
+            self.q_base_ee_tgt = q_base_ee_cur
+        # --- Get current joint positions ---
+        joint_positions = self.robot.get_joints_state().positions
+
+        if isinstance(joint_positions, torch.Tensor):
+            curr_js = joint_positions.detach().cpu().numpy()[self.controller.arm_indices]
+        elif isinstance(joint_positions, np.ndarray):
+            curr_js = joint_positions[self.controller.arm_indices]
+        else:
+            raise TypeError(f"Unsupported joint state type: {type(joint_positions)}")
+        p_base_ee, q_base_ee = self.controller.forward_kinematic(curr_js)
+        cmd = (
+            p_base_ee,
+            q_base_ee,
+            "dummy_forward",
+            {
+                "arm_action": curr_js,
+                "gripper_state": self.skill_cfg.get("gripper_state", 1.0),
+            },
+        )
+        manip_list.append(cmd)
+
+        self.manip_list = manip_list
+
+    # def interp(self, curr_trans, curr_ori, target_trans, target_ori, interp_num, normalize_quaternions=True):
+    #     interp_trans = np.linspace(curr_trans, target_trans, interp_num, axis=0)
+    #     if normalize_quaternions:
+    #         curr_ori = curr_ori / np.linalg.norm(curr_ori)
+    #         target_ori = target_ori / np.linalg.norm(target_ori)
+    #     rotations = R.from_quat([curr_ori, target_ori], scalar_first=True)  # Current  # Target
+    #     slerp = Slerp([0, 1], rotations)
+    #     times = np.linspace(0, 1, interp_num)
+    #     interp_ori = slerp(times).as_quat(scalar_first=True)
+
+    #     return interp_trans, interp_ori
+
+    def is_feasible(self, th=5):
+        return self.controller.num_plan_failed <= th
+
+    def is_subtask_done(self, js_eps=5e-3, t_eps=1e-3, o_eps=5e-3):
+        assert len(self.manip_list) != 0
+        manip_cmd = self.manip_list[0]
+        if manip_cmd[2] == "joint_ctrl":
+            # --- Get current joint positions ---
+            joint_positions = self.robot.get_joints_state().positions
+
+            if isinstance(joint_positions, torch.Tensor):
+                curr_js = joint_positions.detach().cpu().numpy()[self.controller.arm_indices]
+            elif isinstance(joint_positions, np.ndarray):
+                curr_js = joint_positions[self.controller.arm_indices]
+            else:
+                raise TypeError(f"Unsupported joint state type: {type(joint_positions)}")
+            target_js = self.manip_list[0][3]["target"]
+            diff_js = np.linalg.norm(curr_js - target_js)
+            js_flag = diff_js < js_eps
+            self.plan_flag = self.controller.num_last_cmd > 10
+            return np.logical_or(js_flag, self.plan_flag)
+        else:
+            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():
+            self.manip_list.pop(0)
+        return len(self.manip_list) == 0
+
+    def is_success(self):
+        particle_poses = np.array(self.task.particles.GetPointsAttr().Get())
+        container_name = self.skill_cfg.get("container_name", "cup")
+        container = self.task.objects[container_name]
+        container_trans, _ = container.get_world_pose()
+        if isinstance(container_trans, torch.Tensor):
+            container_trans = container_trans.numpy()
+        container_radius = self.skill_cfg.get("container_radius", 0.025)
+        distances_xy = np.linalg.norm(particle_poses[:, :2] - container_trans[:2], axis=1)
+        in_container_mask = distances_xy < container_radius
+        particles_in_container = particle_poses[in_container_mask]
+        fluid_flag = (len(particles_in_container) > self.skill_cfg.get("particle_num_th_min", 50)) and (
+            len(particles_in_container) < self.skill_cfg.get("particle_num_th_max", 300)
+        )
+
+        succ_flag = fluid_flag
+        if self.skill_cfg.get("container_up", []):
+            container_up_list = self.skill_cfg.get("container_up")
+            for container_name, container_up, threshold in container_up_list:
+                container = self.task.objects[container_name]
+                _, container_ori = container.get_world_pose()
+
+                container_ori = R.from_quat(container_ori, scalar_first=True).as_matrix()
+                up_2_idx = {"x": 0, "y": 1, "z": 2}
+                container_flag = container_ori[2, up_2_idx[container_up]] > threshold
+                succ_flag = succ_flag and container_flag
+
+        print(f"[DEBUG] pour: fluid_flag {fluid_flag} | num_particles_in_container {len(particles_in_container)}")
+        return succ_flag

workflows/simbox/core/skills/rotate.py

@@ -0,0 +1,180 @@
+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
+from scipy.spatial.transform import Rotation as R
+from solver.planner import KPAMPlanner
+
+
+# pylint: disable=consider-using-generator,too-many-public-methods,unused-argument
+@register_skill
+class Rotate(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.stage = task.stage
+        self.name = cfg["name"]
+        art_obj_name = cfg["objects"][0]
+        self.art_obj = task.objects[art_obj_name]
+        self.cfg = cfg
+
+        # debug start: KPAMPlanner
+        self.planner_setting = OmegaConf.to_container(cfg["planner_setting"])
+        self.contact_pose_index = self.planner_setting.get("contact_pose_index")
+        self.success_threshold = self.planner_setting.get("success_threshold", 0.785)  # 45 degrees
+        if kwargs:
+            self.world = kwargs["world"]
+            self.draw = kwargs["draw"]
+        # debug end: KPAMPlanner
+        # !!! keyposes should be generated after previous skill is done
+        self.manip_list = []
+
+        if self.cfg.get("obj_info_path", None):
+            self.art_obj.update_articulated_info(self.cfg["obj_info_path"])
+
+    # debug start: KPAMPlanner
+    def setup_kpam(self):
+        self.planner = KPAMPlanner(
+            env=self.world,
+            robot=self.robot,
+            object=self.art_obj,
+            cfg_path=self.planner_setting,
+            controller=self.controller,
+            draw_points=self.draw,
+            stage=self.stage,
+        )
+        if "additional_labels" in self.planner_setting:
+            new_value = self.planner_setting["additional_labels"].get(
+                self.art_obj.asset_relative_path, self.planner.modify_actuation_motion
+            )
+            self.planner.modify_actuation_motion = new_value
+
+    def simple_generate_manip_cmds(self):
+        if self.cfg.get("obj_info_path", None):
+            self.art_obj.update_articulated_info(self.cfg["obj_info_path"])
+
+        self.setup_kpam()
+
+        traj_keyframes, sample_times = self.planner.get_keypose()
+        if len(traj_keyframes) == 0 and len(sample_times) == 0:
+            print("No keyframes found, return empty manip_list")
+            self.manip_list = []
+            return
+        T_world_base = get_relative_transform(
+            get_prim_at_path(self.robot.base_path), get_prim_at_path(self.task.root_prim_path)
+        )
+        self.traj_keyframes = traj_keyframes
+        self.sample_times = sample_times
+        if self.draw:
+            for keypose in traj_keyframes:
+                self.draw.draw_points([(T_world_base @ np.append(keypose[:3, 3], 1))[:3]], [(0, 0, 0, 1)], [7])  # black
+        manip_list = []
+
+        # Update
+        p_base_ee_cur, q_base_ee_cur = self.controller.get_ee_pose()
+        ignore_substring = deepcopy(self.controller.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)
+        # Update
+
+        for i in range(len(self.traj_keyframes)):
+            p_base_ee_tgt = self.traj_keyframes[i][:3, 3]
+            q_base_ee_tgt = R.from_matrix(self.traj_keyframes[i][:3, :3]).as_quat(scalar_first=True)
+            if i <= self.contact_pose_index - 1:
+                cmd = (p_base_ee_tgt, q_base_ee_tgt, "open_gripper", {})
+                manip_list.append(cmd)
+                if i == self.contact_pose_index - 1:
+                    # Update
+                    ignore_substring = deepcopy(self.controller.ignore_substring)
+                    parent_name = self.art_obj.prim_path.split("/")[-2]
+                    ignore_substring.append(parent_name)
+                    cmd = (
+                        p_base_ee_tgt,
+                        q_base_ee_tgt,
+                        "update_specific",
+                        {
+                            "ignore_substring": ignore_substring,
+                            "reference_prim_path": self.controller.reference_prim_path,
+                        },
+                    )
+                    manip_list.append(cmd)
+                    # Update
+            elif i == self.contact_pose_index:
+                if "hearth" in self.art_obj.name:
+                    cmd = (p_base_ee_tgt, q_base_ee_tgt, "open_gripper", {})
+                    manip_list.append(cmd)
+                else:
+                    cmd = (p_base_ee_tgt, q_base_ee_tgt, "close_gripper", {})
+                    for k in range(40):
+                        manip_list.append(cmd)
+            else:
+                cmd = (p_base_ee_tgt, q_base_ee_tgt, "close_gripper", {})
+                manip_list.append(cmd)
+
+        if "hearth" in self.art_obj.name:
+            cmd = (p_base_ee_tgt, q_base_ee_tgt, "close_gripper", {})
+            for k in range(40):
+                manip_list.append(cmd)
+            # Rotate
+            for k in range(100):
+                cmd = (
+                    p_base_ee_tgt,
+                    q_base_ee_tgt,
+                    "in_plane_rotation",
+                    {"target_rotate": self.success_threshold * k / 50},
+                )
+                manip_list.append(cmd)
+            # Rotate
+
+        self.manip_list = manip_list
+
+    def update(self):
+        curr_joint_p = self.art_obj._articulation_view.get_joint_positions()[:, self.art_obj.object_joint_index]
+        self.art_obj._articulation_view.set_joint_position_targets(
+            positions=curr_joint_p, joint_indices=self.art_obj.object_joint_index
+        )
+
+    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():
+            self.manip_list.pop(0)
+            print("POP one manip cmd")
+        if self.is_success():
+            self.manip_list.clear()
+            print("Rotate Done")
+        return len(self.manip_list) == 0
+
+    def is_success(self):
+        curr_joint_p = self.art_obj._articulation_view.get_joint_positions()[:, self.art_obj.object_joint_index]
+        distance = np.abs(curr_joint_p - self.art_obj.articulation_initial_joint_position)
+        return distance >= np.abs(self.success_threshold)

workflows/simbox/core/skills/rotate_obj.py

@@ -0,0 +1,250 @@
+import numpy as np
+import torch
+from core.skills.base_skill import BaseSkill, register_skill
+from core.utils.interpolate_utils import linear_interpolation
+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=unused-argument
+@register_skill
+class Rotate_Obj(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.move_obj = task.objects[cfg["objects"][0]]
+        self.skill_cfg = cfg
+        self.robot_base_path = self.controller.robot_base_path
+        self.T_world_base = get_relative_transform(
+            get_prim_at_path(self.robot_base_path), get_prim_at_path(self.task.root_prim_path)
+        )
+        if "left" in controller.robot_file:
+            self.robot_lr = "left"
+        elif "right" in controller.robot_file:
+            self.robot_lr = "right"
+        self.manip_list = []
+        self.success_threshold_move = self.skill_cfg["success_threshold_move"]
+        self.success_threshold_rotate = self.skill_cfg["success_threshold_rotate"]
+
+    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)
+
+        ignore_substring = 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)
+
+        self.p_base_ee_tgt, self.q_base_ee_tgt = self.getTgtPose()
+
+        self.dummy_forward_cfg = self.skill_cfg.get("dummy_forward", None)
+        if self.dummy_forward_cfg:
+            curr_js, target_js = self.get_tgt_js()
+            interp_js_list = linear_interpolation(curr_js, target_js, self.dummy_forward_cfg.get("num_steps", 10))
+            for js in interp_js_list:
+                p_base_ee, q_base_ee = self.controller.forward_kinematic(js)
+                cmd = (
+                    p_base_ee,
+                    q_base_ee,
+                    "dummy_forward",
+                    {
+                        "arm_action": js,
+                        "gripper_state": self.dummy_forward_cfg.get("gripper_state", 1.0),
+                    },
+                )
+                manip_list.append(cmd)
+
+        first_motion = self.skill_cfg.get("first_motion", None)
+        gripper_state = self.skill_cfg.get("gripper_state", -1.0)
+        gripper_cmd = "close_gripper" if gripper_state == -1.0 else "open_gripper"
+
+        if first_motion == "move":
+            move_offset = self.skill_cfg.get("move_offset", [0, 0, 0])
+            cmd = (
+                self.p_base_ee_tgt + np.array(move_offset),
+                q_base_ee_cur,
+                gripper_cmd,
+                {},
+            )
+            manip_list.append(cmd)
+
+        elif first_motion == "rotate":
+            rotate_offset = self.skill_cfg.get("rotate_offset", [0, 0, 0])
+            cmd = (
+                p_base_ee_cur + np.array(rotate_offset),
+                self.q_base_ee_tgt,
+                gripper_cmd,
+                {},
+            )
+            manip_list.append(cmd)
+
+            if self.skill_cfg.get("rotate_only", False):
+                self.p_base_ee_tgt = p_base_ee_cur + np.array(rotate_offset)
+
+        cmd = (
+            self.p_base_ee_tgt,
+            self.q_base_ee_tgt,
+            gripper_cmd,
+            {},
+        )
+        manip_list.append(cmd)
+
+        self.manip_list = manip_list
+
+    def getTgtPose(self):
+        p_base_ee_cur, q_base_ee_cur = self.controller.get_ee_pose()
+        p_world_move_obj_cur, q_world_move_obj_cur = self.move_obj.get_world_pose()  # w,x,y,z
+
+        T_world_ee_cur = self.T_world_base @ tf_matrix_from_pose(p_base_ee_cur, q_base_ee_cur)
+        T_world_move_obj_cur = tf_matrix_from_pose(p_world_move_obj_cur, q_world_move_obj_cur)
+
+        if self.skill_cfg.get("obj_axis_offset", None):
+            obj_axis_offset = self.skill_cfg.get("obj_axis_offset", None)
+            axis_map = {"x": 0, "y": 1, "z": 2}
+            dT = np.eye(4)
+            for axis, offset in obj_axis_offset:
+                if axis in axis_map:
+                    dT[axis_map[axis], 3] = offset
+            T_world_move_obj_cur = T_world_move_obj_cur @ dT
+
+        q_world_move_obj_tgt = self.getTgtMoveObjOrientation()
+        T_world_move_obj_tgt = T_world_move_obj_cur.copy()
+        T_world_move_obj_tgt[:3, :3] = R.from_quat(q_world_move_obj_tgt, scalar_first=True).as_matrix()
+        p_base_ee_tgt, q_base_ee_tgt = self.calTgtEEPose(
+            T_world_ee_cur,
+            T_world_move_obj_cur,
+            T_world_move_obj_tgt,
+        )
+        trans_offset = self.skill_cfg.get("trans_offset", None)
+        if trans_offset:
+            p_base_ee_tgt += np.array(trans_offset)
+            for _ in range(2):
+                if self.controller.test_single_forward(p_base_ee_tgt, q_base_ee_tgt):
+                    break
+                else:
+                    q_base_ee_tgt[2] -= 0.005
+
+        return p_base_ee_tgt, q_base_ee_tgt
+
+    def calTgtEEPose(self, T_world_ee_cur=None, T_world_move_obj_cur=None, T_world_move_obj_tgt=None):
+        T_base_ee_tgt = (
+            np.linalg.inv(self.T_world_base)
+            @ T_world_move_obj_tgt
+            @ np.linalg.inv(T_world_move_obj_cur)
+            @ T_world_ee_cur
+        )
+
+        return pose_from_tf_matrix(T_base_ee_tgt)
+
+    def getTgtMoveObjOrientation(self):
+        _, q_world_move_obj_cur = self.move_obj.get_world_pose()  # w,x,y,z
+        rotate_obj_euler_delta = np.random.uniform(
+            *self.skill_cfg.get("rotate_obj_euler_delta", [[0, 0, 0], [0, 0, 0]])
+        )
+        q_world_move_obj_tgt = (
+            R.from_quat(q_world_move_obj_cur, scalar_first=True).as_matrix()
+            @ R.from_euler("XYZ", rotate_obj_euler_delta, degrees=True).as_matrix()
+        )
+        q_world_move_obj_tgt = R.from_matrix(q_world_move_obj_tgt).as_quat(scalar_first=True)
+
+        return q_world_move_obj_tgt
+
+    def get_tgt_js(self):
+        """
+        Compute target joint configuration based on current joint state and
+        joint control commands defined in skill configuration.
+
+        Returns:
+            curr_js (np.ndarray): Current joint positions of the controlled arm.
+            target_js (np.ndarray): Target joint positions after applying commands.
+        """
+
+        # --- Get current joint positions ---
+        joint_positions = self.robot.get_joints_state().positions
+
+        if isinstance(joint_positions, torch.Tensor):
+            curr_js = joint_positions.detach().cpu().numpy()[self.controller.arm_indices]
+        elif isinstance(joint_positions, np.ndarray):
+            curr_js = joint_positions[self.controller.arm_indices]
+        else:
+            raise TypeError(f"Unsupported joint state type: {type(joint_positions)}")
+
+        target_js = curr_js.copy()
+
+        # --- Apply joint control commands ---
+        # ctrl_list: list of (joint_index, angle_in_deg, mode), mode in {"abs", "delta"}
+        ctrl_list = self.skill_cfg.get("ctrl_list", [])
+        for joint_idx, angle_deg, mode in ctrl_list:
+            angle_rad = angle_deg * np.pi / 180.0
+            if mode == "abs":
+                target_js[joint_idx] = angle_rad
+            elif mode == "delta":
+                target_js[joint_idx] += angle_rad
+            else:
+                raise ValueError(f"Unknown control mode: {mode}")
+
+        # --- Apply robot-specific joint limits / safety clamps ---
+        robot_file = self.controller.robot_file.lower()
+
+        if "piper" in robot_file:
+            # Example: clamp elbow and wrist joints for Piper robot
+            target_js[2] = min(target_js[2], 0.0)
+            target_js[4] = np.clip(target_js[4], -1.22, 1.22)
+
+        elif "r5a" in robot_file:
+            # Reserved for R5A-specific constraints
+            pass
+
+        return curr_js, target_js
+
+    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_pos = 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_pos < t_eps,
+            diff_ori < o_eps,
+        )
+        self.plan_flag = self.controller.num_last_cmd > 10
+        if self.plan_flag:
+            print(f"move_only plan_flag: {self.plan_flag}, num_last_cmd: {self.controller.num_last_cmd}")
+        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():
+            self.manip_list.pop(0)
+        return len(self.manip_list) == 0
+
+    def is_success(self):
+        p_base_ee_cur, q_base_ee_cur = self.controller.get_ee_pose()
+        distance = np.linalg.norm(p_base_ee_cur - self.p_base_ee_tgt)
+        flag = (distance < self.success_threshold_move) and (len(self.manip_list) == 0)
+        dot = np.clip(np.dot(q_base_ee_cur, self.q_base_ee_tgt), -1.0, 1.0)
+        angle_diff = 2 * np.arccos(np.abs(dot))
+        flag = (angle_diff < self.success_threshold_rotate) and flag
+        return flag

workflows/simbox/core/skills/scan.py

@@ -0,0 +1,142 @@
+import numpy as np
+from core.skills.base_skill import BaseSkill, register_skill
+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 Scan(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
+        self.name = cfg["name"]
+        object_name = cfg["objects"][0]
+        self.pick_obj = task.objects[cfg["objects"][0]]
+        self.manip_list = []
+        lr_arm = "right" if "right" in self.controller.robot_file else "left"
+        self.pickcontact_view = task.pickcontact_views[robot.name][lr_arm][object_name]
+        self.process_valid = True
+        if "right" in self.controller.robot_file:
+            self.robot_ee_path = self.robot.fr_ee_path
+            self.robot_base_path = self.robot.fr_base_path
+        elif "left" in self.controller.robot_file:
+            self.robot_ee_path = self.robot.fl_ee_path
+            self.robot_base_path = self.robot.fl_base_path
+
+    def simple_generate_manip_cmds(self):
+        manip_list = []
+        place_traj = self.sample_place_traj()
+        if len(place_traj) > 1:
+            # Having waypoints
+            for waypoint in place_traj[:-1]:
+                p_base_ee, q_base_ee = waypoint[:3], waypoint[3:]
+                cmd = (p_base_ee, q_base_ee, "close_gripper", {})
+                manip_list.append(cmd)
+        # The last waypoint
+        p_base_ee, q_base_ee = place_traj[-1][:3], place_traj[-1][3:]
+        cmd = (p_base_ee, q_base_ee, "close_gripper", {})
+        manip_list.append(cmd)
+        self.manip_list = manip_list
+
+    def sample_place_traj(self):
+        place_traj = []
+        T_base_ee = get_relative_transform(get_prim_at_path(self.robot_ee_path), get_prim_at_path(self.robot_base_path))
+        T_world_base = get_relative_transform(get_prim_at_path(self.robot_base_path), get_prim_at_path("/World"))
+        T_world_ee = T_world_base @ T_base_ee
+
+        # 1. Objtaining q_world_ee
+        gripper_axis = np.array([1, 0, 0])
+        gripper_axis = gripper_axis / np.linalg.norm(gripper_axis)  # Normalize the vector
+        camera_axis = np.array([0, 1, 1])
+        q_world_ee = self.get_ee_ori(gripper_axis, T_world_ee, camera_axis)
+        # 2. Obtaining p_world_ee
+        p_world_ee_init = self.controller.T_world_ee_init[0:3, 3]  # Getting initial ee position
+        p_world_ee = p_world_ee_init.copy()
+        p_world_ee[0] += np.random.uniform(-0.02, 0.02)
+        p_world_ee[1] += np.random.uniform(0.15, 0.2)
+        p_world_ee[2] += np.random.uniform(-0.14, -0.16)
+        # Place
+        place_traj.append(self.adding_waypoint(p_world_ee, q_world_ee, T_world_base))
+
+        return place_traj
+
+    def get_ee_ori(self, gripper_axis, T_world_ee, camera_axis=None):
+        gripper_x = gripper_axis
+        if camera_axis is not None:
+            gripper_z = camera_axis
+        else:
+            current_z = T_world_ee[0:3, 1]
+            gripper_z = current_z - np.dot(current_z, gripper_x) * gripper_x
+
+        gripper_z = gripper_z / np.linalg.norm(gripper_z)
+        gripper_y = np.cross(gripper_z, gripper_x)
+        gripper_y = gripper_y / np.linalg.norm(gripper_y)
+        gripper_z = np.cross(gripper_x, gripper_y)
+        R_world_ee = np.column_stack((gripper_x, gripper_y, gripper_z))
+        q_world_ee = R.from_matrix(R_world_ee).as_quat(scalar_first=True)
+        return q_world_ee
+
+    def adding_waypoint(self, p_world_ee, q_world_ee, T_world_base):
+        """
+        Adding a waypoint, also transform from wolrd frame to robot frame
+        """
+        T_world_ee = tf_matrix_from_pose(p_world_ee, q_world_ee)
+        T_base_ee = np.linalg.inv(T_world_base) @ T_world_ee
+        p_base_ee, q_base_ee = pose_from_tf_matrix(T_base_ee)
+        waypoint = np.concatenate((p_base_ee, q_base_ee))
+        return waypoint
+
+    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.pick_obj.get_linear_velocity())) < 5
+            )
+
+        flag = flag and self.process_valid
+
+        return flag

workflows/simbox/core/skills/track.py

@@ -0,0 +1,218 @@
+import numpy as np
+from core.skills.base_skill import BaseSkill, register_skill
+from core.utils.transformation_utils import get_orientation, perturb_orientation
+from core.utils.usd_geom_utils import compute_bbox
+from omegaconf import DictConfig
+from omni.isaac.core.controllers import BaseController
+from omni.isaac.core.objects.cylinder import VisualCylinder
+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,
+)
+from scipy.spatial.transform import Rotation as R
+
+
+# pylint: disable=consider-using-generator,too-many-public-methods,unused-argument
+@register_skill
+class Track(BaseSkill):
+    def __init__(
+        self, robot: Robot, controller: BaseController, task: BaseTask, cfg: DictConfig, world, *args, **kwargs
+    ):
+        super().__init__()
+        self.robot = robot
+        self.controller = controller
+        self.task = task
+        self.skill_cfg = cfg
+        self.world = world
+        self.frame = self.skill_cfg.get("frame", "robot")
+        self.robot_base_path = self.controller.robot_base_path
+        self.T_base_2_world = get_relative_transform(
+            get_prim_at_path(self.robot_base_path), get_prim_at_path(self.task.root_prim_path)
+        )
+        self.table_2_base = self.cal_table_2_base()
+        self.way_points = self.sample_waypoints()
+        self.last_target_trans = self.way_points[-1][0]
+        self.last_target_ori = get_orientation(None, self.way_points[-1][1])
+        if "left" in controller.robot_file:
+            self.robot_lr = "left"
+            self.visual_color = {
+                "x": np.array([1.0, 0.0, 0.0]),
+                "y": np.array([0.0, 1.0, 0.0]),
+                "z": np.array([0.0, 0.0, 1.0]),
+            }
+        elif "right" in controller.robot_file:
+            self.robot_lr = "right"
+            self.visual_color = {
+                "x": np.array([1.0, 1.0, 0.0]),
+                "y": np.array([0.0, 1.0, 1.0]),
+                "z": np.array([1.0, 0.0, 1.0]),
+            }
+        self.vs_cylinder_radius = 0.005
+        self.vs_cylinder_height = 0.13
+        self.T_tcp_2_ee = np.array(self.skill_cfg.get("T_tcp_2_ee", np.eye(4)))
+
+        # !!! keyposes should be generated after previous skill is done
+        self.manip_list = []
+
+        T_tcp_2_world = self.get_tcp_pose()
+        trans, ori = self.cal_axis(T_tcp_2_world)
+        for axis in ["x", "y", "z"]:
+            self.task.visuals[f"{axis}_{self.robot_lr}"] = VisualCylinder(
+                prim_path=f"/World/visual/{axis}_{self.robot_lr}",
+                radius=self.vs_cylinder_radius,
+                height=self.vs_cylinder_height,
+                translation=trans[axis],
+                orientation=ori[axis],
+                color=self.visual_color[f"{axis}"],
+            )
+
+    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": [0, 0, 0, 0, 0, 0]})
+        manip_list.append(cmd)
+
+        if self.frame == "robot":
+            for target in self.way_points:
+                cmd = (np.array(target[0]), get_orientation(None, target[1]), "open_gripper", {})
+                manip_list.append(cmd)
+        else:
+            raise NotImplementedError
+
+        self.manip_list = manip_list
+
+    def get_tcp_pose(self, frame: str = "world"):
+        if frame == "world":
+            p_base_ee, q_base_ee = self.controller.get_ee_pose()
+            T_ee_2_base = tf_matrix_from_pose(p_base_ee, q_base_ee)
+            T_tcp_2_world = self.T_base_2_world @ T_ee_2_base @ self.T_tcp_2_ee
+            return T_tcp_2_world
+        else:
+            raise NotImplementedError
+
+    def visualize_target(self, world):
+        if len(self.manip_list) > 0:
+            p_base_ee, q_base_ee, *_ = self.manip_list[0]
+            T_ee_2_base = tf_matrix_from_pose(p_base_ee, q_base_ee)
+            T_tcp_2_world = self.T_base_2_world @ T_ee_2_base @ self.T_tcp_2_ee
+            trans, ori = self.cal_axis(T_tcp_2_world)
+
+            for axis in ["x", "y", "z"]:
+                self.task.visuals[f"{axis}_{self.robot_lr}"].set_world_pose(trans[axis], ori[axis])
+
+            if self.curr_way_points_num > len(self.manip_list):
+                self.curr_way_points_num -= 1
+                for _ in range(40):
+                    world.step(render=True)
+
+        # elif self.task.visuals["x"] and self.task.visuals["y"] and self.task.visuals["z"]:
+        #     for axis in ['x', 'y', 'z']:
+        #         delete_prim(f"/World/visual/{axis}_{self.robot_lr}")
+        #         for i in range(40):
+        #             world.step(render=True)
+
+    def cal_axis(self, T):
+        origin = T[:3, 3]
+        rotation_matrix = T[:3, :3]
+        trans = {}
+        ori = {}
+        axis_table = ["x", "y", "z"]
+
+        for i in range(3):
+            axis = rotation_matrix[:, i]
+
+            axis_start = origin
+            axis_end = origin + axis * self.vs_cylinder_height
+            trans[axis_table[i]] = self.calculate_cylinder_center(axis_start, axis_end)
+            ori[axis_table[i]] = self.calculate_orientation(axis)
+
+        return trans, ori
+
+    def calculate_cylinder_center(self, start_point, end_point):
+        return (start_point + end_point) / 2
+
+    def calculate_orientation(self, axis_vector):
+        default_z_axis = np.array([0, 0, 1])
+        target_axis = axis_vector / np.linalg.norm(axis_vector)
+        rotation = R.align_vectors([target_axis], [default_z_axis])[0]
+        return rotation.as_quat(scalar_first=True)
+
+    def cal_table_2_base(self):
+        tgt = self.task.fixtures["table"]
+        bbox_tgt = compute_bbox(tgt.prim)
+        table_center = (np.asarray(bbox_tgt.min) + np.asarray(bbox_tgt.max)) / 2
+        tgt_z_max = bbox_tgt.max[2]
+        table_center[2] = tgt_z_max
+        base_trans = self.T_base_2_world[:3, 3]
+        table_2_base = table_center - base_trans
+        table_2_base[0], table_2_base[1] = table_2_base[1], -table_2_base[0]
+
+        return table_2_base
+
+    def from_table_2_base(self, trans, ori):
+        return (np.array(trans) + self.table_2_base).tolist(), ori
+
+    def sample_waypoints(self):
+        way_points_num = self.skill_cfg.get("way_points_num", 1)
+        self.curr_way_points_num = way_points_num
+        way_points_trans = self.skill_cfg.get("way_points_trans", None)
+        way_points_ori = np.array(self.skill_cfg.get("way_points_ori", None))
+        trans_min = np.array(way_points_trans["min"])
+        trans_max = np.array(way_points_trans["max"])
+
+        way_points = []
+        for i in range(way_points_num):
+            while True:
+                print(f"... sampling waypoint {i} ...")
+                x = np.random.uniform(trans_min[0], trans_max[0])
+                y = np.random.uniform(trans_min[1], trans_max[1])
+                z = np.random.uniform(trans_min[2], trans_max[2])
+                trans = [x, y, z]
+                ori = perturb_orientation(way_points_ori, self.skill_cfg.get("max_noise_deg", 5))
+                trans, ori = self.from_table_2_base(trans, ori)
+
+                if self.controller.test_single_ik(trans, ori):
+                    way_points.append([trans, ori.tolist()])
+                    break
+
+        for p in way_points:
+            print(p)
+
+        return way_points
+
+    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():
+            self.manip_list.pop(0)
+        return len(self.manip_list) == 0
+
+    def is_success(self, t_eps=5e-3, o_eps=0.087):
+        p_base_ee_cur, q_base_ee_cur = self.controller.get_ee_pose()
+        p_base_ee, q_base_ee = self.last_target_trans, self.last_target_ori
+        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_or(
+            diff_trans < t_eps,
+            diff_ori < o_eps,
+        )
+        return pose_flag

workflows/simbox/core/skills/wait.py

@@ -0,0 +1,88 @@
+import numpy as np
+from core.skills.base_skill import BaseSkill, register_skill
+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
+
+
+# pylint: disable=consider-using-generator,too-many-public-methods,unused-argument
+@register_skill
+class Wait(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
+        self.success_threshold = cfg["success_threshold"]
+        self.name = cfg["name"]
+        self.move_obj = task.objects[cfg["objects"][0]]
+        if "left" in controller.robot_file:
+            self.robot_lr = "left"
+        elif "right" in controller.robot_file:
+            self.robot_lr = "right"
+
+        self.manip_list = []
+
+    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": [0, 0, 1, 0, 0, 0]})
+        manip_list.append(cmd)
+        ignore_substring = 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)
+
+        self.p_base_ee_tgt = p_base_ee_cur
+        self.q_base_ee_tgt = q_base_ee_cur
+
+        cmd = (
+            self.p_base_ee_tgt,
+            self.q_base_ee_tgt,
+            "close_gripper" if self.skill_cfg.get("gripper_state", -1.0) == -1.0 else "open_gripper",
+            {},
+        )
+
+        for _ in range(self.skill_cfg.get("wait_steps", 50)):
+            manip_list.append(cmd)
+
+        self.manip_list = manip_list
+
+    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
+        if self.plan_flag:
+            print(f"move_only plan_flag: {self.plan_flag}, num_last_cmd: {self.controller.num_last_cmd}")
+
+        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():
+            self.manip_list.pop(0)
+        return len(self.manip_list) == 0
+
+    def is_success(self):
+        p_base_ee_cur, _ = self.controller.get_ee_pose()
+        distance = np.linalg.norm(p_base_ee_cur - self.p_base_ee_tgt)
+        flag = (distance < self.success_threshold) and (len(self.manip_list) == 0)
+
+        return flag