diff --git a/mani_skill2/agents/robots/anymal/anymal_c.py b/mani_skill2/agents/robots/anymal/anymal_c.py index aa00e9273..100d76330 100644 --- a/mani_skill2/agents/robots/anymal/anymal_c.py +++ b/mani_skill2/agents/robots/anymal/anymal_c.py @@ -5,11 +5,7 @@ from mani_skill2.agents.base_agent import BaseAgent from mani_skill2.agents.controllers import * from mani_skill2.agents.registration import register_agent -from mani_skill2.utils.sapien_utils import ( - apply_urdf_config, - check_urdf_config, - parse_urdf_config, -) +from mani_skill2.utils import sapien_utils from mani_skill2.utils.structs.articulation import Articulation @@ -83,11 +79,11 @@ def _load_articulation(self): urdf_path = format_path(str(self.urdf_path)) - urdf_config = parse_urdf_config(self.urdf_config, self.scene) - check_urdf_config(urdf_config) + urdf_config = sapien_utils.parse_urdf_config(self.urdf_config, self.scene) + sapien_utils.check_urdf_config(urdf_config) # TODO(jigu): support loading multiple convex collision shapes - apply_urdf_config(loader, urdf_config) + sapien_utils.apply_urdf_config(loader, urdf_config) loader.disable_self_collisions = True self.robot: Articulation = loader.load(urdf_path) assert self.robot is not None, f"Fail to load URDF from {urdf_path}" diff --git a/mani_skill2/agents/robots/dclaw/dclaw.py b/mani_skill2/agents/robots/dclaw/dclaw.py index eb76819a7..abcd39599 100644 --- a/mani_skill2/agents/robots/dclaw/dclaw.py +++ b/mani_skill2/agents/robots/dclaw/dclaw.py @@ -8,7 +8,7 @@ from mani_skill2.agents.controllers import * from mani_skill2.agents.registration import register_agent from mani_skill2.agents.utils import get_active_joint_indices -from mani_skill2.utils.sapien_utils import get_objs_by_names +from mani_skill2.utils import sapien_utils from mani_skill2.utils.structs.joint import Joint from mani_skill2.utils.structs.link import Link from mani_skill2.utils.structs.pose import vectorize_pose @@ -52,7 +52,7 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def _after_init(self): - self.tip_links: List[Link] = get_objs_by_names( + self.tip_links: List[Link] = sapien_utils.get_objs_by_names( self.robot.get_links(), self.tip_link_names ) self.root_joints: List[Joint] = [ diff --git a/mani_skill2/agents/robots/fetch/fetch.py b/mani_skill2/agents/robots/fetch/fetch.py index da65e544a..d034f18d5 100644 --- a/mani_skill2/agents/robots/fetch/fetch.py +++ b/mani_skill2/agents/robots/fetch/fetch.py @@ -11,14 +11,8 @@ from mani_skill2.agents.controllers import * from mani_skill2.agents.registration import register_agent from mani_skill2.sensors.camera import CameraConfig +from mani_skill2.utils import sapien_utils 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, - to_tensor, -) from mani_skill2.utils.structs.actor import Actor from mani_skill2.utils.structs.base import BaseStruct from mani_skill2.utils.structs.joint import Joint @@ -332,19 +326,23 @@ def controller_configs(self): return deepcopy_dict(controller_configs) def _after_init(self): - self.finger1_link: Link = get_obj_by_name( + self.finger1_link: Link = sapien_utils.get_obj_by_name( self.robot.get_links(), "l_gripper_finger_link" ) - self.finger2_link: Link = get_obj_by_name( + self.finger2_link: Link = sapien_utils.get_obj_by_name( self.robot.get_links(), "r_gripper_finger_link" ) - self.tcp: Link = get_obj_by_name(self.robot.get_links(), self.ee_link_name) + self.tcp: Link = sapien_utils.get_obj_by_name( + self.robot.get_links(), self.ee_link_name + ) - self.base_link: Link = get_obj_by_name(self.robot.get_links(), "base_link") - self.l_wheel_link: Link = get_obj_by_name( + self.base_link: Link = sapien_utils.get_obj_by_name( + self.robot.get_links(), "base_link" + ) + self.l_wheel_link: Link = sapien_utils.get_obj_by_name( self.robot.get_links(), "l_wheel_link" ) - self.r_wheel_link: Link = get_obj_by_name( + self.r_wheel_link: Link = sapien_utils.get_obj_by_name( self.robot.get_links(), "r_wheel_link" ) for link in [self.l_wheel_link, self.r_wheel_link]: @@ -354,11 +352,11 @@ def _after_init(self): cg[2] |= FETCH_UNIQUE_COLLISION_BIT cs.set_collision_groups(cg) - self.torso_lift_link: Link = get_obj_by_name( + self.torso_lift_link: Link = sapien_utils.get_obj_by_name( self.robot.get_links(), "torso_lift_link" ) - self.head_camera_link: Link = get_obj_by_name( + self.head_camera_link: Link = sapien_utils.get_obj_by_name( self.robot.get_links(), "head_camera_link" ) @@ -405,25 +403,27 @@ 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( + finger1_contacts = sapien_utils.get_actor_contacts( contacts, self.finger1_link._bodies[0].entity ) - finger2_contacts = get_actor_contacts( + finger2_contacts = sapien_utils.get_actor_contacts( contacts, self.finger2_link._bodies[0].entity ) return ( - np.linalg.norm(compute_total_impulse(finger1_contacts)) + np.linalg.norm(sapien_utils.compute_total_impulse(finger1_contacts)) >= min_impulse - and np.linalg.norm(compute_total_impulse(finger2_contacts)) + and np.linalg.norm( + sapien_utils.compute_total_impulse(finger2_contacts) + ) >= min_impulse ) else: - limpulse = get_pairwise_contact_impulse( + limpulse = sapien_utils.get_pairwise_contact_impulse( contacts, self.finger1_link._bodies[0].entity, object._bodies[0].entity, ) - rimpulse = get_pairwise_contact_impulse( + rimpulse = sapien_utils.get_pairwise_contact_impulse( contacts, self.finger2_link._bodies[0].entity, object._bodies[0].entity, diff --git a/mani_skill2/agents/robots/panda/panda.py b/mani_skill2/agents/robots/panda/panda.py index 096f3e75f..909507f9a 100644 --- a/mani_skill2/agents/robots/panda/panda.py +++ b/mani_skill2/agents/robots/panda/panda.py @@ -10,14 +10,8 @@ from mani_skill2.agents.base_agent import BaseAgent from mani_skill2.agents.controllers import * from mani_skill2.agents.registration import register_agent -from mani_skill2.sensors.camera import CameraConfig +from mani_skill2.utils import sapien_utils 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 @@ -204,15 +198,21 @@ def controller_configs(self): sensor_configs = [] def _after_init(self): - self.finger1_link = get_obj_by_name(self.robot.get_links(), "panda_leftfinger") - self.finger2_link = get_obj_by_name(self.robot.get_links(), "panda_rightfinger") - self.finger1pad_link = get_obj_by_name( + self.finger1_link = sapien_utils.get_obj_by_name( + self.robot.get_links(), "panda_leftfinger" + ) + self.finger2_link = sapien_utils.get_obj_by_name( + self.robot.get_links(), "panda_rightfinger" + ) + self.finger1pad_link = sapien_utils.get_obj_by_name( self.robot.get_links(), "panda_leftfinger_pad" ) - self.finger2pad_link = get_obj_by_name( + self.finger2pad_link = sapien_utils.get_obj_by_name( self.robot.get_links(), "panda_rightfinger_pad" ) - self.tcp = get_obj_by_name(self.robot.get_links(), self.ee_link_name) + self.tcp = sapien_utils.get_obj_by_name( + self.robot.get_links(), self.ee_link_name + ) self.queries: Dict[ str, Tuple[physx.PhysxGpuContactPairImpulseQuery, Tuple[int]] @@ -257,19 +257,25 @@ 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 = sapien_utils.get_actor_contacts( + contacts, self.finger1_link + ) + finger2_contacts = sapien_utils.get_actor_contacts( + contacts, self.finger2_link + ) return ( - np.linalg.norm(compute_total_impulse(finger1_contacts)) + np.linalg.norm(sapien_utils.compute_total_impulse(finger1_contacts)) >= min_impulse - and np.linalg.norm(compute_total_impulse(finger2_contacts)) + and np.linalg.norm( + sapien_utils.compute_total_impulse(finger2_contacts) + ) >= min_impulse ) else: - limpulse = get_pairwise_contact_impulse( + limpulse = sapien_utils.get_pairwise_contact_impulse( contacts, self.finger1_link, object ) - rimpulse = get_pairwise_contact_impulse( + rimpulse = sapien_utils.get_pairwise_contact_impulse( contacts, self.finger2_link, object ) diff --git a/mani_skill2/agents/robots/xarm/xarm7_ability.py b/mani_skill2/agents/robots/xarm/xarm7_ability.py index d115e4bbf..70755cfc0 100644 --- a/mani_skill2/agents/robots/xarm/xarm7_ability.py +++ b/mani_skill2/agents/robots/xarm/xarm7_ability.py @@ -7,7 +7,7 @@ from mani_skill2.agents.base_agent import BaseAgent from mani_skill2.agents.controllers import * from mani_skill2.agents.registration import register_agent -from mani_skill2.utils.sapien_utils import get_obj_by_name, get_objs_by_names +from mani_skill2.utils import sapien_utils @register_agent() @@ -155,7 +155,7 @@ def _after_init(self): "ring_L2", "pinky_L2", ] - self.hand_front_links = get_objs_by_names( + self.hand_front_links = sapien_utils.get_objs_by_names( self.robot.get_links(), hand_front_link_names ) @@ -166,10 +166,12 @@ def _after_init(self): "ring_tip", "pinky_tip", ] - self.finger_tip_links = get_objs_by_names( + self.finger_tip_links = sapien_utils.get_objs_by_names( self.robot.get_links(), finger_tip_link_names ) - self.tcp = get_obj_by_name(self.robot.get_links(), self.ee_link_name) + self.tcp = sapien_utils.get_obj_by_name( + self.robot.get_links(), self.ee_link_name + ) self.queries: Dict[str, Tuple[physx.PhysxGpuContactQuery, Tuple[int]]] = dict() diff --git a/mani_skill2/agents/robots/xmate3/xmate3.py b/mani_skill2/agents/robots/xmate3/xmate3.py index f8d3b2433..a5f2b54c2 100644 --- a/mani_skill2/agents/robots/xmate3/xmate3.py +++ b/mani_skill2/agents/robots/xmate3/xmate3.py @@ -10,13 +10,8 @@ from mani_skill2.agents.controllers import * from mani_skill2.agents.registration import register_agent from mani_skill2.sensors.camera import CameraConfig +from mani_skill2.utils import sapien_utils 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 @@ -71,13 +66,15 @@ def __init__( super().__init__(scene, control_freq, control_mode, fix_root_link, agent_idx) def _after_init(self): - self.finger1_link = get_obj_by_name( + self.finger1_link = sapien_utils.get_obj_by_name( self.robot.get_links(), "left_inner_finger_pad" ) - self.finger2_link = get_obj_by_name( + self.finger2_link = sapien_utils.get_obj_by_name( self.robot.get_links(), "right_inner_finger_pad" ) - self.tcp = get_obj_by_name(self.robot.get_links(), self.ee_link_name) + self.tcp = sapien_utils.get_obj_by_name( + self.robot.get_links(), self.ee_link_name + ) self.queries: Dict[ str, Tuple[physx.PhysxGpuContactPairImpulseQuery, Tuple[int]] ] = dict() @@ -221,19 +218,25 @@ 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 = sapien_utils.get_actor_contacts( + contacts, self.finger1_link + ) + finger2_contacts = sapien_utils.get_actor_contacts( + contacts, self.finger2_link + ) return ( - np.linalg.norm(compute_total_impulse(finger1_contacts)) + np.linalg.norm(sapien_utils.compute_total_impulse(finger1_contacts)) >= min_impulse - and np.linalg.norm(compute_total_impulse(finger2_contacts)) + and np.linalg.norm( + sapien_utils.compute_total_impulse(finger2_contacts) + ) >= min_impulse ) else: - limpulse = get_pairwise_contact_impulse( + limpulse = sapien_utils.get_pairwise_contact_impulse( contacts, self.finger1_link, object ) - rimpulse = get_pairwise_contact_impulse( + rimpulse = sapien_utils.get_pairwise_contact_impulse( contacts, self.finger2_link, object ) diff --git a/mani_skill2/agents/utils.py b/mani_skill2/agents/utils.py index b7a0b1651..e972aeada 100644 --- a/mani_skill2/agents/utils.py +++ b/mani_skill2/agents/utils.py @@ -7,7 +7,7 @@ import sapien.physx as physx from gymnasium import spaces -from mani_skill2.utils.sapien_utils import to_tensor +from mani_skill2.utils import sapien_utils from mani_skill2.utils.structs.articulation import Articulation @@ -15,7 +15,7 @@ def get_active_joint_indices(articulation: Articulation, joint_names: Sequence[s """get the indices of the provided joint names from the Articulation's list of active joints""" all_joint_names = [x.name for x in articulation.get_active_joints()] joint_indices = [all_joint_names.index(x) for x in joint_names] - return to_tensor(joint_indices).int() + return sapien_utils.to_tensor(joint_indices).int() def get_joints_by_names(articulation: Articulation, joint_names: Sequence[str]): diff --git a/mani_skill2/envs/minimal_template.py b/mani_skill2/envs/minimal_template.py index 533f8dcac..449fdc8b3 100644 --- a/mani_skill2/envs/minimal_template.py +++ b/mani_skill2/envs/minimal_template.py @@ -8,8 +8,8 @@ from mani_skill2.agents.robots.panda.panda import Panda from mani_skill2.envs.sapien_env import BaseEnv from mani_skill2.sensors.camera import CameraConfig +from mani_skill2.utils import sapien_utils from mani_skill2.utils.registration import register_env -from mani_skill2.utils.sapien_utils import look_at from mani_skill2.utils.structs.types import GPUMemoryConfig, SimConfig @@ -42,13 +42,13 @@ def __init__(self, *args, robot_uids="panda", robot_init_qpos_noise=0.02, **kwar super().__init__(*args, robot_uids=robot_uids, **kwargs) def _register_sensors(self): - pose = look_at(eye=[0.3, 0, 0.6], target=[-0.1, 0, 0.1]) + pose = sapien_utils.look_at(eye=[0.3, 0, 0.6], target=[-0.1, 0, 0.1]) return [ CameraConfig("base_camera", pose.p, pose.q, 128, 128, np.pi / 2, 0.01, 10) ] def _register_human_render_cameras(self): - pose = look_at([0.6, 0.7, 0.6], [0.0, 0.0, 0.35]) + pose = sapien_utils.look_at([0.6, 0.7, 0.6], [0.0, 0.0, 0.35]) return CameraConfig("render_camera", pose.p, pose.q, 512, 512, 1, 0.01, 10) def _load_actors(self): diff --git a/mani_skill2/envs/sapien_env.py b/mani_skill2/envs/sapien_env.py index 8baf4ff82..6d1faf4c5 100644 --- a/mani_skill2/envs/sapien_env.py +++ b/mani_skill2/envs/sapien_env.py @@ -33,18 +33,12 @@ update_camera_cfgs_from_dict, ) from mani_skill2.sensors.depth_camera import StereoDepthCamera, StereoDepthCameraConfig +from mani_skill2.utils import sapien_utils from mani_skill2.utils.common import ( convert_observation_to_space, dict_merge, flatten_state_dict, ) -from mani_skill2.utils.sapien_utils import ( - batch, - get_obj_by_type, - to_numpy, - to_tensor, - unbatch, -) from mani_skill2.utils.structs.actor import Actor from mani_skill2.utils.structs.articulation import Articulation from mani_skill2.utils.structs.types import Array, SimConfig @@ -238,10 +232,10 @@ def __init__( ) obs, _ = self.reset(seed=2022, options=dict(reconfigure=True)) if physx.is_gpu_enabled(): - obs = to_numpy(obs) + obs = sapien_utils.to_numpy(obs) self._init_raw_obs = obs.copy() """the raw observation returned by the env.reset. Useful for future observation wrappers to use to auto generate observation spaces""" - self._init_raw_state = to_numpy(self.get_state_dict()) + self._init_raw_state = sapien_utils.to_numpy(self.get_state_dict()) """the initial raw state returned by env.get_state. Useful for reconstructing state dictionaries from flattened state vectors""" self.action_space = self.agent.action_space @@ -660,7 +654,7 @@ def reset(self, seed=None, options=None): self._scene._gpu_fetch_all() obs = self.get_obs() if not physx.is_gpu_enabled(): - obs = to_numpy(unbatch(obs)) + obs = sapien_utils.to_numpy(sapien_utils.unbatch(obs)) self._elapsed_steps = 0 return obs, {} @@ -759,12 +753,12 @@ def step(self, action: Union[None, np.ndarray, torch.Tensor, Dict]): ) else: # On CPU sim mode, we always return numpy / python primitives without any batching. - return unbatch( - to_numpy(obs), - to_numpy(reward), - to_numpy(terminated), + return sapien_utils.unbatch( + sapien_utils.to_numpy(obs), + sapien_utils.to_numpy(reward), + sapien_utils.to_numpy(terminated), False, - to_numpy(info), + sapien_utils.to_numpy(info), ) def step_action( @@ -775,7 +769,7 @@ def step_action( if action is None: # simulation without action pass elif isinstance(action, np.ndarray) or isinstance(action, torch.Tensor): - action = to_tensor(action) + action = sapien_utils.to_tensor(action) if action.shape == self._orig_single_action_space.shape: action_is_unbatched = True set_action = True @@ -784,13 +778,13 @@ def step_action( if action["control_mode"] != self.agent.control_mode: self.agent.set_control_mode(action["control_mode"]) self.agent.controller.reset() - action = to_tensor(action["action"]) + action = sapien_utils.to_tensor(action["action"]) else: assert isinstance( self.agent, MultiAgent ), "Received a dictionary for an action but there are not multiple robots in the environment" # assume this is a multi-agent action - action = to_tensor(action) + action = sapien_utils.to_tensor(action) for k, a in action.items(): if a.shape == self._orig_single_action_space[k].shape: action_is_unbatched = True @@ -801,7 +795,7 @@ def step_action( if set_action: if self.num_envs == 1 and action_is_unbatched: - action = batch(action) + action = sapien_utils.batch(action) self.agent.set_action(action) if physx.is_gpu_enabled(): self._scene.px.gpu_apply_articulation_target_position() @@ -990,7 +984,7 @@ def _setup_viewer(self): # CAUTION: `set_scene` should be called after assets are loaded. self._viewer.set_scene(self._scene.sub_scenes[0]) control_window: sapien.utils.viewer.control_window.ControlWindow = ( - get_obj_by_type( + sapien_utils.get_obj_by_type( self._viewer.plugins, sapien.utils.viewer.control_window.ControlWindow ) )