From 1924f0a3b2ecca0e218362864d449d9ab57f7ebf Mon Sep 17 00:00:00 2001 From: arth-shukla <44527268+arth-shukla@users.noreply.github.com> Date: Sun, 21 Jan 2024 12:53:56 -0800 Subject: [PATCH 1/8] Progress: Fetch mobile base --- mani_skill2/agents/robots/fetch/__init__.py | 2 +- mani_skill2/agents/robots/fetch/fetch.py | 107 ++++++++++-------- mani_skill2/assets/robots/fetch/fetch.urdf | 40 +++++++ .../table/table_scene_builder.py | 14 +++ 4 files changed, 117 insertions(+), 46 deletions(-) diff --git a/mani_skill2/agents/robots/fetch/__init__.py b/mani_skill2/agents/robots/fetch/__init__.py index ff32d0c84..92e483603 100644 --- a/mani_skill2/agents/robots/fetch/__init__.py +++ b/mani_skill2/agents/robots/fetch/__init__.py @@ -1 +1 @@ -from .fetch import Fetch +from .fetch import Fetch, FETCH_UNIQUE_COLLISION_BIT diff --git a/mani_skill2/agents/robots/fetch/fetch.py b/mani_skill2/agents/robots/fetch/fetch.py index a05518c17..e7b714709 100644 --- a/mani_skill2/agents/robots/fetch/fetch.py +++ b/mani_skill2/agents/robots/fetch/fetch.py @@ -8,7 +8,7 @@ from mani_skill2.agents.base_agent import BaseAgent from mani_skill2.agents.controllers import * from mani_skill2.sensors.camera import CameraConfig -from mani_skill2.utils.common import np_compute_angle_between +from mani_skill2.utils.common import compute_angle_between from mani_skill2.utils.sapien_utils import ( compute_total_impulse, get_actor_contacts, @@ -16,6 +16,7 @@ get_pairwise_contact_impulse, ) +FETCH_UNIQUE_COLLISION_BIT = 1 << 30 class Fetch(BaseAgent): uid = "fetch" @@ -67,6 +68,12 @@ def __init__(self, *args, **kwargs): self.body_damping = 1e2 self.body_force_limit = 100 + self.base_joint_names = [ + "root_x_axis_joint", + "root_y_axis_joint", + "root_z_rotation_joint", + ] + super().__init__(*args, **kwargs) @property @@ -173,71 +180,61 @@ def controller_configs(self): # -------------------------------------------------------------------------- # body_pd_joint_pos = PDJointPosControllerConfig( self.body_joint_names, - None, - None, + # None, + # None, self.body_stiffness, self.body_damping, self.body_force_limit, normalize_action=False, ) + # -------------------------------------------------------------------------- # + # Base + # -------------------------------------------------------------------------- # + base_pd_joint_vel = PDBaseVelControllerConfig( + self.base_joint_names, + lower=[-0.5, -0.5, -3.14], + upper=[0.5, 0.5, 3.14], + damping=1000, + force_limit=500, + ) + + controller_configs = dict( pd_joint_delta_pos=dict( - arm=arm_pd_joint_delta_pos, - gripper=gripper_pd_joint_pos, - body=body_pd_joint_pos, + arm=arm_pd_joint_delta_pos, gripper=gripper_pd_joint_pos, body=body_pd_joint_pos, base=base_pd_joint_vel, ), pd_joint_pos=dict( - arm=arm_pd_joint_pos, - gripper=gripper_pd_joint_pos, - body=body_pd_joint_pos, + arm=arm_pd_joint_pos, gripper=gripper_pd_joint_pos, body=body_pd_joint_pos, base=base_pd_joint_vel, ), pd_ee_delta_pos=dict( - arm=arm_pd_ee_delta_pos, - gripper=gripper_pd_joint_pos, - body=body_pd_joint_pos, + arm=arm_pd_ee_delta_pos, gripper=gripper_pd_joint_pos, body=body_pd_joint_pos, base=base_pd_joint_vel, ), pd_ee_delta_pose=dict( - arm=arm_pd_ee_delta_pose, - gripper=gripper_pd_joint_pos, - body=body_pd_joint_pos, + arm=arm_pd_ee_delta_pose, gripper=gripper_pd_joint_pos, body=body_pd_joint_pos, base=base_pd_joint_vel, ), pd_ee_delta_pose_align=dict( - arm=arm_pd_ee_delta_pose_align, - gripper=gripper_pd_joint_pos, - body=body_pd_joint_pos, + arm=arm_pd_ee_delta_pose_align, gripper=gripper_pd_joint_pos, body=body_pd_joint_pos, base=base_pd_joint_vel, ), # TODO(jigu): how to add boundaries for the following controllers pd_joint_target_delta_pos=dict( - arm=arm_pd_joint_target_delta_pos, - gripper=gripper_pd_joint_pos, - body=body_pd_joint_pos, + arm=arm_pd_joint_target_delta_pos, gripper=gripper_pd_joint_pos, body=body_pd_joint_pos, base=base_pd_joint_vel, ), pd_ee_target_delta_pos=dict( - arm=arm_pd_ee_target_delta_pos, - gripper=gripper_pd_joint_pos, - body=body_pd_joint_pos, + arm=arm_pd_ee_target_delta_pos, gripper=gripper_pd_joint_pos, body=body_pd_joint_pos, base=base_pd_joint_vel, ), pd_ee_target_delta_pose=dict( - arm=arm_pd_ee_target_delta_pose, - gripper=gripper_pd_joint_pos, - body=body_pd_joint_pos, + arm=arm_pd_ee_target_delta_pose, gripper=gripper_pd_joint_pos, body=body_pd_joint_pos, base=base_pd_joint_vel, ), # Caution to use the following controllers pd_joint_vel=dict( - arm=arm_pd_joint_vel, - gripper=gripper_pd_joint_pos, - body=body_pd_joint_pos, + arm=arm_pd_joint_vel, gripper=gripper_pd_joint_pos, body=body_pd_joint_pos, base=base_pd_joint_vel, ), pd_joint_pos_vel=dict( - arm=arm_pd_joint_pos_vel, - gripper=gripper_pd_joint_pos, - body=body_pd_joint_pos, + arm=arm_pd_joint_pos_vel, gripper=gripper_pd_joint_pos, body=body_pd_joint_pos, base=base_pd_joint_vel, ), pd_joint_delta_pos_vel=dict( - arm=arm_pd_joint_delta_pos_vel, - gripper=gripper_pd_joint_pos, - body=body_pd_joint_pos, + arm=arm_pd_joint_delta_pos_vel, gripper=gripper_pd_joint_pos, body=body_pd_joint_pos, base=base_pd_joint_vel, ), ) @@ -245,13 +242,31 @@ def controller_configs(self): return deepcopy_dict(controller_configs) def _after_init(self): - self.finger1_link = get_obj_by_name( + self.finger1_link: sapien.Entity = get_obj_by_name( self.robot.get_links(), "l_gripper_finger_link" - ) - self.finger2_link = get_obj_by_name( + ).entity + self.finger2_link: sapien.Entity = get_obj_by_name( self.robot.get_links(), "r_gripper_finger_link" + ).entity + self.tcp: physx.PhysxArticulationLinkComponent = get_obj_by_name( + self.robot.get_links(), self.ee_link_name + ) + + self.base_link: physx.PhysxArticulationLinkComponent = get_obj_by_name( + self.robot.get_links(), "base_link" + ) + self.l_wheel_link: physx.PhysxArticulationLinkComponent = get_obj_by_name( + self.robot.get_links(), "l_wheel_link" + ) + self.r_wheel_link: physx.PhysxArticulationLinkComponent = get_obj_by_name( + self.robot.get_links(), "r_wheel_link" ) - self.tcp = get_obj_by_name(self.robot.get_links(), self.ee_link_name) + + for link in [self.base_link, self.l_wheel_link, self.r_wheel_link]: + cs = link.get_collision_shapes()[0] + cg = cs.get_collision_groups() + cg[2] = FETCH_UNIQUE_COLLISION_BIT + cs.set_collision_groups(cg) def is_grasping(self, object: sapien.Entity = None, min_impulse=1e-6, max_angle=85): contacts = self.scene.get_contacts() @@ -272,8 +287,8 @@ def is_grasping(self, object: sapien.Entity = None, min_impulse=1e-6, max_angle= rdirection = self.finger2_link.pose.to_transformation_matrix()[:3, 1] # angle between impulse and open direction - langle = np_compute_angle_between(ldirection, limpulse) - rangle = np_compute_angle_between(rdirection, rimpulse) + langle = compute_angle_between(ldirection, limpulse) + rangle = compute_angle_between(rdirection, rimpulse) lflag = ( np.linalg.norm(limpulse) >= min_impulse @@ -313,7 +328,9 @@ def sensor_configs(self): entity_uid="head_camera_link", ) ] - + @property - def tcp_pose_p(self): - return (self.finger1_link.pose.p + self.finger2_link.pose.p) / 2 + def tcp_pose(self) -> sapien.Pose: + p = (self.finger1_link.pose.p + self.finger2_link.pose.p) / 2 + q = (self.finger1_link.pose.q + self.finger2_link.pose.q) / 2 + return sapien.Pose(p=p, q=q) diff --git a/mani_skill2/assets/robots/fetch/fetch.urdf b/mani_skill2/assets/robots/fetch/fetch.urdf index 974556644..be5901837 100644 --- a/mani_skill2/assets/robots/fetch/fetch.urdf +++ b/mani_skill2/assets/robots/fetch/fetch.urdf @@ -1,4 +1,44 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/mani_skill2/utils/scene_builder/table/table_scene_builder.py b/mani_skill2/utils/scene_builder/table/table_scene_builder.py index 8e5564df0..e9a8566b0 100644 --- a/mani_skill2/utils/scene_builder/table/table_scene_builder.py +++ b/mani_skill2/utils/scene_builder/table/table_scene_builder.py @@ -78,6 +78,20 @@ def initialize(self): ) self.env.agent.reset(qpos) self.env.agent.robot.set_pose(sapien.Pose([-0.562, 0, 0])) + elif self.robot_uid == "fetch": + qpos = np.array( + [0, 0, 0, 0.386, 0, 0, 0, -np.pi / 4, 0, np.pi / 4, 0, np.pi / 3, 0, 0.015, 0.015] + ) + self.agent.reset(qpos) + self.agent.robot.set_pose(sapien.Pose([-0.82, 0, -self.table_height])) + + from mani_skill2.agents.robots.fetch import FETCH_UNIQUE_COLLISION_BIT + cs = self.ground.find_component_by_type( + sapien.physx.PhysxRigidStaticComponent + ).get_collision_shapes()[0] + cg = cs.get_collision_groups() + cg[2] = FETCH_UNIQUE_COLLISION_BIT + cs.set_collision_groups(cg) else: raise NotImplementedError(self.env.robot_uid) From 78f1e3e47a769e3e50e2c5564266a4d303a2456a Mon Sep 17 00:00:00 2001 From: arth-shukla <44527268+arth-shukla@users.noreply.github.com> Date: Sun, 21 Jan 2024 13:00:54 -0800 Subject: [PATCH 2/8] fix --- mani_skill2/agents/robots/fetch/fetch.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mani_skill2/agents/robots/fetch/fetch.py b/mani_skill2/agents/robots/fetch/fetch.py index e7b714709..1a33ba634 100644 --- a/mani_skill2/agents/robots/fetch/fetch.py +++ b/mani_skill2/agents/robots/fetch/fetch.py @@ -180,8 +180,8 @@ def controller_configs(self): # -------------------------------------------------------------------------- # body_pd_joint_pos = PDJointPosControllerConfig( self.body_joint_names, - # None, - # None, + None, + None, self.body_stiffness, self.body_damping, self.body_force_limit, From 52aec2379249be3e44af72509f873b180e004c29 Mon Sep 17 00:00:00 2001 From: arth-shukla <44527268+arth-shukla@users.noreply.github.com> Date: Sun, 21 Jan 2024 13:33:15 -0800 Subject: [PATCH 3/8] fix --- mani_skill2/agents/robots/fetch/fetch.py | 17 +---------------- 1 file changed, 1 insertion(+), 16 deletions(-) diff --git a/mani_skill2/agents/robots/fetch/fetch.py b/mani_skill2/agents/robots/fetch/fetch.py index 1a33ba634..3ccdd5626 100644 --- a/mani_skill2/agents/robots/fetch/fetch.py +++ b/mani_skill2/agents/robots/fetch/fetch.py @@ -34,6 +34,7 @@ class Fetch(BaseAgent): ), ), ) + sensor_configs = [] def __init__(self, *args, **kwargs): self.arm_joint_names = [ @@ -312,22 +313,6 @@ def build_grasp_pose(approaching, closing, center): T[:3, :3] = np.stack([ortho, closing, approaching], axis=1) T[:3, 3] = center return sapien.Pose(T) - - @property - def sensor_configs(self): - return [ - CameraConfig( - uid="fetch_head", - p=[0, 0, 0], - q=[0.9238795, 0, 0.3826834, 0], - width=128, - height=128, - fov=1.57, - near=0.01, - far=10, - entity_uid="head_camera_link", - ) - ] @property def tcp_pose(self) -> sapien.Pose: From f3bf3acc98742907aa3d7d109c6b2b6edc5e7f0b Mon Sep 17 00:00:00 2001 From: arth-shukla <44527268+arth-shukla@users.noreply.github.com> Date: Sun, 21 Jan 2024 14:18:33 -0800 Subject: [PATCH 4/8] unfinished: integrate into gpu sim stuff --- .../agents/controllers/base_controller.py | 7 +- mani_skill2/agents/controllers/pd_base_vel.py | 10 +- mani_skill2/agents/robots/fetch/fetch.py | 131 +++++++++++++----- mani_skill2/envs/pick_and_place/base_env.py | 18 +-- .../table/table_scene_builder.py | 8 +- mani_skill2/utils/structs/joint.py | 6 + test.py | 17 +++ 7 files changed, 132 insertions(+), 65 deletions(-) create mode 100644 test.py diff --git a/mani_skill2/agents/controllers/base_controller.py b/mani_skill2/agents/controllers/base_controller.py index 487e9e8ae..75f555620 100644 --- a/mani_skill2/agents/controllers/base_controller.py +++ b/mani_skill2/agents/controllers/base_controller.py @@ -206,12 +206,7 @@ def _initialize_joints(self): def _assert_fully_actuated(self): active_joints = self.articulation.get_active_joints() - if len(active_joints) != len(self.joints) or not np.all( - [ - active_joint == joint - for active_joint, joint in zip(active_joints, self.joints) - ] - ): + if len(active_joints) != len(self.joints) or set(active_joints) != set(self.joints): print("active_joints:", [x.name for x in active_joints]) print("controlled_joints:", [x.name for x in self.joints]) raise AssertionError("{} is not fully actuated".format(self.articulation)) diff --git a/mani_skill2/agents/controllers/pd_base_vel.py b/mani_skill2/agents/controllers/pd_base_vel.py index 098e741a5..059ec5d70 100644 --- a/mani_skill2/agents/controllers/pd_base_vel.py +++ b/mani_skill2/agents/controllers/pd_base_vel.py @@ -1,4 +1,5 @@ import numpy as np +import torch from mani_skill2.utils.geometry import rotate_2d_vec_by_angle @@ -18,9 +19,16 @@ def _initialize_action_space(self): def set_action(self, action: np.ndarray): action = self._preprocess_action(action) + # TODO (arth): add support for batched qpos and gpu sim + if isinstance(self.qpos, torch.Tensor): + qpos = self.qpos.detach().cpu().numpy() + qpos = qpos[0] + if isinstance(action, torch.Tensor): + action = action.detach().cpu().numpy() + # Convert to ego-centric action # Assume the 3rd DoF stands for orientation - ori = self.qpos[2] + ori = qpos[2] vel = rotate_2d_vec_by_angle(action[:2], ori) new_action = np.hstack([vel, action[2:]]) diff --git a/mani_skill2/agents/robots/fetch/fetch.py b/mani_skill2/agents/robots/fetch/fetch.py index 3ccdd5626..7b52c95cc 100644 --- a/mani_skill2/agents/robots/fetch/fetch.py +++ b/mani_skill2/agents/robots/fetch/fetch.py @@ -1,20 +1,29 @@ from copy import deepcopy +from typing import Dict, Tuple import numpy as np import sapien import sapien.physx as physx +import torch from mani_skill2 import PACKAGE_ASSET_DIR from mani_skill2.agents.base_agent import BaseAgent from mani_skill2.agents.controllers import * from mani_skill2.sensors.camera import CameraConfig -from mani_skill2.utils.common import compute_angle_between +from mani_skill2.utils.common import compute_angle_between, np_compute_angle_between from mani_skill2.utils.sapien_utils import ( compute_total_impulse, get_actor_contacts, get_obj_by_name, get_pairwise_contact_impulse, ) +from mani_skill2.utils.structs.actor import Actor +from mani_skill2.utils.sapien_utils import to_tensor +from mani_skill2.utils.structs.base import BaseStruct +from mani_skill2.utils.structs.joint import Joint +from mani_skill2.utils.structs.link import Link +from mani_skill2.utils.structs.pose import Pose +from mani_skill2.utils.structs.types import Array FETCH_UNIQUE_COLLISION_BIT = 1 << 30 @@ -243,64 +252,110 @@ def controller_configs(self): return deepcopy_dict(controller_configs) def _after_init(self): - self.finger1_link: sapien.Entity = get_obj_by_name( + self.finger1_link: Link = get_obj_by_name( self.robot.get_links(), "l_gripper_finger_link" - ).entity - self.finger2_link: sapien.Entity = get_obj_by_name( + ) + self.finger2_link: Link = get_obj_by_name( self.robot.get_links(), "r_gripper_finger_link" - ).entity - self.tcp: physx.PhysxArticulationLinkComponent = get_obj_by_name( + ) + self.tcp: Link = get_obj_by_name( self.robot.get_links(), self.ee_link_name ) - self.base_link: physx.PhysxArticulationLinkComponent = get_obj_by_name( + self.base_link: Link = get_obj_by_name( self.robot.get_links(), "base_link" ) - self.l_wheel_link: physx.PhysxArticulationLinkComponent = get_obj_by_name( + self.l_wheel_link: Link = get_obj_by_name( self.robot.get_links(), "l_wheel_link" ) - self.r_wheel_link: physx.PhysxArticulationLinkComponent = get_obj_by_name( + self.r_wheel_link: Link = get_obj_by_name( self.robot.get_links(), "r_wheel_link" ) - for link in [self.base_link, self.l_wheel_link, self.r_wheel_link]: - cs = link.get_collision_shapes()[0] + cs = link._bodies[0].get_collision_shapes()[0] cg = cs.get_collision_groups() cg[2] = FETCH_UNIQUE_COLLISION_BIT cs.set_collision_groups(cg) - def is_grasping(self, object: sapien.Entity = None, min_impulse=1e-6, max_angle=85): - contacts = self.scene.get_contacts() - if object is None: - finger1_contacts = get_actor_contacts(contacts, self.finger1_link) - finger2_contacts = get_actor_contacts(contacts, self.finger2_link) - return ( - np.linalg.norm(compute_total_impulse(finger1_contacts)) >= min_impulse - and np.linalg.norm(compute_total_impulse(finger2_contacts)) - >= min_impulse - ) - else: - limpulse = get_pairwise_contact_impulse(contacts, self.finger1_link, object) - rimpulse = get_pairwise_contact_impulse(contacts, self.finger2_link, object) + self.queries: Dict[str, Tuple[physx.PhysxGpuContactQuery, Tuple[int]]] = dict() - # direction to open the gripper - ldirection = -self.finger1_link.pose.to_transformation_matrix()[:3, 1] - rdirection = self.finger2_link.pose.to_transformation_matrix()[:3, 1] + def is_grasping(self, object: Actor = None, min_impulse=1e-6, max_angle=85): + # TODO (stao): is_grasping code needs to be updated for new GPU sim + if physx.is_gpu_enabled(): + if object.name not in self.queries: + body_pairs = list(zip(self.finger1_link._bodies, object._bodies)) + body_pairs += list(zip(self.finger2_link._bodies, object._bodies)) + self.queries[object.name] = ( + self.scene.px.gpu_create_contact_query(body_pairs), + (len(object._bodies), 3), + ) + print(f"Create query for Panda grasp({object.name})") + query, contacts_shape = self.queries[object.name] + self.scene.px.gpu_query_contacts(query) + # query.cuda_contacts # (num_unique_pairs * num_envs, 3) + contacts = query.cuda_contacts.clone().reshape((-1, *contacts_shape)) + lforce = torch.linalg.norm(contacts[0], axis=1) + rforce = torch.linalg.norm(contacts[1], axis=1) - # angle between impulse and open direction - langle = compute_angle_between(ldirection, limpulse) - rangle = compute_angle_between(rdirection, rimpulse) + # NOTE (stao): 0.5 * time_step is a decent value when tested on a pick cube task. + min_force = 0.5 * self.scene.px.timestep - lflag = ( - np.linalg.norm(limpulse) >= min_impulse - and np.rad2deg(langle) <= max_angle + # direction to open the gripper + ldirection = -self.finger1_link.pose.to_transformation_matrix()[..., :3, 1] + rdirection = self.finger2_link.pose.to_transformation_matrix()[..., :3, 1] + langle = compute_angle_between(ldirection, contacts[0]) + rangle = compute_angle_between(rdirection, contacts[1]) + lflag = torch.logical_and( + lforce >= min_force, torch.rad2deg(langle) <= max_angle ) - rflag = ( - np.linalg.norm(rimpulse) >= min_impulse - and np.rad2deg(rangle) <= max_angle + rflag = torch.logical_and( + rforce >= min_force, torch.rad2deg(rangle) <= max_angle ) - return all([lflag, rflag]) + return torch.logical_and(lflag, rflag) + else: + contacts = self.scene.get_contacts() + + if object is None: + finger1_contacts = get_actor_contacts(contacts, self.finger1_link) + finger2_contacts = get_actor_contacts(contacts, self.finger2_link) + return ( + np.linalg.norm(compute_total_impulse(finger1_contacts)) + >= min_impulse + and np.linalg.norm(compute_total_impulse(finger2_contacts)) + >= min_impulse + ) + else: + limpulse = get_pairwise_contact_impulse( + contacts, self.finger1_link, object + ) + rimpulse = get_pairwise_contact_impulse( + contacts, self.finger2_link, object + ) + + # direction to open the gripper + ldirection = -self.finger1_link.pose.to_transformation_matrix()[ + ..., :3, 1 + ] + rdirection = self.finger2_link.pose.to_transformation_matrix()[ + ..., :3, 1 + ] + + # TODO Convert this to batched code + # angle between impulse and open direction + langle = np_compute_angle_between(ldirection[0], limpulse) + rangle = np_compute_angle_between(rdirection[0], rimpulse) + + lflag = ( + np.linalg.norm(limpulse) >= min_impulse + and np.rad2deg(langle) <= max_angle + ) + rflag = ( + np.linalg.norm(rimpulse) >= min_impulse + and np.rad2deg(rangle) <= max_angle + ) + + return all([lflag, rflag]) @staticmethod def build_grasp_pose(approaching, closing, center): @@ -315,7 +370,7 @@ def build_grasp_pose(approaching, closing, center): return sapien.Pose(T) @property - def tcp_pose(self) -> sapien.Pose: + def tcp_pose(self) -> Pose: p = (self.finger1_link.pose.p + self.finger2_link.pose.p) / 2 q = (self.finger1_link.pose.q + self.finger2_link.pose.q) / 2 return sapien.Pose(p=p, q=q) diff --git a/mani_skill2/envs/pick_and_place/base_env.py b/mani_skill2/envs/pick_and_place/base_env.py index 2b93cff00..8751ef717 100644 --- a/mani_skill2/envs/pick_and_place/base_env.py +++ b/mani_skill2/envs/pick_and_place/base_env.py @@ -75,25 +75,11 @@ def _initialize_agent(self): self.agent.reset(qpos) self.agent.robot.set_pose(Pose([-0.562, 0, 0])) elif self.robot_uid == "fetch": - # TODO (arth): Better fetch init qpos = np.array( - [ - 0.386, - 0, - 0, - 0, - -np.pi / 4, - 0, - np.pi / 4, - 0, - np.pi / 3, - 0, - 0.015, - 0.015, - ] + [0, 0, 0, 0.386, 0, 0, 0, -np.pi / 4, 0, np.pi / 4, 0, np.pi / 3, 0, 0.015, 0.015] ) self.agent.reset(qpos) - self.agent.robot.set_pose(Pose([-0.82, 0, -0.920])) + self.agent.robot.set_pose(sapien.Pose([-0.82, 0, -0.920])) else: raise NotImplementedError(self.robot_uid) diff --git a/mani_skill2/utils/scene_builder/table/table_scene_builder.py b/mani_skill2/utils/scene_builder/table/table_scene_builder.py index e9a8566b0..7ff1110de 100644 --- a/mani_skill2/utils/scene_builder/table/table_scene_builder.py +++ b/mani_skill2/utils/scene_builder/table/table_scene_builder.py @@ -78,15 +78,15 @@ def initialize(self): ) self.env.agent.reset(qpos) self.env.agent.robot.set_pose(sapien.Pose([-0.562, 0, 0])) - elif self.robot_uid == "fetch": + elif self.env.robot_uid == "fetch": qpos = np.array( [0, 0, 0, 0.386, 0, 0, 0, -np.pi / 4, 0, np.pi / 4, 0, np.pi / 3, 0, 0.015, 0.015] ) - self.agent.reset(qpos) - self.agent.robot.set_pose(sapien.Pose([-0.82, 0, -self.table_height])) + self.env.agent.reset(qpos) + self.env.agent.robot.set_pose(sapien.Pose([-0.82, 0, -self.table_height])) from mani_skill2.agents.robots.fetch import FETCH_UNIQUE_COLLISION_BIT - cs = self.ground.find_component_by_type( + cs = self.ground._objs[0].find_component_by_type( sapien.physx.PhysxRigidStaticComponent ).get_collision_shapes()[0] cg = cs.get_collision_groups() diff --git a/mani_skill2/utils/structs/joint.py b/mani_skill2/utils/structs/joint.py index 47bacd8a0..90a4fe73d 100644 --- a/mani_skill2/utils/structs/joint.py +++ b/mani_skill2/utils/structs/joint.py @@ -34,6 +34,10 @@ class Joint(BaseStruct[physx.PhysxJointComponent]): _data_index: slice = None name: str = None + # TODO (arth): might need better hash in future but this is fine for now + def __hash__(self): + return hash(self.name) + @classmethod def create( cls, @@ -210,7 +214,9 @@ def drive_velocity_target(self) -> torch.Tensor: @drive_velocity_target.setter def drive_velocity_target(self, arg1: Array) -> None: + print('before totensor', arg1) arg1 = to_tensor(arg1) + print('after totensor', arg1) if physx.is_gpu_enabled(): raise NotImplementedError( "Cannot set drive velocity targets at the moment in GPU simulation" diff --git a/test.py b/test.py new file mode 100644 index 000000000..0eeb0f54f --- /dev/null +++ b/test.py @@ -0,0 +1,17 @@ +import gymnasium as gym +import mani_skill2.envs + +env = gym.make( + 'PickCube-v0', + robot_uid='fetch', + # render_mode='human', +) + +import numpy as np + +env.reset() +while True: + # action = env.action_space.sample() + action = np.zeros(env.action_space.shape) + env.step(action=action) + env.render() From faa5f6168527366d418cb38700e64a14cd2fc82b Mon Sep 17 00:00:00 2001 From: arth-shukla <44527268+arth-shukla@users.noreply.github.com> Date: Sun, 21 Jan 2024 18:55:58 -0800 Subject: [PATCH 5/8] Fix: gpu sim compatibility issues --- mani_skill2/agents/controllers/pd_base_vel.py | 2 +- mani_skill2/agents/robots/fetch/fetch.py | 2 +- mani_skill2/utils/sapien_utils.py | 24 +++++++++---------- test.py | 17 ------------- 4 files changed, 14 insertions(+), 31 deletions(-) delete mode 100644 test.py diff --git a/mani_skill2/agents/controllers/pd_base_vel.py b/mani_skill2/agents/controllers/pd_base_vel.py index 059ec5d70..e5322fe2b 100644 --- a/mani_skill2/agents/controllers/pd_base_vel.py +++ b/mani_skill2/agents/controllers/pd_base_vel.py @@ -33,7 +33,7 @@ def set_action(self, action: np.ndarray): new_action = np.hstack([vel, action[2:]]) for i, joint in enumerate(self.joints): - joint.set_drive_velocity_target(new_action[i]) + joint.set_drive_velocity_target(np.array([new_action[i]])) class PDBaseVelControllerConfig(PDJointVelControllerConfig): diff --git a/mani_skill2/agents/robots/fetch/fetch.py b/mani_skill2/agents/robots/fetch/fetch.py index 7b52c95cc..7c22338b6 100644 --- a/mani_skill2/agents/robots/fetch/fetch.py +++ b/mani_skill2/agents/robots/fetch/fetch.py @@ -289,7 +289,7 @@ def is_grasping(self, object: Actor = None, min_impulse=1e-6, max_angle=85): self.scene.px.gpu_create_contact_query(body_pairs), (len(object._bodies), 3), ) - print(f"Create query for Panda grasp({object.name})") + print(f"Create query for Fetch grasp({object.name})") query, contacts_shape = self.queries[object.name] self.scene.px.gpu_query_contacts(query) # query.cuda_contacts # (num_unique_pairs * num_envs, 3) diff --git a/mani_skill2/utils/sapien_utils.py b/mani_skill2/utils/sapien_utils.py index 2989f45af..e39988435 100644 --- a/mani_skill2/utils/sapien_utils.py +++ b/mani_skill2/utils/sapien_utils.py @@ -294,13 +294,13 @@ def get_pairwise_contacts( pairwise_contacts = [] for contact in contacts: if ( - contact.components[0].entity == actor0 - and contact.components[1].entity == actor1 + contact.bodies[0].entity == actor0 + and contact.bodies[1].entity == actor1 ): pairwise_contacts.append((contact, True)) elif ( - contact.components[0].entity == actor1 - and contact.components[1].entity == actor0 + contact.bodies[0].entity == actor1 + and contact.bodies[1].entity == actor0 ): pairwise_contacts.append((contact, False)) return pairwise_contacts @@ -328,9 +328,9 @@ def get_actor_contacts( ) -> List[Tuple[physx.PhysxContact, bool]]: entity_contacts = [] for contact in contacts: - if contact.components[0].entity == actor: + if contact.bodies[0].entity == actor: entity_contacts.append((contact, True)) - elif contact.components[1].entity == actor: + elif contact.bodies[1].entity == actor: entity_contacts.append((contact, False)) return entity_contacts @@ -348,16 +348,16 @@ def get_articulation_contacts( if included_links is None: included_links = links for contact in contacts: - if contact.components[0] in included_links: - if contact.components[1] in links: + if contact.bodies[0] in included_links: + if contact.bodies[1] in links: continue - if contact.components[1].entity in excluded_entities: + if contact.bodies[1].entity in excluded_entities: continue articulation_contacts.append((contact, True)) - elif contact.components[1] in included_links: - if contact.components[0] in links: + elif contact.bodies[1] in included_links: + if contact.bodies[0] in links: continue - if contact.components[0].entity in excluded_entities: + if contact.bodies[0].entity in excluded_entities: continue articulation_contacts.append((contact, False)) return articulation_contacts diff --git a/test.py b/test.py deleted file mode 100644 index 0eeb0f54f..000000000 --- a/test.py +++ /dev/null @@ -1,17 +0,0 @@ -import gymnasium as gym -import mani_skill2.envs - -env = gym.make( - 'PickCube-v0', - robot_uid='fetch', - # render_mode='human', -) - -import numpy as np - -env.reset() -while True: - # action = env.action_space.sample() - action = np.zeros(env.action_space.shape) - env.step(action=action) - env.render() From a473a0d670b54dfabbbafd26241c5851aa4e2a96 Mon Sep 17 00:00:00 2001 From: arth-shukla <44527268+arth-shukla@users.noreply.github.com> Date: Sun, 21 Jan 2024 19:26:14 -0800 Subject: [PATCH 6/8] Fix: fetch sensor cfg --- mani_skill2/agents/robots/fetch/fetch.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/mani_skill2/agents/robots/fetch/fetch.py b/mani_skill2/agents/robots/fetch/fetch.py index 7c22338b6..78a2c5287 100644 --- a/mani_skill2/agents/robots/fetch/fetch.py +++ b/mani_skill2/agents/robots/fetch/fetch.py @@ -43,7 +43,19 @@ class Fetch(BaseAgent): ), ), ) - sensor_configs = [] + sensor_configs = [ + CameraConfig( + uid="fetch_head", + p=[0, 0, 0], + q=[1, 0, 0, 0], + width=128, + height=128, + fov=1.57, + near=0.01, + far=10, + entity_uid="head_camera_link", + ) + ] def __init__(self, *args, **kwargs): self.arm_joint_names = [ From 3ea9c06a7c259c21269966479b6d67ebfe642cb2 Mon Sep 17 00:00:00 2001 From: arth-shukla <44527268+arth-shukla@users.noreply.github.com> Date: Sun, 21 Jan 2024 21:37:32 -0800 Subject: [PATCH 7/8] Fix: pd_ee_pose controller ik, env step cpu sim output, utils, fetch pose --- mani_skill2/agents/controllers/pd_ee_pose.py | 10 +++--- mani_skill2/agents/robots/fetch/fetch.py | 2 +- mani_skill2/envs/sapien_env.py | 8 +++-- mani_skill2/utils/sapien_utils.py | 34 ++++++++++++++++++-- 4 files changed, 44 insertions(+), 10 deletions(-) diff --git a/mani_skill2/agents/controllers/pd_ee_pose.py b/mani_skill2/agents/controllers/pd_ee_pose.py index d6a10cc4f..1aaeeffb8 100644 --- a/mani_skill2/agents/controllers/pd_ee_pose.py +++ b/mani_skill2/agents/controllers/pd_ee_pose.py @@ -8,7 +8,7 @@ from scipy.spatial.transform import Rotation from mani_skill2.utils.common import clip_and_scale_action -from mani_skill2.utils.sapien_utils import get_obj_by_name +from mani_skill2.utils.sapien_utils import get_obj_by_name, to_numpy, to_tensor from mani_skill2.utils.structs.pose import vectorize_pose from .base_controller import BaseController, ControllerConfig @@ -64,15 +64,17 @@ def reset(self): def compute_ik(self, target_pose, max_iterations=100): # Assume the target pose is defined in the base frame + # TODO (arth): currently ik only supports cpu, so input/output is managed as such + # in future, need to change input/output processing per gpu implementation result, success, error = self.pmodel.compute_inverse_kinematics( self.ee_link_idx, - target_pose, - initial_qpos=self.articulation.get_qpos(), + target_pose.sp, + initial_qpos=to_numpy(self.articulation.get_qpos()).squeeze(0), active_qmask=self.qmask, max_iterations=max_iterations, ) if success: - return result[self.joint_indices] + return to_tensor([result[self.joint_indices]]) else: return None diff --git a/mani_skill2/agents/robots/fetch/fetch.py b/mani_skill2/agents/robots/fetch/fetch.py index 78a2c5287..53e6a447d 100644 --- a/mani_skill2/agents/robots/fetch/fetch.py +++ b/mani_skill2/agents/robots/fetch/fetch.py @@ -385,4 +385,4 @@ def build_grasp_pose(approaching, closing, center): def tcp_pose(self) -> Pose: p = (self.finger1_link.pose.p + self.finger2_link.pose.p) / 2 q = (self.finger1_link.pose.q + self.finger2_link.pose.q) / 2 - return sapien.Pose(p=p, q=q) + return Pose.create_from_pq(p=p, q=q) diff --git a/mani_skill2/envs/sapien_env.py b/mani_skill2/envs/sapien_env.py index 33a1bffe3..36332c12c 100644 --- a/mani_skill2/envs/sapien_env.py +++ b/mani_skill2/envs/sapien_env.py @@ -32,7 +32,7 @@ get_component_meshes, merge_meshes, ) -from mani_skill2.utils.sapien_utils import get_obj_by_type, to_numpy, to_tensor +from mani_skill2.utils.sapien_utils import get_obj_by_type, to_numpy, to_tensor, unbatch from mani_skill2.utils.structs.types import Array from mani_skill2.utils.visualization.misc import observations_to_images, tile_images @@ -627,7 +627,11 @@ def step(self, action: Union[None, np.ndarray, Dict]): if self.num_envs == 1: terminated = terminated[0] reward = reward[0] - return obs, reward, terminated, False, info + + if physx.is_gpu_enabled(): + return obs, reward, terminated, torch.Tensor(False), info + else: + return unbatch(obs, reward, terminated.item(), False, to_numpy(info)) def step_action(self, action): set_action = False diff --git a/mani_skill2/utils/sapien_utils.py b/mani_skill2/utils/sapien_utils.py index e39988435..e1edb7acd 100644 --- a/mani_skill2/utils/sapien_utils.py +++ b/mani_skill2/utils/sapien_utils.py @@ -35,13 +35,17 @@ def to_tensor(array: Union[torch.Tensor, np.array, Sequence]): elif get_backend_name() == "numpy": if isinstance(array, np.ndarray): return torch.from_numpy(array) + # TODO (arth): better way to address torch "UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow" ? + elif isinstance(array, list) and isinstance(array[0], np.ndarray): + return torch.from_numpy(np.array(array)) + elif np.iterable(array): + return torch.Tensor(array) else: return torch.tensor(array) - -def to_numpy(array: Union[Array, Sequence]): +def _to_numpy(array: Union[Array, Sequence]) -> np.ndarray: if isinstance(array, (dict)): - return {k: to_numpy(v) for k, v in array.items()} + return {k: _to_numpy(v) for k, v in array.items()} if isinstance(array, str): return array if torch is not None: @@ -52,6 +56,30 @@ def to_numpy(array: Union[Array, Sequence]): else: return np.array(array) +def to_numpy(array: Union[Array, Sequence], dtype=None) -> np.ndarray: + array = _to_numpy(array) + if dtype is not None: + return array.astype(dtype) + return array + +def _unbatch(array: Union[Array, Sequence]): + if isinstance(array, (dict)): + return {k: _unbatch(v) for k, v in array.items()} + if isinstance(array, str): + return array + if torch is not None: + if isinstance(array, torch.Tensor): + return array.squeeze(0) + if isinstance(array, np.ndarray): + if np.iterable(array) and array.shape[0] == 1: + return array.squeeze(0) + if isinstance(array, list): + if len(array) == 1: + return array[0] + return array + +def unbatch(*args: Tuple[Union[Array, Sequence]]): + return tuple([_unbatch(x) for x in args]) def clone_tensor(array: Array): if torch is not None and isinstance(array, torch.Tensor): From b79389a9624337fd5a3659774fc1ee1aaef3eb65 Mon Sep 17 00:00:00 2001 From: arth-shukla <44527268+arth-shukla@users.noreply.github.com> Date: Mon, 22 Jan 2024 00:28:55 -0800 Subject: [PATCH 8/8] Fix: fetch is_grasping cpu sim --- mani_skill2/agents/robots/fetch/fetch.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mani_skill2/agents/robots/fetch/fetch.py b/mani_skill2/agents/robots/fetch/fetch.py index 53e6a447d..b8d9f00c2 100644 --- a/mani_skill2/agents/robots/fetch/fetch.py +++ b/mani_skill2/agents/robots/fetch/fetch.py @@ -329,8 +329,8 @@ def is_grasping(self, object: Actor = None, min_impulse=1e-6, max_angle=85): contacts = self.scene.get_contacts() if object is None: - finger1_contacts = get_actor_contacts(contacts, self.finger1_link) - finger2_contacts = get_actor_contacts(contacts, self.finger2_link) + finger1_contacts = get_actor_contacts(contacts, self.finger1_link._bodies[0].entity) + finger2_contacts = get_actor_contacts(contacts, self.finger2_link._bodies[0].entity) return ( np.linalg.norm(compute_total_impulse(finger1_contacts)) >= min_impulse @@ -339,10 +339,10 @@ def is_grasping(self, object: Actor = None, min_impulse=1e-6, max_angle=85): ) else: limpulse = get_pairwise_contact_impulse( - contacts, self.finger1_link, object + contacts, self.finger1_link._bodies[0].entity, object._bodies[0].entity ) rimpulse = get_pairwise_contact_impulse( - contacts, self.finger2_link, object + contacts, self.finger2_link._bodies[0].entity, object._bodies[0].entity ) # direction to open the gripper