diff --git a/mani_skill2/agents/__init__.py b/mani_skill2/agents/__init__.py index e69de29bb..7dc854f6a 100644 --- a/mani_skill2/agents/__init__.py +++ b/mani_skill2/agents/__init__.py @@ -0,0 +1 @@ +from .registration import REGISTERED_AGENTS diff --git a/mani_skill2/agents/registration.py b/mani_skill2/agents/registration.py new file mode 100644 index 000000000..10ba3a921 --- /dev/null +++ b/mani_skill2/agents/registration.py @@ -0,0 +1,40 @@ +from dataclasses import dataclass +from typing import Dict + +from mani_skill2 import logger +from mani_skill2.agents.base_agent import BaseAgent + + +@dataclass +class AgentSpec: + """Agent specifications. At the moment it is a simple wrapper around the agent_cls but the dataclass is used in case we may need additional metadata""" + + agent_cls: type[BaseAgent] + + +REGISTERED_AGENTS: Dict[str, AgentSpec] = {} + + +def register_agent(override=False): + """A decorator to register agents into ManiSkill so they can be used easily by string uid. + + Args: + uid (str): unique id of the agent. + override (bool): whether to override the agent if it is already registered. + """ + + def _register_agent(agent_cls: type[BaseAgent]): + if agent_cls.uid in REGISTERED_AGENTS: + if override: + logger.warn(f"Overriding registered agent {agent_cls.uid}") + REGISTERED_AGENTS.pop(agent_cls.uid) + else: + logger.warn( + f"Agent {agent_cls.uid} is already registered. Skip registration." + ) + return agent_cls + + REGISTERED_AGENTS[agent_cls.uid] = AgentSpec(agent_cls=agent_cls) + return agent_cls + + return _register_agent diff --git a/mani_skill2/agents/robots/__init__.py b/mani_skill2/agents/robots/__init__.py index f3db0e3e1..4bbb8ecd8 100644 --- a/mani_skill2/agents/robots/__init__.py +++ b/mani_skill2/agents/robots/__init__.py @@ -1,24 +1,24 @@ from mani_skill2.agents.robots.anymal.anymal_c import ANYmalC + +from .allegro_hand import AllegroHandLeft, AllegroHandRight, AllegroHandRightTouch from .dclaw import DClaw from .fetch import Fetch -from .mobile_panda import MobilePandaDualArm, MobilePandaSingleArm from .panda import Panda from .xarm import XArm7Ability from .xmate3 import Xmate3Robotiq -from .allegro_hand import AllegroHandRightTouch, AllegroHandRight, AllegroHandLeft -ROBOTS = { - "panda": Panda, - "mobile_panda_dual_arm": MobilePandaDualArm, - "mobile_panda_single_arm": MobilePandaSingleArm, - "xmate3_robotiq": Xmate3Robotiq, - "fetch": Fetch, - # Dexterous Hand - "dclaw": DClaw, - "xarm7_ability": XArm7Ability, - "allegro_hand_right": AllegroHandRight, - "allegro_hand_left": AllegroHandLeft, - "allegro_hand_right_touch": AllegroHandRightTouch, - # Locomotion - "anymal-c": ANYmalC, -} +# = { +# "panda": Panda, +# "mobile_panda_dual_arm": MobilePandaDualArm, +# "mobile_panda_single_arm": MobilePandaSingleArm, +# "xmate3_robotiq": Xmate3Robotiq, +# "fetch": Fetch, +# # Dexterous Hand +# "dclaw": DClaw, +# "xarm7_ability": XArm7Ability, +# "allegro_hand_right": AllegroHandRight, +# "allegro_hand_left": AllegroHandLeft, +# "allegro_hand_right_touch": AllegroHandRightTouch, +# # Locomotion +# "anymal-c": ANYmalC, +# } diff --git a/mani_skill2/agents/robots/allegro_hand/allegro.py b/mani_skill2/agents/robots/allegro_hand/allegro.py index e9cd288f4..db0efb879 100644 --- a/mani_skill2/agents/robots/allegro_hand/allegro.py +++ b/mani_skill2/agents/robots/allegro_hand/allegro.py @@ -7,13 +7,12 @@ from mani_skill2 import PACKAGE_ASSET_DIR from mani_skill2.agents.base_agent import BaseAgent from mani_skill2.agents.controllers import * -from mani_skill2.utils.sapien_utils import ( - get_obj_by_name, -) -from mani_skill2.utils.sapien_utils import get_objs_by_names +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.structs.pose import vectorize_pose +@register_agent() class AllegroHandRight(BaseAgent): uid = "allegro_hand_right" urdf_path = f"{PACKAGE_ASSET_DIR}/robots/allegro/allegro_hand_right_glb.urdf" @@ -146,6 +145,7 @@ def palm_pose(self): return vectorize_pose(self.palm_link.pose) +@register_agent() class AllegroHandLeft(AllegroHandRight): uid = "allegro_hand_left" urdf_path = f"{PACKAGE_ASSET_DIR}/robots/allegro/allegro_hand_left.urdf" diff --git a/mani_skill2/agents/robots/allegro_hand/allegro_touch.py b/mani_skill2/agents/robots/allegro_hand/allegro_touch.py index f08c284bf..d2c976ef5 100644 --- a/mani_skill2/agents/robots/allegro_hand/allegro_touch.py +++ b/mani_skill2/agents/robots/allegro_hand/allegro_touch.py @@ -1,5 +1,5 @@ import itertools -from typing import List, Dict, Tuple, Optional +from typing import Dict, List, Optional, Tuple import numpy as np import sapien @@ -7,16 +7,18 @@ from sapien import physx from mani_skill2 import PACKAGE_ASSET_DIR +from mani_skill2.agents.registration import register_agent from mani_skill2.agents.robots.allegro_hand.allegro import AllegroHandRight from mani_skill2.utils.sapien_utils import ( compute_total_impulse, - get_multiple_pairwise_contacts, get_actors_contacts, + get_multiple_pairwise_contacts, + get_objs_by_names, ) -from mani_skill2.utils.sapien_utils import get_objs_by_names from mani_skill2.utils.structs.actor import Actor +@register_agent() class AllegroHandRightTouch(AllegroHandRight): uid = "allegro_hand_right_touch" urdf_path = f"{PACKAGE_ASSET_DIR}/robots/allegro/variation/allegro_hand_right_fsr_simple.urdf" diff --git a/mani_skill2/agents/robots/anymal/anymal_c.py b/mani_skill2/agents/robots/anymal/anymal_c.py index 1edad7887..122511ab7 100644 --- a/mani_skill2/agents/robots/anymal/anymal_c.py +++ b/mani_skill2/agents/robots/anymal/anymal_c.py @@ -3,6 +3,7 @@ from mani_skill2 import PACKAGE_ASSET_DIR, format_path 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, @@ -11,6 +12,7 @@ from mani_skill2.utils.structs.articulation import Articulation +@register_agent() class ANYmalC(BaseAgent): uid = "anymal-c" urdf_path = f"{PACKAGE_ASSET_DIR}/robots/anymal-c/urdf/anymal.urdf" diff --git a/mani_skill2/agents/robots/dclaw/dclaw.py b/mani_skill2/agents/robots/dclaw/dclaw.py index 77464e0a6..eb76819a7 100644 --- a/mani_skill2/agents/robots/dclaw/dclaw.py +++ b/mani_skill2/agents/robots/dclaw/dclaw.py @@ -6,17 +6,15 @@ from mani_skill2 import PACKAGE_ASSET_DIR from mani_skill2.agents.base_agent import BaseAgent from mani_skill2.agents.controllers import * -from mani_skill2.agents.utils import ( - get_active_joint_indices, -) -from mani_skill2.utils.sapien_utils import ( - get_objs_by_names, -) +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.structs.joint import Joint from mani_skill2.utils.structs.link import Link from mani_skill2.utils.structs.pose import vectorize_pose +@register_agent() class DClaw(BaseAgent): uid = "dclaw" urdf_path = f"{PACKAGE_ASSET_DIR}/robots/dclaw/dclaw_gripper_glb.urdf" diff --git a/mani_skill2/agents/robots/fetch/fetch.py b/mani_skill2/agents/robots/fetch/fetch.py index d2d6dcdf8..da65e544a 100644 --- a/mani_skill2/agents/robots/fetch/fetch.py +++ b/mani_skill2/agents/robots/fetch/fetch.py @@ -9,6 +9,7 @@ from mani_skill2 import PACKAGE_ASSET_DIR 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.common import compute_angle_between, np_compute_angle_between from mani_skill2.utils.sapien_utils import ( @@ -28,6 +29,7 @@ FETCH_UNIQUE_COLLISION_BIT = 1 << 30 +@register_agent() class Fetch(BaseAgent): uid = "fetch" urdf_path = f"{PACKAGE_ASSET_DIR}/robots/fetch/fetch.urdf" @@ -66,10 +68,28 @@ class Fetch(BaseAgent): near=0.01, far=10, entity_uid="gripper_link", - ) + ), ] REACHABLE_DIST = 1.5 - RESTING_QPOS = np.array([0, 0, 0, 0.386, 0, -0.370, 0.562, -1.032, 0.695, 0.955, -0.1, 2.077, 0, 0.015, 0.015]) + RESTING_QPOS = np.array( + [ + 0, + 0, + 0, + 0.386, + 0, + -0.370, + 0.562, + -1.032, + 0.695, + 0.955, + -0.1, + 2.077, + 0, + 0.015, + 0.015, + ] + ) def __init__(self, *args, **kwargs): self.arm_joint_names = [ diff --git a/mani_skill2/agents/robots/mobile_panda/__init__.py b/mani_skill2/agents/robots/mobile_panda/__init__.py deleted file mode 100644 index 0f1f5212b..000000000 --- a/mani_skill2/agents/robots/mobile_panda/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .base_mobile_agent import DummyMobileAgent -from .mobile_panda_dual_arm import MobilePandaDualArm -from .mobile_panda_single_arm import MobilePandaSingleArm diff --git a/mani_skill2/agents/robots/mobile_panda/base_mobile_agent.py b/mani_skill2/agents/robots/mobile_panda/base_mobile_agent.py deleted file mode 100644 index 3c1c08bb0..000000000 --- a/mani_skill2/agents/robots/mobile_panda/base_mobile_agent.py +++ /dev/null @@ -1,73 +0,0 @@ -""" -NOTE (stao): note that the existence of a dummy mobile agent file here suggests likely there is a lot of code for mobile agents we can abstract away / share -and avoid using a base_class.py kind of file -""" - -import numpy as np -from sapien import Pose -from transforms3d.euler import euler2quat - -from mani_skill2.agents.base_agent import BaseAgent -from mani_skill2.agents.controllers import * -from mani_skill2.utils.sapien_utils import get_obj_by_name - - -class DummyMobileAgent(BaseAgent): - def __init__( - self, scene, control_freq, control_mode=None, fix_root_link=True, agent_idx=None - ): - if control_mode is None: # if user did not specify a control_mode - control_mode = "base_pd_joint_vel_arm_pd_joint_vel" - super().__init__( - scene, - control_freq, - control_mode=control_mode, - fix_root_link=fix_root_link, - agent_idx=agent_idx, - ) - - def _after_init(self): - super()._after_init() - - # Sanity check - active_joints = self.robot.get_active_joints() - assert active_joints[0].name == "root_x_axis_joint" - assert active_joints[1].name == "root_y_axis_joint" - assert active_joints[2].name == "root_z_rotation_joint" - - # Dummy base - self.base_link = self.robot.get_links()[3] - - # Ignore collision between the adjustable body and ground - bodies = get_obj_by_name(self.robot.get_links(), "adjustable_body") - for body in bodies._objs: - s = body.get_collision_shapes()[0] - gs = s.get_collision_groups() - gs[2] = 1 << 30 - s.set_collision_groups(gs) - - def get_proprioception(self): - state_dict = super().get_proprioception() - qpos, qvel = state_dict["qpos"], state_dict["qvel"] - base_pos, base_orientation, arm_qpos = qpos[:, :2], qpos[:, 2], qpos[:, 3:] - base_vel, base_ang_vel, arm_qvel = qvel[:, :2], qvel[:, 2], qvel[:, 3:] - - state_dict["qpos"] = arm_qpos - state_dict["qvel"] = arm_qvel - state_dict["base_pos"] = base_pos - state_dict["base_orientation"] = base_orientation - state_dict["base_vel"] = base_vel - state_dict["base_ang_vel"] = base_ang_vel - return state_dict - - @property - def base_pose(self): - qpos = self.robot.get_qpos() - x, y, ori = qpos[:3] - return Pose([x, y, 0], euler2quat(0, 0, ori)) - - def set_base_pose(self, xy, ori): - qpos = self.robot.get_qpos() - qpos[0:2] = xy - qpos[2] = ori - self.robot.set_qpos(qpos) diff --git a/mani_skill2/agents/robots/mobile_panda/mobile_panda_dual_arm.py b/mani_skill2/agents/robots/mobile_panda/mobile_panda_dual_arm.py deleted file mode 100644 index c63120b25..000000000 --- a/mani_skill2/agents/robots/mobile_panda/mobile_panda_dual_arm.py +++ /dev/null @@ -1,320 +0,0 @@ -from copy import deepcopy - -import numpy as np -from sapien import Pose - -from mani_skill2 import PACKAGE_ASSET_DIR -from mani_skill2.agents.controllers import * -from mani_skill2.sensors.camera import CameraConfig -from mani_skill2.utils.sapien_utils import get_objs_by_names - -from .base_mobile_agent import DummyMobileAgent - - -class MobilePandaDualArm(DummyMobileAgent): - uid = "mobile_panda_dual_arm" - urdf_path = f"{PACKAGE_ASSET_DIR}/robots/panda/mobile_panda_dual_arm.urdf" - urdf_config = dict( - _materials=dict( - gripper=dict(static_friction=2.0, dynamic_friction=2.0, restitution=0.0) - ), - link=dict( - right_panda_leftfinger=dict( - material="gripper", patch_radius=0.1, min_patch_radius=0.1 - ), - right_panda_rightfinger=dict( - material="gripper", patch_radius=0.1, min_patch_radius=0.1 - ), - left_panda_leftfinger=dict( - material="gripper", patch_radius=0.1, min_patch_radius=0.1 - ), - left_panda_rightfinger=dict( - material="gripper", patch_radius=0.1, min_patch_radius=0.1 - ), - ), - ) - - def __init__(self, scene, control_freq, control_mode=None, fix_root_link=True): - self.base_joint_names = [ - "root_x_axis_joint", - "root_y_axis_joint", - "root_z_rotation_joint", - "linear_actuator_height", - ] - - arm_joint_names = [ - "panda_joint1", - "panda_joint2", - "panda_joint3", - "panda_joint4", - "panda_joint5", - "panda_joint6", - "panda_joint7", - ] - self.arm_joint_names = { - "left": ["left_" + x for x in arm_joint_names], - "right": ["right_" + x for x in arm_joint_names], - } - self.arm_stiffness = 1e3 - self.arm_damping = 1e2 - self.arm_force_limit = 100 - self.arm_joint_delta = 0.1 - self.arm_ee_delta = 0.1 - - gripper_joint_names = [ - "panda_finger_joint1", - "panda_finger_joint2", - ] - self.gripper_joint_names = { - "left": ["left_" + x for x in gripper_joint_names], - "right": ["right_" + x for x in gripper_joint_names], - } - self.gripper_stiffness = 1e3 - self.gripper_damping = 1e2 - self.gripper_force_limit = 100 - - self.ee_link_name = { - "left": "left_panda_hand_tcp", - "right": "right_panda_hand_tcp", - } - - self.camera_h = 1.5 - super().__init__(scene, control_freq, control_mode, fix_root_link) - - def _after_init(self): - super()._after_init() - - ( - self.rfinger1_joint, - self.rfinger2_joint, - self.lfinger1_joint, - self.lfinger2_joint, - ) = get_objs_by_names( - self.robot.get_joints(), - [ - "right_panda_finger_joint1", - "right_panda_finger_joint2", - "left_panda_finger_joint1", - "left_panda_finger_joint2", - ], - ) - ( - self.rfinger1_link, - self.rfinger2_link, - self.lfinger1_link, - self.lfinger2_link, - ) = get_objs_by_names( - self.robot.get_links(), - [ - "right_panda_leftfinger", - "right_panda_rightfinger", - "left_panda_leftfinger", - "left_panda_rightfinger", - ], - ) - - self.rhand, self.lhand = get_objs_by_names( - self.robot.get_links(), ["right_panda_hand", "left_panda_hand"] - ) - - # Define some useful functions for this robot - def get_fingers_info(self): - fingers_pos = self.get_ee_coords().flatten() - fingers_vel = self.get_ee_vels().flatten() - return { - "fingers_pos": fingers_pos, - "fingers_vel": fingers_vel, - } - - def get_ee_coords(self): - finger_tips = [ - (self.rfinger2_joint.get_global_pose() * Pose([0, 0.035, 0])).p, - (self.rfinger1_joint.get_global_pose() * Pose([0, -0.035, 0])).p, - (self.lfinger2_joint.get_global_pose() * Pose([0, 0.035, 0])).p, - (self.lfinger1_joint.get_global_pose() * Pose([0, -0.035, 0])).p, - ] - return np.array(finger_tips) - - def get_ee_vels(self): - finger_vels = [ - self.rfinger1_link.get_linear_velocity(), - self.rfinger2_link.get_linear_velocity(), - self.lfinger1_link.get_linear_velocity(), - self.lfinger2_link.get_linear_velocity(), - ] - return np.array(finger_vels) - - @property - def controller_configs(self): - _C = {} # controller configs for each component - - # -------------------------------------------------------------------------- # - # Mobile Base - # -------------------------------------------------------------------------- # - # fmt: off - _C["base"] = dict( - # PD ego-centric joint velocity - base_pd_joint_vel=PDBaseVelControllerConfig( - self.base_joint_names, - lower=[-0.5, -0.5, -3.14, -0.5], - upper=[0.5, 0.5, 3.14, 0.5], - damping=1000, - force_limit=500, - ) - ) - # fmt: on - - sides = list(self.arm_joint_names.keys()) - - # -------------------------------------------------------------------------- # - # Arm - # -------------------------------------------------------------------------- # - for side in sides: - # PD joint velocity - arm_pd_joint_vel = PDJointVelControllerConfig( - self.arm_joint_names[side], - lower=-3.14, - upper=3.14, - damping=self.arm_damping, - force_limit=self.arm_force_limit, - ) - - # PD joint position - arm_pd_joint_delta_pos = PDJointPosControllerConfig( - self.arm_joint_names[side], - -self.arm_joint_delta, - self.arm_joint_delta, - stiffness=self.arm_stiffness, - damping=self.arm_damping, - force_limit=self.arm_force_limit, - use_delta=True, - ) - arm_pd_joint_pos = PDJointPosControllerConfig( - self.arm_joint_names[side], - None, - None, - stiffness=self.arm_stiffness, - damping=self.arm_damping, - force_limit=self.arm_force_limit, - normalize_action=False, - ) - arm_pd_joint_target_delta_pos = deepcopy(arm_pd_joint_delta_pos) - arm_pd_joint_target_delta_pos.use_target = True - - # PD ee pose - arm_pd_ee_delta_pos = PDEEPosControllerConfig( - self.arm_joint_names[side], - -self.arm_ee_delta, - self.arm_ee_delta, - stiffness=self.arm_stiffness, - damping=self.arm_damping, - force_limit=self.arm_force_limit, - ee_link=self.ee_link_name[side], - ) - arm_pd_ee_delta_pose = PDEEPoseControllerConfig( - self.arm_joint_names[side], - -self.arm_ee_delta, - self.arm_ee_delta, - rot_bound=self.arm_ee_delta, - stiffness=self.arm_stiffness, - damping=self.arm_damping, - force_limit=self.arm_force_limit, - ee_link=self.ee_link_name[side], - ) - - arm_pd_ee_target_delta_pos = deepcopy(arm_pd_ee_delta_pos) - arm_pd_ee_target_delta_pos.use_target = True - arm_pd_ee_target_delta_pose = deepcopy(arm_pd_ee_delta_pose) - arm_pd_ee_target_delta_pose.use_target = True - - # PD joint position and velocity - arm_pd_joint_pos_vel = PDJointPosVelControllerConfig( - self.arm_joint_names[side], - None, - None, - stiffness=self.arm_stiffness, - damping=self.arm_damping, - force_limit=self.arm_force_limit, - normalize_action=False, - ) - arm_pd_joint_delta_pos_vel = PDJointPosVelControllerConfig( - self.arm_joint_names[side], - -self.arm_joint_delta, - self.arm_joint_delta, - stiffness=self.arm_stiffness, - damping=self.arm_damping, - force_limit=self.arm_force_limit, - use_delta=True, - ) - - _C[side + "_arm"] = dict( - arm_pd_joint_vel=arm_pd_joint_vel, - arm_pd_joint_delta_pos=arm_pd_joint_delta_pos, - arm_pd_joint_pos=arm_pd_joint_pos, - arm_pd_joint_target_delta_pos=arm_pd_joint_target_delta_pos, - arm_pd_ee_delta_pos=arm_pd_ee_delta_pos, - arm_pd_ee_delta_pose=arm_pd_ee_delta_pose, - arm_pd_ee_target_delta_pos=arm_pd_ee_target_delta_pos, - arm_pd_ee_target_delta_pose=arm_pd_ee_target_delta_pose, - arm_pd_joint_pos_vel=arm_pd_joint_pos_vel, - arm_pd_joint_delta_pos_vel=arm_pd_joint_delta_pos_vel, - ) - - # -------------------------------------------------------------------------- # - # Gripper - # -------------------------------------------------------------------------- # - for side in sides: - _C[side + "_gripper"] = dict( - gripper_pd_joint_pos=PDJointPosMimicControllerConfig( - self.gripper_joint_names[side], - -0.01, # a trick to have force when the object is thin - 0.04, - self.gripper_stiffness, - self.gripper_damping, - self.gripper_force_limit, - ), - ) - - controller_configs = {} - for base_controller_name in _C["base"]: - # Assume right arm always be there - for arm_controller_name in _C["right_arm"]: - c = {"base": _C["base"][base_controller_name]} - for side in sides: - c[side + "_arm"] = _C[side + "_arm"][arm_controller_name] - # Assume gripper_pd_joint_pos only - c[side + "_gripper"] = _C[side + "_gripper"]["gripper_pd_joint_pos"] - combined_name = base_controller_name + "_" + arm_controller_name - controller_configs[combined_name] = c - - # Make a deepcopy in case users modify any config - return deepcopy_dict(controller_configs) - - sensor_configs = [] - - # TODO (stao): Remove this @property and make sensor configs completely statically defined. - # The expectation is that a robot should not physically ever be changed usually. If you do - # want to adapt the robot, you should inherit a robot class and make appropriate changes - # @property - # def sensor_configs(self): - # sensors = [] - # qs = [ - # [0.9238795, 0, 0.3826834, 0], - # [0.46193977, 0.33141357, 0.19134172, -0.80010315], - # [-0.46193977, 0.33141357, -0.19134172, -0.80010315], - # ] - # for i in range(3): - # q = qs[i] - # camera = CameraConfig( - # f"overhead_camera_{i}", - # p=[0, 0, self.camera_h], - # q=q, - # width=400, - # height=160, - # near=0.1, - # far=10, - # fov=np.pi / 3, - # entity_uid="mobile_base", - # ) - # sensors.append(camera) - # return sensors diff --git a/mani_skill2/agents/robots/mobile_panda/mobile_panda_single_arm.py b/mani_skill2/agents/robots/mobile_panda/mobile_panda_single_arm.py deleted file mode 100644 index 5fa7d32fb..000000000 --- a/mani_skill2/agents/robots/mobile_panda/mobile_panda_single_arm.py +++ /dev/null @@ -1,100 +0,0 @@ -import numpy as np -import sapien.physx as physx -from sapien import Pose - -from mani_skill2 import PACKAGE_ASSET_DIR -from mani_skill2.utils.sapien_utils import get_obj_by_name, get_objs_by_names - -from .mobile_panda_dual_arm import MobilePandaDualArm - - -class MobilePandaSingleArm(MobilePandaDualArm): - # single arm robot shares the same controllers and sensor config code - uid = "mobile_panda_single_arm" - urdf_path = f"{PACKAGE_ASSET_DIR}/robots/panda/mobile_panda_single_arm.urdf" - urdf_config = dict( - _materials=dict( - gripper=dict(static_friction=2.0, dynamic_friction=2.0, restitution=0.0) - ), - link=dict( - right_panda_leftfinger=dict( - material="gripper", patch_radius=0.1, min_patch_radius=0.1 - ), - right_panda_rightfinger=dict( - material="gripper", patch_radius=0.1, min_patch_radius=0.1 - ), - ), - ) - - def _after_loading_articulation(self): - self.arm_joint_names = { - "right": self.arm_joint_names["right"], - } - self.gripper_joint_names = { - "right": self.gripper_joint_names["right"], - } - self.ee_link_name = { - "right": self.ee_link_name["right"], - } - - def _after_init(self): - super()._after_init() - - self.finger1_joint, self.finger2_joint = get_objs_by_names( - self.robot.get_joints(), - ["right_panda_finger_joint1", "right_panda_finger_joint2"], - ) - self.finger1_link, self.finger2_link = get_objs_by_names( - self.robot.get_links(), - ["right_panda_leftfinger", "right_panda_rightfinger"], - ) - self.hand: physx.PhysxArticulationLinkComponent = get_obj_by_name( - self.robot.get_links(), "right_panda_hand" - ) - - def get_fingers_info(self): - fingers_pos = self.get_ee_coords().flatten() - fingers_vel = self.get_ee_vels().flatten() - return { - "fingers_pos": fingers_pos, - "fingers_vel": fingers_vel, - } - - def get_ee_coords(self): - finger_tips = [ - (self.finger2_joint.get_global_pose() * Pose([0, 0.035, 0])).p, - (self.finger1_joint.get_global_pose() * Pose([0, -0.035, 0])).p, - ] - return np.array(finger_tips) - - def get_ee_vels(self): - finger_vels = [ - self.finger2_link.get_linear_velocity(), - self.finger1_link.get_linear_velocity(), - ] - return np.array(finger_vels) - - def get_ee_coords_sample(self): - l = 0.0355 - r = 0.052 - ret = [] - for i in range(10): - x = (l * i + (4 - i) * r) / 4 - finger_tips = [ - (self.finger2_joint.get_global_pose() * Pose([0, x, 0])).p, - (self.finger1_joint.get_global_pose() * Pose([0, -x, 0])).p, - ] - ret.append(finger_tips) - return np.array(ret).transpose((1, 0, 2)) - - @staticmethod - def build_grasp_pose(approaching, closing, center): - """Build a grasp pose (panda_hand).""" - assert np.abs(1 - np.linalg.norm(approaching)) < 1e-3 - assert np.abs(1 - np.linalg.norm(closing)) < 1e-3 - assert np.abs(approaching @ closing) <= 1e-3 - ortho = np.cross(closing, approaching) - T = np.eye(4) - T[:3, :3] = np.stack([ortho, closing, approaching], axis=1) - T[:3, 3] = center - return Pose(T) diff --git a/mani_skill2/agents/robots/panda/panda.py b/mani_skill2/agents/robots/panda/panda.py index 9bdc00c3d..096f3e75f 100644 --- a/mani_skill2/agents/robots/panda/panda.py +++ b/mani_skill2/agents/robots/panda/panda.py @@ -9,6 +9,7 @@ from mani_skill2 import PACKAGE_ASSET_DIR 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.common import compute_angle_between, np_compute_angle_between from mani_skill2.utils.sapien_utils import ( @@ -20,6 +21,7 @@ from mani_skill2.utils.structs.actor import Actor +@register_agent() class Panda(BaseAgent): uid = "panda" urdf_path = f"{PACKAGE_ASSET_DIR}/robots/panda/panda_v2.urdf" diff --git a/mani_skill2/agents/robots/panda/panda_realsensed435.py b/mani_skill2/agents/robots/panda/panda_realsensed435.py index a34f61610..726f600ab 100644 --- a/mani_skill2/agents/robots/panda/panda_realsensed435.py +++ b/mani_skill2/agents/robots/panda/panda_realsensed435.py @@ -1,9 +1,11 @@ from mani_skill2 import PACKAGE_ASSET_DIR +from mani_skill2.agents.registration import register_agent from mani_skill2.sensors.camera import CameraConfig from .panda import Panda +@register_agent() class PandaRealSensed435(Panda): uid = "panda_realsensed435" urdf_path = f"{PACKAGE_ASSET_DIR}/robots/panda/panda_v3.urdf" diff --git a/mani_skill2/agents/robots/xarm/xarm7_ability.py b/mani_skill2/agents/robots/xarm/xarm7_ability.py index 448f7f95c..d115e4bbf 100644 --- a/mani_skill2/agents/robots/xarm/xarm7_ability.py +++ b/mani_skill2/agents/robots/xarm/xarm7_ability.py @@ -6,12 +6,11 @@ from mani_skill2 import PACKAGE_ASSET_DIR from mani_skill2.agents.base_agent import BaseAgent from mani_skill2.agents.controllers import * -from mani_skill2.utils.sapien_utils import ( - get_obj_by_name, - get_objs_by_names, -) +from mani_skill2.agents.registration import register_agent +from mani_skill2.utils.sapien_utils import get_obj_by_name, get_objs_by_names +@register_agent() class XArm7Ability(BaseAgent): uid = "xarm7_ability" urdf_path = f"{PACKAGE_ASSET_DIR}/robots/xarm7/xarm7_ability_right_hand.urdf" diff --git a/mani_skill2/agents/robots/xmate3/xmate3.py b/mani_skill2/agents/robots/xmate3/xmate3.py index f76f62064..f8d3b2433 100644 --- a/mani_skill2/agents/robots/xmate3/xmate3.py +++ b/mani_skill2/agents/robots/xmate3/xmate3.py @@ -8,6 +8,7 @@ from mani_skill2 import ASSET_DIR 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.common import compute_angle_between, np_compute_angle_between from mani_skill2.utils.sapien_utils import ( @@ -19,6 +20,7 @@ from mani_skill2.utils.structs.actor import Actor +@register_agent() class Xmate3Robotiq(BaseAgent): uid = "xmate3_robotiq" urdf_path = f"{ASSET_DIR}/robots/xmate3_robotiq/xmate3_robotiq.urdf" diff --git a/mani_skill2/envs/__init__.py b/mani_skill2/envs/__init__.py index e9c81a4a5..d3467a5a6 100644 --- a/mani_skill2/envs/__init__.py +++ b/mani_skill2/envs/__init__.py @@ -1,4 +1,2 @@ -# from .ms1 import * -from .ms2 import * from .scenes import * from .tasks import * diff --git a/mani_skill2/envs/ms1/README.md b/mani_skill2/envs/ms1/README.md deleted file mode 100644 index 18bd95ffe..000000000 --- a/mani_skill2/envs/ms1/README.md +++ /dev/null @@ -1,11 +0,0 @@ -# ManiSkill1 Environments - -We migrate 4 [ManiSkill1](https://github.com/haosulab/ManiSkill) environments to ManiSkill2. - -The major changes that can change behaviors are listed as follows - -- Assets are refined with an advanced convex decomposition method [CoACD](https://github.com/SarahWeiii/CoACD). -- The robot URDF is updated with official mass and inertia, as well as finger collisions. -- The success metric does not depend on multiple timesteps. -- The GT segmentation masks are removed from observations. -- If the model does not change, the simulation scene will be reused when the env is reset. diff --git a/mani_skill2/envs/ms1/__init__.py b/mani_skill2/envs/ms1/__init__.py deleted file mode 100644 index 05ee49c82..000000000 --- a/mani_skill2/envs/ms1/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from . import move_bucket, open_cabinet_door_drawer, push_chair diff --git a/mani_skill2/envs/ms1/base_env.py b/mani_skill2/envs/ms1/base_env.py deleted file mode 100644 index 8c21f6cd8..000000000 --- a/mani_skill2/envs/ms1/base_env.py +++ /dev/null @@ -1,216 +0,0 @@ -from collections import OrderedDict -from pathlib import Path -from typing import Dict, List, Tuple - -import numpy as np -import sapien -import sapien.physx as physx -from sapien import Pose - -from mani_skill2 import format_path -from mani_skill2.agents.robots.mobile_panda import DummyMobileAgent -from mani_skill2.envs.sapien_env import BaseEnv -from mani_skill2.sensors.camera import CameraConfig -from mani_skill2.utils.building.ground import build_ground -from mani_skill2.utils.common import random_choice -from mani_skill2.utils.io_utils import load_json -from mani_skill2.utils.sapien_utils import ( - apply_urdf_config, - get_actor_state, - get_articulation_padded_state, - parse_urdf_config, -) -from mani_skill2.utils.structs.pose import vectorize_pose - - -class MS1BaseEnv(BaseEnv): - DEFAULT_MODEL_JSON: str - ASSET_UID: str - agent: DummyMobileAgent - - def __init__( - self, - *args, - asset_root: str = "{ASSET_DIR}/partnet_mobility/dataset", - model_json: str = None, - model_ids: List[str] = (), - **kwargs, - ): - self.asset_root = Path(format_path(asset_root)) - - if model_json is None: - model_json = self.DEFAULT_MODEL_JSON - model_json = format_path(model_json) - self.model_db: Dict[str, Dict] = load_json(model_json) - - if isinstance(model_ids, str): - model_ids = [model_ids] - if len(model_ids) == 0: - model_ids = sorted(self.model_db.keys()) - assert len(model_ids) > 0, model_json - - self.model_ids = model_ids - self.model_id = None - - self.model_urdf_paths = {} - for model_id in self.model_ids: - self.model_urdf_paths[model_id] = self.find_urdf_path(model_id) - - super().__init__(*args, **kwargs) - - def find_urdf_path(self, model_id): - model_dir = self.asset_root / str(model_id) - - urdf_names = ["mobility_cvx.urdf", "mobility_fixed.urdf"] - for urdf_name in urdf_names: - urdf_path = model_dir / urdf_name - if urdf_path.exists(): - return urdf_path - - raise FileNotFoundError( - f"No valid URDF is found for {model_id}." - "Please download Partnet-Mobility (ManiSkill2):" - "`python -m mani_skill2.utils.download_asset {}`".format(self.ASSET_UID) - ) - - # -------------------------------------------------------------------------- # - # Reset - # -------------------------------------------------------------------------- # - def reset(self, seed=None, options=None): - if options is None: - options = dict() - self._prev_actor_pose = None - self._set_episode_rng(seed) - model_id = options.pop("model_id", None) - reconfigure = options.pop("reconfigure", False) - _reconfigure = self._set_model(model_id) - reconfigure = _reconfigure or reconfigure - options["reconfigure"] = reconfigure - return super().reset(seed=self._episode_seed, options=options) - - def _set_model(self, model_id): - """Set the model id. If not provided, choose one randomly.""" - reconfigure = False - - # Model ID - if model_id is None: - model_id = random_choice(self.model_ids, self._episode_rng) - if model_id != self.model_id: - reconfigure = True - self.model_id = model_id - self.model_info = self.model_db[self.model_id] - - return reconfigure - - def _load_actors(self): - # Create a collision ground plane - ground = build_ground(self._scene) - # TODO (stao): This is quite hacky. Future we expect the robot to be an actual well defined robot without needing to intersect the ground. We should probably deprecate the old ms1 envs eventually - # Specify a collision (ignore) group to avoid collision with robot torso - cs = ground.find_component_by_type( - physx.PhysxRigidStaticComponent - ).get_collision_shapes()[0] - cg = cs.get_collision_groups() - cg[2] = cg[2] | 1 << 30 - cs.set_collision_groups(cg) - - def _load_partnet_mobility( - self, fix_root_link=True, scale=1.0, urdf_config: dict = None - ): - loader = self._scene.create_urdf_loader() - loader.fix_root_link = fix_root_link - loader.scale = scale - loader.load_multiple_collisions_from_file = True - - urdf_path = self.model_urdf_paths[self.model_id] - urdf_config = parse_urdf_config(urdf_config or {}, self._scene) - apply_urdf_config(loader, urdf_config) - articulation: physx.PhysxArticulation = loader.load(str(urdf_path)) - return articulation - - def _register_render_cameras(self): - p = [-1.5, 0, 1.5] - q = [0.9238795, 0, 0.3826834, 0] - return CameraConfig("render_camera", p, q, 512, 512, 1, 0.01, 10) - - def _setup_viewer(self): - super()._setup_viewer() - self._viewer.set_camera_xyz(-2, 0, 3) - self._viewer.set_camera_rpy(0, -0.8, 0) - - # -------------------------------------------------------------------------- # - # Success - # -------------------------------------------------------------------------- # - def check_link_static( - self, link: physx.PhysxArticulationLinkComponent, max_v=None, max_ang_v=None - ): - """Check whether the link is static by finite difference. - Note that the angular velocity is normalized by pi due to legacy issues. - """ - from mani_skill2.utils.geometry import angle_distance - - pose = link.get_pose() - - if self._elapsed_steps <= 1: - flag_v = (max_v is None) or ( - np.linalg.norm(link.get_linear_velocity()) <= max_v - ) - flag_ang_v = (max_ang_v is None) or ( - np.linalg.norm(link.get_angular_velocity()) <= max_ang_v - ) - else: - dt = 1.0 / self._control_freq - flag_v = (max_v is None) or ( - np.linalg.norm(pose.p - self._prev_actor_pose.p) <= max_v * dt - ) - flag_ang_v = (max_ang_v is None) or ( - angle_distance(self._prev_actor_pose, pose) <= max_ang_v * dt - ) - - # CAUTION: carefully deal with it for MPC - self._prev_actor_pose = pose - return flag_v and flag_ang_v - - # ---------------------------------------------------------------------------- # - # Observation - # ---------------------------------------------------------------------------- # - def _get_obs_agent(self): - obs = super()._get_obs_agent() - if self._obs_mode not in ["state", "state_dict"]: - obs["base_pose"] = vectorize_pose(self.agent.base_pose) - return obs - - def _get_obs_extra(self) -> OrderedDict: - obs = OrderedDict() - if self._obs_mode in ["state", "state_dict"]: - obs.update(self._get_obs_priviledged()) - return obs - - def _get_obs_priviledged(self): - obs = OrderedDict() - actors = self._get_task_actors() - if len(actors) > 0: - actors_state = np.hstack([get_actor_state(actor) for actor in actors]) - obs["actors"] = actors_state - - arts_and_dofs = self._get_task_articulations() - arts_state = np.hstack( - [ - get_articulation_padded_state(art, max_dof) - for art, max_dof in arts_and_dofs - ] - ) - obs["articulations"] = arts_state - - obs.update(self.agent.get_fingers_info()) - return obs - - def _get_task_actors(self) -> List[sapien.Entity]: - """Get task-relevant actors (for privileged states).""" - return [] - - def _get_task_articulations(self) -> List[Tuple[physx.PhysxArticulation, int]]: - """Get task-relevant articulations (for privileged states). - Each element is (art_obj, max_dof). - """ - return [] diff --git a/mani_skill2/envs/ms1/move_bucket.py b/mani_skill2/envs/ms1/move_bucket.py deleted file mode 100644 index ad02a9446..000000000 --- a/mani_skill2/envs/ms1/move_bucket.py +++ /dev/null @@ -1,429 +0,0 @@ -import numpy as np -import sapien -import sapien.physx as physx -import sapien.render -import trimesh -from sapien import Pose -from scipy.spatial import distance as sdist -from transforms3d.euler import euler2quat -from transforms3d.quaternions import quat2mat - -from mani_skill2.agents.robots.mobile_panda import MobilePandaDualArm -from mani_skill2.utils.common import np_random -from mani_skill2.utils.geometry import ( - angle_between_vec, - get_axis_aligned_bbox_for_articulation, - get_local_axis_aligned_bbox_for_link, - transform_points, -) -from mani_skill2.utils.geometry.trimesh_utils import get_actor_visual_mesh -from mani_skill2.utils.registration import register_env -from mani_skill2.utils.sapien_utils import get_obj_by_name -from mani_skill2.utils.structs.pose import vectorize_pose - -from .base_env import MS1BaseEnv - - -@register_env("MoveBucket-v1", max_episode_steps=200) -class MoveBucketEnv(MS1BaseEnv): - DEFAULT_MODEL_JSON = ( - "{PACKAGE_ASSET_DIR}/partnet_mobility/meta/info_bucket_train.json" - ) - ASSET_UID = "partnet_mobility_bucket" - - def _register_render_cameras(self): - cam_cfg = super()._register_render_cameras() - cam_cfg.p = [0, 0, 4] - cam_cfg.q = [0.70710678, 0.0, 0.70710678, 0.0] - return cam_cfg - - # -------------------------------------------------------------------------- # - # Reconfigure - # -------------------------------------------------------------------------- # - def _load_articulations(self): - urdf_config = dict( - material=dict(static_friction=0.5, dynamic_friction=0.5, restitution=0), - density=1000, - ) - self.bucket = self._load_partnet_mobility( - fix_root_link=False, - scale=self.model_info.get("scale", 1), - urdf_config=urdf_config, - ) - self.bucket.set_name(self.model_id) - - # Restrict the handle's range of motion - lim = self.bucket.get_active_joints()[0].get_limit() - v = (lim[0, 1] - lim[0, 0]) * 0.1 - lim[0, 0] += v - lim[0, 1] -= v - self.bucket.get_active_joints()[0].set_limit(lim) - - if self._reward_mode in ["dense", "normalized_dense"]: - self._set_bucket_links_mesh() - - def _set_bucket_links_mesh(self): - self.links_info = {} - for link in self.bucket.get_links(): - mesh = get_actor_visual_mesh(link.entity) - if mesh is None: - continue - self.links_info[link.name] = [link, mesh] - - def _load_actors(self): - super()._load_actors() - - # place a target platform on ground - box_half_size = [0.3, 0.3, 0.1] - builder = self._scene.create_actor_builder() - - white_diffuse = self._renderer.create_material() - white_diffuse.base_color = [0.8, 0.8, 0.8, 1] - white_diffuse.metallic = 0.0 - white_diffuse.roughness = 0.9 - white_diffuse.specular = 0 - builder.add_box_visual(half_size=box_half_size, material=white_diffuse) - - obj_material = self._scene.create_physical_material(0.5, 0.5, 0) - builder.add_box_collision( - half_size=box_half_size, material=obj_material, density=1000 - ) - self.target_platform = builder.set_physx_body_type("static").build( - name="target_platform" - ) - - # balls - R = 0.05 - self.balls_radius = R - builder = self._scene.create_actor_builder() - builder.add_sphere_collision(radius=R, density=1000) - builder.add_sphere_visual( - radius=R, material=sapien.render.RenderMaterial(base_color=[0, 1, 1, 1]) - ) - self.balls = [] - self.GX = self.GY = self.GZ = 1 - for i in range(self.GX * self.GY * self.GZ): - actor = builder.build(name="ball_{:d}".format(i + 1)) - self.balls.append(actor) - - def _load_agent(self): - self.agent = MobilePandaDualArm( - self._scene, self._control_freq, self._control_mode - ) - - links = self.agent.robot.get_links() - self.left_tcp: physx.PhysxArticulationLinkComponent = get_obj_by_name( - links, "left_panda_hand_tcp" - ) - self.right_tcp: physx.PhysxArticulationLinkComponent = get_obj_by_name( - links, "right_panda_hand_tcp" - ) - - # -------------------------------------------------------------------------- # - # Reset - # -------------------------------------------------------------------------- # - def _initialize_task(self): - self.root_link = self.bucket.get_links()[0] - - self._set_target() - self._initialize_bucket() - self._initialize_robot() - self._initialize_balls() - if self._reward_mode in ["dense", "normalized_dense"]: - self._set_bucket_links_pcd() - - for _ in range(25): - self._scene.step() - - def _set_target(self): - self.target_xy = np.zeros(2) - target_orientation = 0 - target_q = euler2quat(target_orientation, 0, 0, "szyx") - self.target_p = np.zeros(3) - self.target_p[:2] = self.target_xy - self.target_platform.set_pose(Pose(p=self.target_p, q=target_q)) - - def _initialize_bucket(self): - pose = Pose(p=[0, 0, 2], q=[1, 0, 0, 0]) - self.bucket.set_pose(pose) - bb = np.array( - get_axis_aligned_bbox_for_articulation(self.bucket) - ) # bb in world - - # Sample a position around the target - center = self.target_xy - dist = self._episode_rng.uniform(low=0.8, high=1.2) - theta = self._episode_rng.uniform(low=-np.pi, high=np.pi) - self.init_bucket_to_target_theta = theta - delta = np.array([np.cos(theta), np.sin(theta)]) * dist - pos_xy = center + delta - - self.bucket_center_offset = (bb[1, 2] - bb[0, 2]) / 5 - self.bucket_body_link = self.bucket.get_active_joints()[0].get_parent_link() - self.bb_local = np.array( - get_local_axis_aligned_bbox_for_link(self.bucket_body_link) - ) - self.center_local = (self.bb_local[0] + self.bb_local[1]) / 2 - - pose.set_p([pos_xy[0], pos_xy[1], pose.p[2] - bb[0, 2]]) - - # Sample an orientation for the bucket - ax = ay = 0 - az = self._episode_rng.uniform(low=-np.pi, high=np.pi) - q = euler2quat(ax, ay, az, "sxyz") - pose.set_q(q) - - # Finalize bucket pose - self.bucket.set_pose(pose) - self.init_bucket_height = ( - self.bucket_body_link.get_pose() - * self.bucket_body_link.get_cmass_local_pose() - ).p[2] - - # Finalize the bucket joint state - self.bucket.set_qpos(self.bucket.get_qlimit()[:, 0]) - self.bucket.set_qvel(np.zeros(self.bucket.dof)) - - def _initialize_robot(self): - # Base position - center = self.bucket.get_pose().p - dist = self._episode_rng.uniform(low=0.6, high=0.8) - theta = self._episode_rng.uniform(low=-0.4 * np.pi, high=0.4 * np.pi) - theta += self.init_bucket_to_target_theta - delta = np.array([np.cos(theta), np.sin(theta)]) * dist - base_pos = center[:2] + delta - - # Base orientation - perturb_orientation = self._episode_rng.uniform( - low=-0.05 * np.pi, high=0.05 * np.pi - ) - base_theta = -np.pi + theta + perturb_orientation - - # Torso height - h = 0.5 - - # Arm - arm_qpos = [0, 0, 0, -1.5, 0, 3, 0.78, 0.02, 0.02] - - qpos = np.hstack([base_pos, base_theta, h, arm_qpos, arm_qpos]) - self.agent.reset(qpos) - - def _initialize_balls(self): - bb = np.array(get_axis_aligned_bbox_for_articulation(self.bucket)) - R = self.balls_radius - - ball_idx = 0 - for i in range(self.GX): - for j in range(self.GY): - for k in range(self.GZ): - dx = -self.GX * R * 2 / 2 + R + 2 * R * i - dy = -self.GY * R * 2 / 2 + R + 2 * R * j - dz = R + R * 2 * k - pose = self.bucket.pose - pose = Pose( - [pose.p[0] + dx, pose.p[1] + dy, bb[1, 2] - bb[0, 2] + dz] - ) - actor = self.balls[ball_idx] - actor.set_pose(pose) - ball_idx += 1 - - def _set_bucket_links_pcd(self): - for name, info in self.links_info.items(): - mesh: trimesh.Trimesh = info[1] - with np_random(self._episode_seed): - pcd = mesh.sample(512) - self.links_info[name].append(pcd) - - # -------------------------------------------------------------------------- # - # Success metric and shaped reward - # -------------------------------------------------------------------------- # - def evaluate(self, **kwargs): - w2b = ( - self.bucket_body_link.pose.inv().to_transformation_matrix() - ) # world to bucket - - in_bucket = True - for b in self.balls: - p = w2b[:3, :3] @ b.pose.p + w2b[:3, 3] - if not np.all((p > self.bb_local[0]) * (p < self.bb_local[1])): - in_bucket = False - break - - z_axis_world = np.array([0, 0, 1]) - z_axis_bucket = quat2mat(self.root_link.get_pose().q) @ z_axis_world - bucket_tilt = abs(angle_between_vec(z_axis_world, z_axis_bucket)) - - dist_bucket_to_target = np.linalg.norm( - self.root_link.get_pose().p[:2] - self.target_xy - ) - - vel_norm = np.linalg.norm(self.root_link.linear_velocity) - ang_vel_norm = np.linalg.norm(self.root_link.angular_velocity) - - flags = { - "balls_in_bucket": in_bucket, - "bucket_above_platform": dist_bucket_to_target < 0.3, - "bucket_standing": bucket_tilt < 0.1 * np.pi, - # "bucket_static": (vel_norm < 0.1 and ang_vel_norm < 0.2), - "bucket_static": self.check_link_static( - self.root_link, max_v=0.1, max_ang_v=0.2 - ), - } - - return dict( - success=all(flags.values()), - **flags, - dist_bucket_to_target=dist_bucket_to_target, - bucket_tilt=bucket_tilt, - bucket_vel_norm=vel_norm, - bucket_ang_vel_norm=ang_vel_norm, - ) - - def _get_bucket_pcd(self): - """Get the point cloud of the bucket given its current joint positions.""" - links_pcd = [] - for name, info in self.links_info.items(): - link: physx.PhysxArticulationLinkComponent = info[0] - pcd: np.ndarray = info[2] - T = link.pose.to_transformation_matrix() - pcd = transform_points(T, pcd) - links_pcd.append(pcd) - chair_pcd = np.concatenate(links_pcd, axis=0) - return chair_pcd - - def compute_dense_reward(self, action, info: dict, **kwargs): - reward = -20.0 - - actor = self.root_link - ee_coords = np.array(self.agent.get_ee_coords()) - ee_mids = np.array([ee_coords[:2].mean(0), ee_coords[2:].mean(0)]) - # ee_vels = np.array(self.agent.get_ee_vels()) - bucket_pcd = self._get_bucket_pcd() - - # EE approach bucket - dist_ees_to_bucket = sdist.cdist(ee_coords, bucket_pcd) # [M, N] - dist_ees_to_bucket = dist_ees_to_bucket.min(1) # [M] - dist_ee_to_bucket = dist_ees_to_bucket.mean() - log_dist_ee_to_bucket = np.log(dist_ee_to_bucket + 1e-5) - reward += -dist_ee_to_bucket - np.clip(log_dist_ee_to_bucket, -10, 0) - - # EE adjust height - bucket_mid = ( - self.bucket_body_link.get_pose() - * self.bucket_body_link.get_cmass_local_pose() - ).p - bucket_mid[2] += self.bucket_center_offset - v1 = ee_mids[0] - bucket_mid - v2 = ee_mids[1] - bucket_mid - ees_oppo = sdist.cosine(v1, v2) - ees_height_diff = abs( - (quat2mat(self.root_link.get_pose().q).T @ (ee_mids[0] - ee_mids[1]))[2] - ) - log_ees_height_diff = np.log(ees_height_diff + 1e-5) - reward += -np.clip(log_ees_height_diff, -10, 0) * 0.2 - - # Keep bucket standing - bucket_tilt = info["bucket_tilt"] - log_dist_ori = np.log(bucket_tilt + 1e-5) - reward += -bucket_tilt * 0.2 - - # Penalize action - # Assume action is relative and normalized. - action_norm = np.linalg.norm(action) - reward -= action_norm * 1e-6 - - # Bucket velocity - actor_vel = actor.get_linear_velocity() - actor_vel_norm = np.linalg.norm(actor_vel) - disp_bucket_to_target = self.root_link.get_pose().p[:2] - self.target_xy - actor_vel_dir = sdist.cosine(actor_vel[:2], disp_bucket_to_target) - actor_ang_vel_norm = np.linalg.norm(actor.get_angular_velocity()) - actor_vel_up = actor_vel[2] - - # Stage reward - # NOTE(jigu): stage reward can also be used to debug which stage it is - stage_reward = 0 - - bucket_height = ( - self.bucket_body_link.get_pose() - * self.bucket_body_link.get_cmass_local_pose() - ).p[2] - dist_bucket_height = np.linalg.norm( - bucket_height - self.init_bucket_height - 0.2 - ) - dist_bucket_to_target = info["dist_bucket_to_target"] - - if dist_ee_to_bucket < 0.1: - stage_reward += 2 - reward += ees_oppo * 2 # - np.clip(log_ees_height_diff, -10, 0) * 0.2 - - bucket_height = ( - self.bucket_body_link.get_pose() - * self.bucket_body_link.get_cmass_local_pose() - ).p[2] - dist_bucket_height = np.linalg.norm( - bucket_height - self.init_bucket_height - 0.2 - ) - if dist_bucket_height < 0.03: - stage_reward += 2 - reward -= np.clip(log_dist_ori, -4, 0) - if dist_bucket_to_target <= 0.3: - stage_reward += 2 - reward += ( - np.exp(-actor_vel_norm * 10) * 2 - ) # + np.exp(-actor_ang_vel_norm) * 0.5 - if actor_vel_norm <= 0.1 and actor_ang_vel_norm <= 0.2: - stage_reward += 2 - if bucket_tilt <= 0.1 * np.pi: - stage_reward += 2 - else: - reward_vel = (actor_vel_dir - 1) * actor_vel_norm - reward += ( - np.clip(1 - np.exp(-reward_vel), -1, np.inf) * 2 - - dist_bucket_to_target * 2 - ) - else: - reward += ( - np.clip(1 - np.exp(-actor_vel_up), -1, np.inf) * 2 - - dist_bucket_height * 20 - ) - - if bucket_tilt > 0.4 * np.pi: - stage_reward -= 2 - - reward = reward + stage_reward - - # Update info - info.update( - action_norm=action_norm, - stage_reward=stage_reward, - dist_ee_to_bucket=dist_ee_to_bucket, - bucket_height=bucket_height, - ees_oppo=ees_oppo, - ees_height_diff=ees_height_diff, - ) - return reward - - def compute_normalized_dense_reward(self, **kwargs): - return self.compute_dense_reward(**kwargs) / 20.0 - - # ---------------------------------------------------------------------------- # - # Observation - # ---------------------------------------------------------------------------- # - def _get_task_actors(self): - return self.balls - - def _get_task_articulations(self): - # bucket max dof is 1 in our data - return [(self.bucket, 2)] - - def set_state_dict(self, state: np.ndarray): - super().set_state_dict(state) - self._prev_actor_pose = self.bucket.pose - - def _get_obs_extra(self): - obs = super()._get_obs_extra() - if self._obs_mode not in ["state", "state_dict"]: - obs["left_tcp_pose"] = vectorize_pose(self.left_tcp.pose) - obs["right_tcp_pose"] = vectorize_pose(self.right_tcp.pose) - return obs diff --git a/mani_skill2/envs/ms1/open_cabinet_door_drawer.py b/mani_skill2/envs/ms1/open_cabinet_door_drawer.py deleted file mode 100644 index 1bfffe7d7..000000000 --- a/mani_skill2/envs/ms1/open_cabinet_door_drawer.py +++ /dev/null @@ -1,448 +0,0 @@ -from collections import OrderedDict -from typing import List - -import numpy as np -import sapien -import sapien.physx as physx -import sapien.render -import trimesh -from sapien import Pose -from scipy.spatial import distance as sdist - -from mani_skill2.agents.robots.mobile_panda import MobilePandaSingleArm -from mani_skill2.utils.common import np_random, random_choice -from mani_skill2.utils.geometry import angle_distance, transform_points -from mani_skill2.utils.geometry.trimesh_utils import ( - get_articulation_meshes, - get_render_body_meshes, - get_render_shape_meshes, - merge_meshes, -) -from mani_skill2.utils.registration import register_env -from mani_skill2.utils.sapien_utils import get_obj_by_name -from mani_skill2.utils.structs.pose import vectorize_pose - -from .base_env import MS1BaseEnv - - -def clip_and_normalize(x, a_min, a_max=None): - if a_max is None: - a_max = np.abs(a_min) - a_min = -a_max - return (np.clip(x, a_min, a_max) - a_min) / (a_max - a_min) - - -class OpenCabinetEnv(MS1BaseEnv): - ASSET_UID = "partnet_mobility_cabinet" - MAX_DOF = 8 - agent: MobilePandaSingleArm - - def __init__(self, *args, fixed_target_link_idx: int = None, **kwargs): - # The index in target links (not all links) - self._fixed_target_link_idx = fixed_target_link_idx - self._cache_bboxes = {} - super().__init__(*args, **kwargs) - - def _register_render_cameras(self): - cam_cfg = super()._register_render_cameras() - cam_cfg.p = [-1.5, 0, 1.5] - cam_cfg.q = [0.9238795, 0, 0.3826834, 0] - return cam_cfg - - # -------------------------------------------------------------------------- # - # Reconfigure - # -------------------------------------------------------------------------- # - def _load_articulations(self): - urdf_config = dict( - material=dict(static_friction=1, dynamic_friction=1, restitution=0), - ) - scale = self.model_info["scale"] - self.cabinet = self._load_partnet_mobility( - fix_root_link=True, scale=scale, urdf_config=urdf_config - ) - self.cabinet.set_name(self.model_id) - - assert self.cabinet.dof <= self.MAX_DOF, self.cabinet.dof - self._set_cabinet_handles() - self._ignore_collision() - - if self._reward_mode in ["dense", "normalized_dense"]: - # NOTE(jigu): Explicit `set_pose` is needed. - self.cabinet.set_pose(Pose()) - self._set_cabinet_handles_mesh() - self._compute_handles_grasp_poses() - - def _set_cabinet_handles(self, joint_types: List[str]): - self.target_links = [] - self.target_joints = [] - self.target_handles = [] - - # NOTE(jigu): links and their parent joints. - for link, joint in zip(self.cabinet.get_links(), self.cabinet.get_joints()): - if joint.type not in joint_types: - continue - handles = [] - for render_shape in link.entity.find_component_by_type( - sapien.render.RenderBodyComponent - ).render_shapes: - if "handle" not in render_shape.name: - continue - handles.append(render_shape) - if len(handles) > 0: - self.target_links.append(link) - self.target_joints.append(joint) - self.target_handles.append(handles) - - def _set_cabinet_handles_mesh(self): - self.target_handles_mesh = [] - - for handle_visuals in self.target_handles: - meshes = [] - for render_shape in handle_visuals: - meshes.extend(get_render_shape_meshes(render_shape)) - handle_mesh = merge_meshes(meshes) - # Legacy issue: convex hull is assumed for further computation - handle_mesh = trimesh.convex.convex_hull(handle_mesh) - self.target_handles_mesh.append(handle_mesh) - - def _compute_grasp_poses(self, mesh: trimesh.Trimesh, pose: sapien.Pose): - # NOTE(jigu): only for axis-aligned horizontal and vertical cases - mesh2: trimesh.Trimesh = mesh.copy() - # Assume the cabinet is axis-aligned canonically - mesh2.apply_transform(pose.to_transformation_matrix()) - - extents = mesh2.extents - if extents[1] > extents[2]: # horizontal handle - closing = np.array([0, 0, 1]) - else: # vertical handle - closing = np.array([0, 1, 0]) - - # Only rotation of grasp poses are used. Thus, center is dummy. - approaching = [1, 0, 0] - grasp_poses = [ - self.agent.build_grasp_pose(approaching, closing, [0, 0, 0]), - self.agent.build_grasp_pose(approaching, -closing, [0, 0, 0]), - ] - - # # Visualization - # grasp_T = grasp_poses[0].to_transformation_matrix() - # coord_frame = trimesh.creation.axis( - # transform=grasp_T, origin_size=0.001, axis_radius=0.001, axis_length=0.01 - # ) - # trimesh.Scene([mesh2, coord_frame]).show() - - pose_inv = pose.inv() - grasp_poses = [pose_inv * x for x in grasp_poses] - - return grasp_poses - - def _compute_handles_grasp_poses(self): - self.target_handles_grasp_poses = [] - for i in range(len(self.target_handles)): - link = self.target_links[i] - mesh = self.target_handles_mesh[i] - grasp_poses = self._compute_grasp_poses(mesh, link.pose) - self.target_handles_grasp_poses.append(grasp_poses) - - def _ignore_collision(self): - """Ignore collision within the articulation to avoid impact from imperfect collision shapes.""" - # The legacy version only ignores collision of child links of active joints. - for link in self.cabinet.get_links(): - for s in link.get_collision_shapes(): - g0, g1, g2, g3 = s.get_collision_groups() - s.set_collision_groups([g0, g1, g2 | 1 << 31, g3]) - - def _load_agent(self): - self.agent = MobilePandaSingleArm( - self._scene, self._control_freq, self._control_mode - ) - - links = self.agent.robot.get_links() - self.tcp: physx.PhysxArticulationLinkComponent = get_obj_by_name( - links, "right_panda_hand_tcp" - ) - - # -------------------------------------------------------------------------- # - # Reset - # -------------------------------------------------------------------------- # - def reset(self, seed=None, options=None): - if options is None: - options = dict() - return super().reset(seed=seed, options=options) - - def _initialize_task(self): - self._initialize_cabinet() - self._initialize_robot() - self._set_target_link() - self._set_joint_physical_parameters() - - def _compute_cabinet_bbox(self): - mesh = merge_meshes(get_articulation_meshes(self.cabinet)) - return mesh.bounds # [2, 3] - - def _initialize_cabinet(self): - # Set joint positions to lower bounds - qlimits = self.cabinet.get_qlimits() # [N, 2] - assert not np.isinf(qlimits).any(), qlimits - qpos = np.ascontiguousarray(qlimits[:, 0]) - # NOTE(jigu): must use a contiguous array for `set_qpos` - self.cabinet.set_qpos(qpos) - - # If the scale can change, caching does not work. - bounds = self._cache_bboxes.get(self.model_id, None) - if bounds is None: - # The bound is computed based on current qpos. - # NOTE(jigu): Make sure the box is computed at a canoncial pose. - self.cabinet.set_pose(Pose()) - bounds = self._compute_cabinet_bbox() - self._cache_bboxes[self.model_id] = bounds - self.cabinet.set_pose(Pose([0, 0, -bounds[0, 2]])) - - def _initialize_robot(self): - # Base position - # The forward direction of cabinets is -x. - center = np.array([0, 0.8]) - dist = self._episode_rng.uniform(1.6, 1.8) - theta = self._episode_rng.uniform(0.9 * np.pi, 1.1 * np.pi) - direction = np.array([np.cos(theta), np.sin(theta)]) - xy = center + direction * dist - - # Base orientation - noise_ori = self._episode_rng.uniform(-0.05 * np.pi, 0.05 * np.pi) - ori = (theta - np.pi) + noise_ori - - h = 1e-4 - arm_qpos = np.array([0, 0, 0, -1.5, 0, 3, 0.78, 0.02, 0.02]) - - qpos = np.hstack([xy, ori, h, arm_qpos]) - self.agent.reset(qpos) - - def _set_joint_physical_parameters(self): - for joint in self.cabinet.get_active_joints(): - joint.set_friction(self._episode_rng.uniform(0.05, 0.15)) - joint.set_drive_property( - stiffness=0, damping=self._episode_rng.uniform(5, 20) - ) - - def _set_target_link(self): - if self._fixed_target_link_idx is None: - indices = np.arange(len(self.target_links)) - self.target_link_idx = random_choice(indices, rng=self._episode_rng) - else: - self.target_link_idx = self._fixed_target_link_idx - assert self.target_link_idx < len(self.target_links), self.target_link_idx - - self.target_link: physx.PhysxArticulationLinkComponent = self.target_links[ - self.target_link_idx - ] - self.target_joint: physx.PhysxJointComponent = self.target_joints[ - self.target_link_idx - ] - # The index in active joints - self.target_joint_idx_q = self.cabinet.get_active_joints().index( - self.target_joint - ) - - qmin, qmax = self.target_joint.get_limits()[0] - self.target_qpos = qmin + (qmax - qmin) * 0.9 - self.target_angle_diff = self.target_qpos - qmin - - # One-hot indicator for which link is target - self.target_indicator = np.zeros(self.MAX_DOF, np.float32) - self.target_indicator[self.target_joint_idx_q] = 1 - - # x-axis is the revolute/prismatic joint direction - joint_pose = self.target_joint.get_global_pose().to_transformation_matrix() - self.target_joint_axis = joint_pose[:3, 0] - # It is not handle position - cmass_pose = self.target_link.pose * self.target_link.cmass_local_pose - self.target_link_pos = cmass_pose.p - - # Cache handle point cloud - if self._reward_mode in ["dense", "normalized_dense"]: - self._set_target_handle_info() - - def _set_target_handle_info(self): - self.target_handle_mesh = self.target_handles_mesh[self.target_link_idx] - with np_random(self._episode_seed): - self.target_handle_pcd = self.target_handle_mesh.sample(100) - self.target_handle_sdf = trimesh.proximity.ProximityQuery( - self.target_handle_mesh - ) - - # -------------------------------------------------------------------------- # - # Success metric and shaped reward - # -------------------------------------------------------------------------- # - @property - def link_qpos(self): - return self.cabinet.get_qpos()[self.target_joint_idx_q] - - @property - def link_qvel(self): - return self.cabinet.get_qvel()[self.target_joint_idx_q] - - def evaluate(self, **kwargs) -> dict: - vel_norm = np.linalg.norm(self.target_link.linear_velocity) - ang_vel_norm = np.linalg.norm(self.target_link.angular_velocity) - link_qpos = self.link_qpos - - flags = dict( - # cabinet_static=vel_norm <= 0.1 and ang_vel_norm <= 1, - cabinet_static=self.check_link_static( - self.target_link, max_v=0.1, max_ang_v=1 - ), - open_enough=link_qpos >= self.target_qpos, - ) - - return dict( - success=all(flags.values()), - **flags, - link_vel_norm=vel_norm, - link_ang_vel_norm=ang_vel_norm, - link_qpos=link_qpos - ) - - def compute_dense_reward(self, *args, info: dict, **kwargs): - reward = 0.0 - - # -------------------------------------------------------------------------- # - # The end-effector should be close to the target pose - # -------------------------------------------------------------------------- # - handle_pose = self.target_link.pose - ee_pose = self.agent.hand.pose - - # Position - ee_coords = self.agent.get_ee_coords_sample() # [2, 10, 3] - handle_pcd = transform_points( - handle_pose.to_transformation_matrix(), self.target_handle_pcd - ) - # trimesh.PointCloud(handle_pcd).show() - disp_ee_to_handle = sdist.cdist(ee_coords.reshape(-1, 3), handle_pcd) - dist_ee_to_handle = disp_ee_to_handle.reshape(2, -1).min(-1) # [2] - reward_ee_to_handle = -dist_ee_to_handle.mean() * 2 - reward += reward_ee_to_handle - - # Encourage grasping the handle - ee_center_at_world = ee_coords.mean(0) # [10, 3] - ee_center_at_handle = transform_points( - handle_pose.inv().to_transformation_matrix(), ee_center_at_world - ) - # self.ee_center_at_handle = ee_center_at_handle - dist_ee_center_to_handle = self.target_handle_sdf.signed_distance( - ee_center_at_handle - ) - # print("SDF", dist_ee_center_to_handle) - dist_ee_center_to_handle = dist_ee_center_to_handle.max() - reward_ee_center_to_handle = ( - clip_and_normalize(dist_ee_center_to_handle, -0.01, 4e-3) - 1 - ) - reward += reward_ee_center_to_handle - - # pointer = trimesh.creation.icosphere(radius=0.02, color=(1, 0, 0)) - # trimesh.Scene([self.target_handle_mesh, trimesh.PointCloud(ee_center_at_handle)]).show() - - # Rotation - target_grasp_poses = self.target_handles_grasp_poses[self.target_link_idx] - target_grasp_poses = [handle_pose * x for x in target_grasp_poses] - angles_ee_to_grasp_poses = [ - angle_distance(ee_pose, x) for x in target_grasp_poses - ] - ee_rot_reward = -min(angles_ee_to_grasp_poses) / np.pi * 3 - reward += ee_rot_reward - - # -------------------------------------------------------------------------- # - # Stage reward - # -------------------------------------------------------------------------- # - coeff_qvel = 1.5 # joint velocity - coeff_qpos = 0.5 # joint position distance - stage_reward = -5 - (coeff_qvel + coeff_qpos) - # Legacy version also abstract coeff_qvel + coeff_qpos. - - link_qpos = info["link_qpos"] - link_qvel = self.link_qvel - link_vel_norm = info["link_vel_norm"] - link_ang_vel_norm = info["link_ang_vel_norm"] - - ee_close_to_handle = ( - dist_ee_to_handle.max() <= 0.01 and dist_ee_center_to_handle > 0 - ) - if ee_close_to_handle: - stage_reward += 0.5 - - # Distance between current and target joint positions - # TODO(jigu): the lower bound 0 is problematic? should we use lower bound of joint limits? - reward_qpos = ( - clip_and_normalize(link_qpos, 0, self.target_qpos) * coeff_qpos - ) - reward += reward_qpos - - if not info["open_enough"]: - # Encourage positive joint velocity to increase joint position - reward_qvel = clip_and_normalize(link_qvel, -0.1, 0.5) * coeff_qvel - reward += reward_qvel - else: - # Add coeff_qvel for smooth transition of stagess - stage_reward += 2 + coeff_qvel - reward_static = -(link_vel_norm + link_ang_vel_norm * 0.5) - reward += reward_static - - # Legacy version uses static from info, which is incompatible with MPC. - # if info["cabinet_static"]: - if link_vel_norm <= 0.1 and link_ang_vel_norm <= 1: - stage_reward += 1 - - # Update info - info.update(ee_close_to_handle=ee_close_to_handle, stage_reward=stage_reward) - - reward += stage_reward - return reward - - def compute_normalized_dense_reward(self, *args, info: dict, **kwargs): - return self.compute_dense_reward(*args, info=info, **kwargs) / 10.0 - - # -------------------------------------------------------------------------- # - # Observations - # -------------------------------------------------------------------------- # - def _get_obs_extra(self) -> OrderedDict: - obs = super()._get_obs_extra() - if self._obs_mode not in ["state", "state_dict"]: - obs.update( - target_angle_diff=self.target_angle_diff, - target_joint_axis=self.target_joint_axis, - target_link_pos=self.target_link_pos, - ) - obs["tcp_pose"] = vectorize_pose(self.tcp.pose) - return obs - - def _get_obs_priviledged(self): - obs = super()._get_obs_priviledged() - obs["target_indicator"] = self.target_indicator - return obs - - def _get_task_articulations(self): - # The maximum DoF is 6 in our data. - return [(self.cabinet, 8)] - - def set_state_dict(self, state: np.ndarray): - super().set_state_dict(state) - self._prev_actor_pose = self.target_link.pose - - -@register_env("OpenCabinetDoor-v1", max_episode_steps=200) -class OpenCabinetDoorEnv(OpenCabinetEnv): - DEFAULT_MODEL_JSON = ( - "{PACKAGE_ASSET_DIR}/partnet_mobility/meta/info_cabinet_door_train.json" - ) - - def _set_cabinet_handles(self): - super()._set_cabinet_handles(["revolute", "revolute_unwrapped"]) - - -@register_env("OpenCabinetDrawer-v1", max_episode_steps=200) -class OpenCabinetDrawerEnv(OpenCabinetEnv): - DEFAULT_MODEL_JSON = ( - "{PACKAGE_ASSET_DIR}/partnet_mobility/meta/info_cabinet_drawer_train.json" - ) - - def _set_cabinet_handles(self): - super()._set_cabinet_handles("prismatic") diff --git a/mani_skill2/envs/ms1/push_chair.py b/mani_skill2/envs/ms1/push_chair.py deleted file mode 100644 index 2353ff5b9..000000000 --- a/mani_skill2/envs/ms1/push_chair.py +++ /dev/null @@ -1,384 +0,0 @@ -import numpy as np -import sapien -import sapien.physx as physx -import sapien.render -import trimesh -from sapien import Pose -from scipy.spatial import distance as sdist -from transforms3d.euler import euler2quat, quat2euler - -from mani_skill2.agents.robots.mobile_panda import MobilePandaDualArm -from mani_skill2.utils.common import np_random -from mani_skill2.utils.geometry import transform_points -from mani_skill2.utils.geometry.trimesh_utils import get_actor_visual_mesh -from mani_skill2.utils.registration import register_env -from mani_skill2.utils.sapien_utils import get_obj_by_name -from mani_skill2.utils.structs.pose import vectorize_pose - -from .base_env import MS1BaseEnv - - -@register_env("PushChair-v1", max_episode_steps=200) -class PushChairEnv(MS1BaseEnv): - DEFAULT_MODEL_JSON = ( - "{PACKAGE_ASSET_DIR}/partnet_mobility/meta/info_chair_train.json" - ) - ASSET_UID = "partnet_mobility_chair" - - def _get_default_scene_config(self): - scene_config = super()._get_default_scene_config() - scene_config.solver_iterations = 15 - return scene_config - - def _register_render_cameras(self): - cam_cfg = super()._register_render_cameras() - cam_cfg.p = [0, 0, 4] - cam_cfg.q = [0.70710678, 0.0, 0.70710678, 0.0] - return cam_cfg - - # -------------------------------------------------------------------------- # - # Reconfigure - # -------------------------------------------------------------------------- # - def _load_articulations(self): - urdf_config = dict( - material=dict(static_friction=0.1, dynamic_friction=0.1, restitution=0), - density=200, - ) - self.chair = self._load_partnet_mobility( - fix_root_link=False, scale=0.8, urdf_config=urdf_config - ) - self.chair.set_name(self.model_id) - - self._set_chair_links() - self._ignore_collision() - - if self._reward_mode in ["dense", "normalized_dense"]: - self._set_chair_links_mesh() - - @staticmethod - def _check_link_types(link: physx.PhysxArticulationLinkComponent): - link_types = [] - comp = link.entity.find_component_by_type(sapien.render.RenderBodyComponent) - if comp is None: - return link_types - for shape in comp.render_shapes: - name = shape.name - if "wheel" in name: - link_types.append("wheel") - if "seat" in name: - link_types.append("seat") - if "leg" in name or "foot" in name: - link_types.append("support") - return link_types - - def _set_chair_links(self): - chair_links = self.chair.get_links() - - # Infer link types - self.root_link = chair_links[0] - self.wheel_links = [] - self.seat_link = None - self.support_link = None - for link in chair_links: - link_types = self._check_link_types(link) - if "wheel" in link_types: - self.wheel_links.append(link) - if "seat" in link_types: - assert self.seat_link is None, (self.seat_link, link) - self.seat_link = link - if "support" in link_types: - assert self.support_link is None, (self.support_link, link) - self.support_link = link - - # Set the physical material for wheels - wheel_material = self._scene.create_physical_material( - static_friction=1, dynamic_friction=1, restitution=0 - ) - for link in self.wheel_links: - for s in link.get_collision_shapes(): - s.set_physical_material(wheel_material) - - def _ignore_collision(self): - """Ignore collision within the articulation to avoid impact from imperfect collision shapes.""" - if self.seat_link == self.support_link: # sometimes they are the same one - return - - for link in [self.seat_link, self.support_link]: - shapes = link.get_collision_shapes() - for s in shapes: - g0, g1, g2, g3 = s.get_collision_groups() - s.set_collision_groups([g0, g1, g2 | 1 << 31, g3]) - - def _load_actors(self): - super()._load_actors() - - # A red sphere to indicate the target to push the chair. - builder = self._scene.create_actor_builder() - builder.add_sphere_visual( - radius=0.15, material=sapien.render.RenderMaterial(base_color=(1, 0, 0, 1)) - ) - self.target_indicator = builder.set_physx_body_type("static").build( - name="target_indicator" - ) - - def _load_agent(self): - self.agent = MobilePandaDualArm( - self._scene, self._control_freq, self._control_mode - ) - self.agent.camera_h = 2 - - links = self.agent.robot.get_links() - self.left_tcp: physx.PhysxArticulationLinkComponent = get_obj_by_name( - links, "left_panda_hand_tcp" - ) - self.right_tcp: physx.PhysxArticulationLinkComponent = get_obj_by_name( - links, "right_panda_hand_tcp" - ) - - def _set_chair_links_mesh(self): - self.links_info = {} - for link in self.chair.get_links(): - mesh = get_actor_visual_mesh(link.entity) - if mesh is None: - continue - self.links_info[link.name] = [link, mesh] - - # -------------------------------------------------------------------------- # - # Reset - # -------------------------------------------------------------------------- # - def _initialize_task(self): - self._set_target() - # NOTE(jigu): Initialize articulations and agent after the target is determined. - self._initialize_chair() - self._initialize_robot() - - if self._reward_mode in ["dense", "normalized_dense"]: - self._set_chair_links_pcd() - self._set_joint_physical_parameters() - - def _set_target(self): - self.target_xy = np.zeros(2, dtype=np.float32) - self.target_p = np.zeros(3, dtype=np.float32) - self.target_p[:2] = self.target_xy - self.target_indicator.set_pose(Pose(p=self.target_p)) - - def _initialize_chair(self): - # Load resting states from model info - p = np.float32(self.model_info["position"]) - q = np.float32(self.model_info["rotation"]) - init_qpos = self.model_info["initial_qpos"] - self.chair.set_pose(Pose(p, q)) - self.chair.set_qpos(init_qpos) - - # Sample a position around the target - r = self._episode_rng.uniform(0.8, 1.2) - theta = self._episode_rng.uniform(-np.pi, np.pi) - p[0] = np.cos(theta) * r - p[1] = np.sin(theta) * r - - # The chair faces towards the origin. - chair_init_ori = theta - np.pi - # Add a perturbation - noise_ori = self._episode_rng.uniform(-0.4 * np.pi, 0.4 * np.pi) - chair_init_ori = chair_init_ori + noise_ori - self.chair_init_ori = chair_init_ori - - # NOTE(jigu): The base link may not be axis-aligned, - # while the seat link is assumed to be aligned (z forward and y up). - # seat_dir = self.seat_link.pose.to_transformation_matrix()[:3, 2] - # seat_ori = np.arctan2(seat_dir[1], seat_dir[0]) - # # Compute the relative angle from seat to chair orientation - # delta_ori = self.chair_init_ori - seat_ori - # ax, ay, az = quat2euler(q, "sxyz") - # az = az + delta_ori - # q = euler2quat(ax, ay, az, "sxyz") - - # Legacy way to compute orientation (less readable) - _, _, az_seat = quat2euler(self.seat_link.get_pose().q, "sxyz") - ax, ay, az = quat2euler(q, "sxyz") - az = az - az_seat + np.pi * 1.5 + theta # face to target - az = az + noise_ori - q = euler2quat(ax, ay, az, "sxyz") - - # Finalize chair pose - self.chair.set_pose(Pose(p, q)) - - def _initialize_robot(self): - # Base position (around the chair, away from target) - xy = self.chair.pose.p[:2] - r = self._episode_rng.uniform(0.8, 1.2) - theta = self._episode_rng.uniform(-0.2 * np.pi, 0.2 * np.pi) - theta = (self.chair_init_ori + np.pi) + theta - xy[0] += np.cos(theta) * r - xy[1] += np.sin(theta) * r - - # Base orientation - noise_ori = self._episode_rng.uniform(-0.05 * np.pi, 0.05 * np.pi) - ori = (theta - np.pi) + noise_ori - - # Torso height - h = 0.9 - - # Arm - arm_qpos = [0, 0, 0, -1.5, 0, 3, 0.78, 0.02, 0.02] - - qpos = np.hstack([xy, ori, h, arm_qpos, arm_qpos]) - self.agent.reset(qpos) - - def _set_joint_physical_parameters(self): - for joint in self.chair.get_active_joints(): - parent_link = joint.get_parent_link() - - # revolute joint between seat and support - if parent_link is not None and "helper" in parent_link.name: - # assert joint.type == "revolute", (self.model_id, joint.type) - joint.set_friction(self._episode_rng.uniform(0.05, 0.15)) - joint.set_drive_property( - stiffness=0, damping=self._episode_rng.uniform(5, 15) - ) - else: - joint.set_friction(self._episode_rng.uniform(0.0, 0.1)) - joint.set_drive_property( - stiffness=0, damping=self._episode_rng.uniform(0, 0.5) - ) - - def _set_chair_links_pcd(self): - for name, info in self.links_info.items(): - mesh: trimesh.Trimesh = info[1] - with np_random(self._episode_seed): - pcd = mesh.sample(512) - self.links_info[name].append(pcd) - - # -------------------------------------------------------------------------- # - # Success metric and shaped reward - # -------------------------------------------------------------------------- # - def evaluate(self, **kwargs): - disp_chair_to_target = self.chair.pose.p[:2] - self.target_xy - dist_chair_to_target = np.linalg.norm(disp_chair_to_target) - - # z-axis of chair should be upward - z_axis_chair = self.root_link.pose.to_transformation_matrix()[:3, 2] - chair_tilt = np.arccos(z_axis_chair[2]) - - vel_norm = np.linalg.norm(self.root_link.linear_velocity) - ang_vel_norm = np.linalg.norm(self.root_link.angular_velocity) - - flags = dict( - chair_close_to_target=dist_chair_to_target < 0.15, - chair_standing=chair_tilt < 0.05 * np.pi, - # chair_static=(vel_norm < 0.1 and ang_vel_norm < 0.2), - chair_static=self.check_link_static( - self.root_link, max_v=0.1, max_ang_v=0.2 - ), - ) - return dict( - success=all(flags.values()), - **flags, - dist_chair_to_target=dist_chair_to_target, - chair_tilt=chair_tilt, - chair_vel_norm=vel_norm, - chair_ang_vel_norm=ang_vel_norm, - ) - - def _get_chair_pcd(self): - """Get the point cloud of the chair given its current joint positions.""" - links_pcd = [] - for name, info in self.links_info.items(): - link: physx.PhysxArticulationLinkComponent = info[0] - pcd: np.ndarray = info[2] - T = link.pose.to_transformation_matrix() - pcd = transform_points(T, pcd) - links_pcd.append(pcd) - chair_pcd = np.concatenate(links_pcd, axis=0) - return chair_pcd - - def compute_dense_reward(self, action: np.ndarray, info: dict, **kwargs): - reward = 0 - - # Compute distance between end-effectors and chair surface - ee_coords = np.array(self.agent.get_ee_coords()) # [M, 3] - chair_pcd = self._get_chair_pcd() # [N, 3] - - # EE approach chair - dist_ees_to_chair = sdist.cdist(ee_coords, chair_pcd) # [M, N] - dist_ees_to_chair = dist_ees_to_chair.min(1) # [M] - dist_ee_to_chair = dist_ees_to_chair.mean() - log_dist_ee_to_chair = np.log(dist_ee_to_chair + 1e-5) - reward += -dist_ee_to_chair - np.clip(log_dist_ee_to_chair, -10, 0) - - # Keep chair standing - chair_tilt = info["chair_tilt"] - reward += -chair_tilt * 0.2 - - # Penalize action - # Assume action is relative and normalized. - action_norm = np.linalg.norm(action) - reward -= action_norm * 1e-6 - - # Chair velocity - # Legacy version uses full velocity instead of xy-plane velocity - chair_vel = self.root_link.linear_velocity[:2] - chair_vel_norm = np.linalg.norm(chair_vel) - disp_chair_to_target = self.root_link.get_pose().p[:2] - self.target_xy - cos_chair_vel_to_target = sdist.cosine(disp_chair_to_target, chair_vel) - chair_ang_vel_norm = info["chair_ang_vel_norm"] - - # Stage reward - # NOTE(jigu): stage reward can also be used to debug which stage it is - stage_reward = -10 - # -18 can guarantee the reward is negative - dist_chair_to_target = info["dist_chair_to_target"] - - if chair_tilt < 0.2 * np.pi: - # Chair is standing - if dist_ee_to_chair < 0.1: - # EE is close to chair - stage_reward += 2 - if dist_chair_to_target <= 0.15: - # Chair is close to target - stage_reward += 2 - # Try to keep chair static - reward += np.exp(-chair_vel_norm * 10) * 2 - # Legacy: Note that the static condition here is different from success metric - if chair_vel_norm <= 0.1 and chair_ang_vel_norm <= 0.2: - stage_reward += 2 - else: - # Try to increase velocity along direction to the target - # Compute directional velocity - x = (1 - cos_chair_vel_to_target) * chair_vel_norm - reward += max(-1, 1 - np.exp(x)) * 2 - dist_chair_to_target * 2 - else: - stage_reward = -5 - - reward = reward + stage_reward - - # Update info - info.update( - dist_ee_to_chair=dist_ee_to_chair, - action_norm=action_norm, - chair_vel_norm=chair_vel_norm, - cos_chair_vel_to_target=cos_chair_vel_to_target, - stage_reward=stage_reward, - ) - return reward - - def compute_normalized_dense_reward(self, **kwargs): - return self.compute_dense_reward(**kwargs) / 10.0 - - # ---------------------------------------------------------------------------- # - # Observations - # ---------------------------------------------------------------------------- # - def _get_task_articulations(self): - # The maximum DoF is 20 in our data. - return [(self.chair, 25)] - - def set_state_dict(self, state: np.ndarray): - super().set_state_dict(state) - self._prev_actor_pose = self.root_link.pose - - def _get_obs_extra(self): - obs = super()._get_obs_extra() - if self._obs_mode not in ["state", "state_dict"]: - obs["left_tcp_pose"] = vectorize_pose(self.left_tcp.pose) - obs["right_tcp_pose"] = vectorize_pose(self.right_tcp.pose) - return obs diff --git a/mani_skill2/envs/ms2/__init__.py b/mani_skill2/envs/ms2/__init__.py deleted file mode 100644 index 20f833390..000000000 --- a/mani_skill2/envs/ms2/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from .assembly import * -from .misc import * - -# from .mpm import * -# from .pick_and_place import * diff --git a/mani_skill2/envs/ms2/assembly/__init__.py b/mani_skill2/envs/ms2/assembly/__init__.py deleted file mode 100644 index 6682a5966..000000000 --- a/mani_skill2/envs/ms2/assembly/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# isort: off -from . import assembling_kits -from . import peg_insertion_side -from . import plug_charger diff --git a/mani_skill2/envs/ms2/assembly/assembling_kits.py b/mani_skill2/envs/ms2/assembly/assembling_kits.py deleted file mode 100644 index 9682b85ce..000000000 --- a/mani_skill2/envs/ms2/assembly/assembling_kits.py +++ /dev/null @@ -1,242 +0,0 @@ -from pathlib import Path -from typing import List - -import numpy as np -import sapien -from transforms3d.euler import euler2quat, quat2euler - -from mani_skill2 import format_path -from mani_skill2.utils.io_utils import load_json -from mani_skill2.utils.registration import register_env -from mani_skill2.utils.sapien_utils import look_at -from mani_skill2.utils.scene_builder import TableSceneBuilder -from mani_skill2.utils.structs.pose import vectorize_pose - -from .base_env import StationaryManipulationEnv - - -@register_env("AssemblingKits-v0", max_episode_steps=200) -class AssemblingKitsEnv(StationaryManipulationEnv): - def __init__(self, asset_root="{ASSET_DIR}/assembling_kits", **kwargs): - self.asset_root = Path(format_path(asset_root)) - - self._kit_dir = self.asset_root / "kits" - self._models_dir = self.asset_root / "models" - if not (self._kit_dir.exists() and self._models_dir.exists()): - raise FileNotFoundError( - "The objects/kits are not found." - "Please download (ManiSkill2) AssemblingKits assets:" - "`python -m mani_skill2.utils.download_asset assembling_kits`." - ) - - self._episode_json = load_json(self.asset_root / "episodes.json") - self._episodes = self._episode_json["episodes"] - self.episode_idx = None - - self.symmetry = self._episode_json["config"]["symmetry"] - self.color = self._episode_json["config"]["color"] - self.object_scale = self._episode_json["config"]["object_scale"] - - super().__init__(**kwargs) - - def reset(self, seed=None, options=None): - if options is None: - options = dict() - self._set_episode_rng(seed) - episode_idx = options.pop("episode_idx", None) - reconfigure = options.pop("reconfigure", False) - if episode_idx is None: - episode_idx = self._episode_rng.randint(len(self._episodes)) - if self.episode_idx != episode_idx: - reconfigure = True - self.episode_idx = episode_idx - options["reconfigure"] = reconfigure - - episode = self._episodes[episode_idx] - self.kit_id: int = episode["kit"] - self.spawn_pos = np.float32(episode["spawn_pos"]) - self.object_id: int = episode["obj_to_place"] - self._other_objects_id: List[int] = episode["obj_in_place"] - - return super().reset(seed=self._episode_seed, options=options) - - def _parse_json(self, path): - """Parse kit JSON information""" - kit_json = load_json(path) - - self._objects_id = [o["object_id"] for o in kit_json["objects"]] - self.objects_pos = { - o["object_id"]: np.float32(o["pos"]) for o in kit_json["objects"] - } - # z-axis orientation - self.objects_rot = {o["object_id"]: o["rot"] for o in kit_json["objects"]} - - def _load_kit(self): - self._parse_json(self._kit_dir / f"{self.kit_id}.json") - - builder = self._scene.create_actor_builder() - kit_path = str(self._kit_dir / f"{self.kit_id}.obj") - builder.add_nonconvex_collision_from_file(kit_path) - - material = self._renderer.create_material() - material.set_base_color([0.27807487, 0.20855615, 0.16934046, 1.0]) - material.metallic = 0.0 - material.roughness = 0.5 - material.specular = 0.0 - builder.add_visual_from_file(kit_path, material=material) - - return builder.build_static(name="kit") - - def _load_object(self, object_id): - collision_path = self._models_dir / "collision" / f"{object_id:02d}.obj" - visual_path = self._models_dir / "visual" / f"{object_id:02d}.obj" - - builder = self._scene.create_actor_builder() - builder.add_multiple_convex_collisions_from_file( - str(collision_path), scale=self.object_scale - ) - - material = self._renderer.create_material() - material.set_base_color(self.color[self._episode_rng.choice(len(self.color))]) - material.metallic = 0.0 - material.roughness = 0.1 - material.specular = 0.0 - builder.add_visual_from_file( - str(visual_path), scale=self.object_scale, material=material - ) - - return builder.build(f"obj_{object_id:02d}") - - def _load_actors(self): - TableSceneBuilder().build(self._scene) - - self.kit = self._load_kit() - self.obj = self._load_object(self.object_id) - self._other_objs = [self._load_object(i) for i in self._other_objects_id] - - def _initialize_actors(self): - self.kit.set_pose(sapien.Pose([0, 0, 0.01])) - self.obj.set_pose( - sapien.Pose( - self.spawn_pos, - euler2quat(0, 0, self._episode_rng.rand() * np.pi * 2), - ) - ) - - for i, o in enumerate(self._other_objs): - obj_id = self._other_objects_id[i] - o.set_pose( - sapien.Pose( - np.array(self.objects_pos[obj_id]), - euler2quat(0, 0, self.objects_rot[obj_id]), - ) - ) - - def _initialize_task(self): - self._obj_init_pos = np.float32(self.spawn_pos) - self._obj_goal_pos = np.float32(self.objects_pos[self.object_id]) - - def _get_obs_extra(self): - obs = dict( - tcp_pose=vectorize_pose(self.agent.tcp.pose), - obj_init_pos=self._obj_init_pos, - obj_goal_pos=self._obj_goal_pos, - ) - if self._obs_mode in ["state", "state_dict"]: - obs.update( - obj_pose=vectorize_pose(self.obj.pose), - tcp_to_obj_pos=self.obj.pose.p - self.agent.tcp.pose.p, - obj_to_goal_pos=self.objects_pos[self.object_id] - self.obj.pose.p, - ) - return obs - - def _check_pos_diff(self, pos_eps=2e-2): - pos_diff = self.objects_pos[self.object_id][:2] - self.obj.get_pose().p[:2] - pos_diff_norm = np.linalg.norm(pos_diff) - return pos_diff, pos_diff_norm, pos_diff_norm < pos_eps - - def _check_rot_diff(self, rot_eps=np.deg2rad(4)): - rot = quat2euler(self.obj.get_pose().q)[-1] # Check z-axis rotation - rot_diff = 0 - if self.symmetry[self.object_id] > 0: - rot_diff = ( - np.abs(rot - self.objects_rot[self.object_id]) - % self.symmetry[self.object_id] - ) - if rot_diff > (self.symmetry[self.object_id] / 2): - rot_diff = self.symmetry[self.object_id] - rot_diff - return rot_diff, rot_diff < rot_eps - - def _check_in_slot(self, obj: sapien.Entity, height_eps=3e-3): - return obj.pose.p[2] < height_eps - - def evaluate(self, **kwargs) -> dict: - pos_diff, pos_diff_norm, pos_correct = self._check_pos_diff() - rot_diff, rot_correct = self._check_rot_diff() - in_slot = self._check_in_slot(self.obj) - return { - "pos_diff": pos_diff, - "pos_diff_norm": pos_diff_norm, - "pos_correct": pos_correct, - "rot_diff": rot_diff, - "rot_correct": rot_correct, - "in_slot": in_slot, - "success": pos_correct and rot_correct and in_slot, - } - - def compute_dense_reward(self, info, **kwargs): - if info["success"]: - return 10.0 - - reward = 0.0 - gripper_width = ( - self.agent.robot.get_qlimits()[-1, 1] * 2 - ) # NOTE: hard-coded with panda - - # reaching reward - gripper_to_obj_dist = np.linalg.norm(self.agent.tcp.pose.p - self.obj.pose.p) - reaching_reward = 1 - np.tanh(4.0 * np.maximum(gripper_to_obj_dist - 0.01, 0.0)) - reward += reaching_reward - - # object position and rotation reward - pos_diff_norm_reward = 1 - np.tanh(4.0 * info["pos_diff_norm"]) - rot_diff_reward = 1 - np.tanh(0.4 * info["rot_diff"]) - rot_diff_reward += 1.0 - np.tanh( - 1.0 - self.obj.pose.to_transformation_matrix()[2, 2] - ) # encourage object to be parallel to xy-plane - object_well_positioned = info["pos_correct"] and info["rot_correct"] - if object_well_positioned: - reward += 1.0 - - # grasp reward - is_grasped = self.agent.is_grasping( - self.obj, max_angle=30 - ) # max_angle ensures that the gripper grasps the object appropriately, not in a strange pose - if is_grasped or object_well_positioned: - reward += 2.0 - reward += pos_diff_norm_reward - reward += rot_diff_reward - if object_well_positioned and rot_diff_reward > 0.98: - # ungrasp the object - if not is_grasped: - reward += 1.0 - else: - reward = ( - reward - + 1.0 * np.sum(self.agent.robot.get_qpos()[-2:]) / gripper_width - ) - - return reward - - def compute_normalized_dense_reward(self, **kwargs): - return self.compute_dense_reward(**kwargs) / 10.0 - - def _register_sensors(self): - cam_cfg = super()._register_sensors() - cam_cfg.pose = look_at([0.2, 0, 0.4], [0, 0, 0]) - return cam_cfg - - def _register_human_render_cameras(self): - cam_cfg = super()._register_render_cameras() - cam_cfg.pose = look_at([0.3, 0.3, 0.8], [0.0, 0.0, 0.1]) - return cam_cfg diff --git a/mani_skill2/envs/ms2/assembly/base_env.py b/mani_skill2/envs/ms2/assembly/base_env.py deleted file mode 100644 index da7b31f9a..000000000 --- a/mani_skill2/envs/ms2/assembly/base_env.py +++ /dev/null @@ -1,68 +0,0 @@ -from typing import Union - -import numpy as np -import sapien -from sapien import Pose - -from mani_skill2.agents.base_agent import BaseAgent -from mani_skill2.agents.robots.panda import Panda, PandaRealSensed435 -from mani_skill2.agents.robots.xmate3 import Xmate3Robotiq -from mani_skill2.envs.sapien_env import BaseEnv -from mani_skill2.sensors.camera import CameraConfig -from mani_skill2.utils.sapien_utils import look_at -from mani_skill2.utils.structs.pose import vectorize_pose - - -class StationaryManipulationEnv(BaseEnv): - agent: Union[Panda, Xmate3Robotiq] - - def __init__( - self, *args, robot_uids=PandaRealSensed435, robot_init_qpos_noise=0.02, **kwargs - ): - self.robot_init_qpos_noise = robot_init_qpos_noise - super().__init__(*args, robot_uids=robot_uids, **kwargs) - - def _initialize_agent(self): - if self.robot_uids == "panda" or self.robot_uids == "panda_realsensed435": - # fmt: off - # EE at [0.615, 0, 0.17] - qpos = np.array( - [0.0, np.pi / 8, 0, -np.pi * 5 / 8, 0, np.pi * 3 / 4, np.pi / 4, 0.04, 0.04] - ) - # fmt: on - qpos[:-2] += self._episode_rng.normal( - 0, self.robot_init_qpos_noise, len(qpos) - 2 - ) - self.agent.reset(qpos) - self.agent.robot.set_pose(Pose([-0.615, 0, 0])) - elif self.robot_uids == "xmate3_robotiq": - qpos = np.array( - [0, np.pi / 6, 0, np.pi / 3, 0, np.pi / 2, -np.pi / 2, 0.04, 0.04] - ) - qpos[:-2] += self._episode_rng.normal( - 0, self.robot_init_qpos_noise, len(qpos) - 2 - ) - self.agent.reset(qpos) - self.agent.robot.set_pose(Pose([-0.562, 0, 0])) - else: - raise NotImplementedError(self.robot_uids) - - def _register_sensors(self): - pose = look_at([0.2, 0, 0.4], [0, 0, 0]) - 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([1.0, 1.0, 0.8], [0.0, 0.0, 0.5]) - return CameraConfig("render_camera", pose.p, pose.q, 512, 512, 1, 0.01, 10) - - def _setup_viewer(self): - super()._setup_viewer() - self._viewer.set_camera_xyz(0.8, 0, 1.0) - self._viewer.set_camera_rpy(0, -0.5, 3.14) - - def _get_obs_agent(self): - obs = self.agent.get_proprioception() - obs["base_pose"] = vectorize_pose(self.agent.robot.pose) - return obs diff --git a/mani_skill2/envs/ms2/assembly/peg_insertion_side.py b/mani_skill2/envs/ms2/assembly/peg_insertion_side.py deleted file mode 100644 index baef66106..000000000 --- a/mani_skill2/envs/ms2/assembly/peg_insertion_side.py +++ /dev/null @@ -1,293 +0,0 @@ -from collections import OrderedDict - -import numpy as np -import sapien -import torch -from sapien import Pose -from transforms3d.euler import euler2quat - -from mani_skill2.utils.registration import register_env -from mani_skill2.utils.sapien_utils import hex2rgba, look_at -from mani_skill2.utils.scene_builder import TableSceneBuilder -from mani_skill2.utils.structs.pose import vectorize_pose - -from .base_env import StationaryManipulationEnv - - -@register_env("PegInsertionSide-v0", max_episode_steps=200) -class PegInsertionSideEnv(StationaryManipulationEnv): - _clearance = 0.003 - - def reset(self, seed=None, options=None): - if options is None: - options = {} - if options.get("reconfigure") is None: - options["reconfigure"] = True - return super().reset(seed, options) - - def _build_box_with_hole( - self, inner_radius, outer_radius, depth, center=(0, 0), name="box_with_hole" - ): - builder = self._scene.create_actor_builder() - thickness = (outer_radius - inner_radius) * 0.5 - # x-axis is hole direction - half_center = [x * 0.5 for x in center] - half_sizes = [ - [depth, thickness - half_center[0], outer_radius], - [depth, thickness + half_center[0], outer_radius], - [depth, outer_radius, thickness - half_center[1]], - [depth, outer_radius, thickness + half_center[1]], - ] - offset = thickness + inner_radius - poses = [ - Pose([0, offset + half_center[0], 0]), - Pose([0, -offset + half_center[0], 0]), - Pose([0, 0, offset + half_center[1]]), - Pose([0, 0, -offset + half_center[1]]), - ] - - mat = self._renderer.create_material() - mat.set_base_color(hex2rgba("#FFD289")) - mat.metallic = 0.0 - mat.roughness = 0.5 - mat.specular = 0.5 - - for half_size, pose in zip(half_sizes, poses): - builder.add_box_collision(pose, half_size) - builder.add_box_visual(pose, half_size, material=mat) - - return builder.build_static(name) - - def _load_actors(self): - self.table_scene = TableSceneBuilder(env=self) - self.table_scene.build() - - # peg - # length, radius = 0.1, 0.02 - length = self._episode_rng.uniform(0.075, 0.125) - radius = self._episode_rng.uniform(0.015, 0.025) - builder = self._scene.create_actor_builder() - builder.add_box_collision(half_size=[length, radius, radius]) - - # peg head - mat = self._renderer.create_material() - mat.set_base_color(hex2rgba("#EC7357")) - mat.metallic = 0.0 - mat.roughness = 0.5 - mat.specular = 0.5 - builder.add_box_visual( - Pose([length / 2, 0, 0]), - half_size=[length / 2, radius, radius], - material=mat, - ) - - # peg tail - mat = self._renderer.create_material() - mat.set_base_color(hex2rgba("#EDF6F9")) - mat.metallic = 0.0 - mat.roughness = 0.5 - mat.specular = 0.5 - builder.add_box_visual( - Pose([-length / 2, 0, 0]), - half_size=[length / 2, radius, radius], - material=mat, - ) - - self.peg = builder.build("peg") - self.peg_head_offset = Pose([length, 0, 0]) - self.peg_half_size = np.float32([length, radius, radius]) - - # box with hole - center = 0.5 * (length - radius) * self._episode_rng.uniform(-1, 1, size=2) - inner_radius, outer_radius, depth = radius + self._clearance, length, length - self.box = self._build_box_with_hole( - inner_radius, outer_radius, depth, center=center - ) - self.box_hole_offset = Pose(np.hstack([0, center])) - self.box_hole_radius = inner_radius - - def _initialize_actors(self): - self.table_scene.initialize() - xy = self._episode_rng.uniform([-0.1, -0.3], [0.1, 0]) - pos = np.hstack([xy, self.peg_half_size[2]]) - ori = np.pi / 2 + self._episode_rng.uniform(-np.pi / 3, np.pi / 3) - quat = euler2quat(0, 0, ori) - self.peg.set_pose(Pose(pos, quat)) - - xy = self._episode_rng.uniform([-0.05, 0.2], [0.05, 0.4]) - pos = np.hstack([xy, self.peg_half_size[0]]) - ori = np.pi / 2 + self._episode_rng.uniform(-np.pi / 8, np.pi / 8) - quat = euler2quat(0, 0, ori) - self.box.set_pose(Pose(pos, quat)) - - def _initialize_agent(self): - if self.robot_uids == "panda_realsensed435": - # fmt: off - qpos = np.array( - [0.0, np.pi / 8, 0, -np.pi * 5 / 8, 0, np.pi * 3 / 4, -np.pi / 4, 0.04, 0.04] - ) - # fmt: on - qpos[:-2] += self._episode_rng.normal( - 0, self.robot_init_qpos_noise, len(qpos) - 2 - ) - self.agent.reset(qpos) - self.agent.robot.set_pose(Pose([-0.615, 0, 0])) - else: - raise NotImplementedError(self.robot_uids) - - @property - def peg_head_pos(self): - return self.peg.pose.p + self.peg_head_offset.p - - @property - def peg_head_pose(self): - return self.peg.pose * self.peg_head_offset - - @property - def box_hole_pose(self): - return self.box.pose * self.box_hole_offset - - def _initialize_task(self): - self.goal_pos = self.box_hole_pose.p # goal of peg head inside the hole - # NOTE(jigu): The goal pose is computed based on specific geometries used in this task. - # Only consider one side - self.goal_pose = ( - self.box.pose * self.box_hole_offset * self.peg_head_offset.inv() - ) - - def _get_obs_extra(self) -> OrderedDict: - obs = OrderedDict(tcp_pose=vectorize_pose(self.agent.tcp.pose)) - if self._obs_mode in ["state", "state_dict"]: - obs.update( - peg_pose=vectorize_pose(self.peg.pose), - peg_half_size=self.peg_half_size, - box_hole_pose=vectorize_pose(self.box_hole_pose), - box_hole_radius=self.box_hole_radius, - ) - return obs - - def has_peg_inserted(self): - # Only head position is used in fact - peg_head_pos_at_hole = (self.box_hole_pose.inv() * self.peg_head_pose).p - # x-axis is hole direction - x_flag = -0.015 <= peg_head_pos_at_hole[:, 0] - y_flag = ( - -self.box_hole_radius <= peg_head_pos_at_hole[:, 1] <= self.box_hole_radius - ) - z_flag = ( - -self.box_hole_radius <= peg_head_pos_at_hole[:, 2] <= self.box_hole_radius - ) - return ( - torch.logical_and(torch.logical_and(x_flag, y_flag), z_flag), - peg_head_pos_at_hole, - ) - - def evaluate(self, **kwargs) -> dict: - success, peg_head_pos_at_hole = self.has_peg_inserted() - return dict(success=success, peg_head_pos_at_hole=peg_head_pos_at_hole) - - def compute_dense_reward(self, info, **kwargs): - reward = 0.0 - - if info["success"]: - return 25.0 - - # grasp pose rotation reward - tcp_pose_wrt_peg = self.peg.pose.inv() * self.agent.tcp.pose - tcp_rot_wrt_peg = tcp_pose_wrt_peg.to_transformation_matrix()[:3, :3] - gt_rot_1 = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]]) - gt_rot_2 = np.array([[-1, 0, 0], [0, 1, 0], [0, 0, -1]]) - grasp_rot_loss_fxn = lambda A: np.arcsin( - np.clip(1 / (2 * np.sqrt(2)) * np.sqrt(np.trace(A.T @ A)), 0, 1) - ) - grasp_rot_loss = np.minimum( - grasp_rot_loss_fxn(gt_rot_1 - tcp_rot_wrt_peg), - grasp_rot_loss_fxn(gt_rot_2 - tcp_rot_wrt_peg), - ) / (np.pi / 2) - rotated_properly = grasp_rot_loss < 0.2 - reward += 1 - grasp_rot_loss - - gripper_pos = self.agent.tcp.pose.p - tgt_gripper_pose = self.peg.pose - offset = sapien.Pose( - [-0.06, 0, 0] - ) # account for panda gripper width with a bit more leeway - tgt_gripper_pose = tgt_gripper_pose * (offset) - if rotated_properly: - # reaching reward - gripper_to_peg_dist = np.linalg.norm(gripper_pos - tgt_gripper_pose.p) - reaching_reward = 1 - np.tanh( - 4.0 * np.maximum(gripper_to_peg_dist - 0.015, 0.0) - ) - # reaching_reward = 1 - np.tanh(10.0 * gripper_to_peg_dist) - reward += reaching_reward - - # grasp reward - is_grasped = self.agent.is_grasping( - self.peg, max_angle=20 - ) # max_angle ensures that the gripper grasps the peg appropriately, not in a strange pose - if is_grasped: - reward += 2.0 - - # pre-insertion award, encouraging both the peg center and the peg head to match the yz coordinates of goal_pose - pre_inserted = False - if is_grasped: - peg_head_wrt_goal = self.goal_pose.inv() * self.peg_head_pose - peg_head_wrt_goal_yz_dist = np.linalg.norm(peg_head_wrt_goal.p[1:]) - peg_wrt_goal = self.goal_pose.inv() * self.peg.pose - peg_wrt_goal_yz_dist = np.linalg.norm(peg_wrt_goal.p[1:]) - if peg_head_wrt_goal_yz_dist < 0.01 and peg_wrt_goal_yz_dist < 0.01: - pre_inserted = True - reward += 3.0 - pre_insertion_reward = 3 * ( - 1 - - np.tanh( - 0.5 * (peg_head_wrt_goal_yz_dist + peg_wrt_goal_yz_dist) - + 4.5 - * np.maximum(peg_head_wrt_goal_yz_dist, peg_wrt_goal_yz_dist) - ) - ) - reward += pre_insertion_reward - - # insertion reward - if is_grasped and pre_inserted: - peg_head_wrt_goal_inside_hole = ( - self.box_hole_pose.inv() * self.peg_head_pose - ) - insertion_reward = 5 * ( - 1 - np.tanh(5.0 * np.linalg.norm(peg_head_wrt_goal_inside_hole.p)) - ) - reward += insertion_reward - else: - reward = reward - 10 * np.maximum( - self.peg.pose.p[2] - + self.peg_half_size[2] - + 0.01 - - self.agent.tcp.pose.p[2], - 0.0, - ) - reward = reward - 10 * np.linalg.norm( - tgt_gripper_pose.p[:2] - self.agent.tcp.pose.p[:2] - ) - - return reward - - def compute_normalized_dense_reward(self, **kwargs): - return self.compute_dense_reward(**kwargs) / 25.0 - - def _register_sensors(self): - cam_cfg = super()._register_sensors() - cam_cfg.pose = look_at([0, -0.3, 0.2], [0, 0, 0.1]) - return cam_cfg - - def _register_human_render_cameras(self): - cam_cfg = super()._register_render_cameras() - cam_cfg.pose = look_at([0.5, -0.5, 0.8], [0.05, -0.1, 0.4]) - return cam_cfg - - def set_state_dict(self, state): - super().set_state_dict(state) - # NOTE(xuanlin): This way is specific to how we compute goals. - # The general way is to handle variables explicitly - # TODO (stao): can we refactor this out - self._initialize_task() diff --git a/mani_skill2/envs/ms2/assembly/plug_charger.py b/mani_skill2/envs/ms2/assembly/plug_charger.py deleted file mode 100644 index c082cf24b..000000000 --- a/mani_skill2/envs/ms2/assembly/plug_charger.py +++ /dev/null @@ -1,317 +0,0 @@ -from collections import OrderedDict - -import numpy as np -import sapien -import sapien.physx as physx -from sapien import Pose -from transforms3d.euler import euler2quat -from transforms3d.quaternions import qinverse, qmult, quat2axangle - -from mani_skill2.utils.registration import register_env -from mani_skill2.utils.sapien_utils import hex2rgba, look_at -from mani_skill2.utils.scene_builder import TableSceneBuilder -from mani_skill2.utils.structs.pose import vectorize_pose - -from .base_env import StationaryManipulationEnv - - -@register_env("PlugCharger-v0", max_episode_steps=200) -class PlugChargerEnv(StationaryManipulationEnv): - _base_size = [2e-2, 1.5e-2, 1.2e-2] # charger base half size - _peg_size = [8e-3, 0.75e-3, 3.2e-3] # charger peg half size - _peg_gap = 7e-3 # charger peg gap - _clearance = 5e-4 # single side clearance - _receptacle_size = [1e-2, 5e-2, 5e-2] # receptacle half size - - def _build_charger(self, peg_size, base_size, gap): - builder = self._scene.create_actor_builder() - - # peg - mat = self._renderer.create_material() - mat.set_base_color(hex2rgba("#FFFFFF")) - mat.metallic = 1.0 - mat.roughness = 0.0 - mat.specular = 1.0 - builder.add_box_collision(Pose([peg_size[0], gap, 0]), peg_size) - builder.add_box_visual(Pose([peg_size[0], gap, 0]), peg_size, material=mat) - builder.add_box_collision(Pose([peg_size[0], -gap, 0]), peg_size) - builder.add_box_visual(Pose([peg_size[0], -gap, 0]), peg_size, material=mat) - - # base - mat = self._renderer.create_material() - mat.set_base_color(hex2rgba("#FFFFFF")) - mat.metallic = 0.0 - mat.roughness = 0.1 - builder.add_box_collision(Pose([-base_size[0], 0, 0]), base_size) - builder.add_box_visual(Pose([-base_size[0], 0, 0]), base_size, material=mat) - - return builder.build(name="charger") - - def _build_receptacle(self, peg_size, receptacle_size, gap): - builder = self._scene.create_actor_builder() - - sy = 0.5 * (receptacle_size[1] - peg_size[1] - gap) - sz = 0.5 * (receptacle_size[2] - peg_size[2]) - dx = -receptacle_size[0] - dy = peg_size[1] + gap + sy - dz = peg_size[2] + sz - - mat = self._renderer.create_material() - mat.set_base_color(hex2rgba("#FFFFFF")) - mat.metallic = 0.0 - mat.roughness = 0.1 - - poses = [ - Pose([dx, 0, dz]), - Pose([dx, 0, -dz]), - Pose([dx, dy, 0]), - Pose([dx, -dy, 0]), - ] - half_sizes = [ - [receptacle_size[0], receptacle_size[1], sz], - [receptacle_size[0], receptacle_size[1], sz], - [receptacle_size[0], sy, receptacle_size[2]], - [receptacle_size[0], sy, receptacle_size[2]], - ] - for pose, half_size in zip(poses, half_sizes): - builder.add_box_collision(pose, half_size) - builder.add_box_visual(pose, half_size, material=mat) - - # Fill the gap - pose = Pose([-receptacle_size[0], 0, 0]) - half_size = [receptacle_size[0], gap - peg_size[1], peg_size[2]] - builder.add_box_collision(pose, half_size) - builder.add_box_visual(pose, half_size, material=mat) - - # Add dummy visual for hole - mat = self._renderer.create_material() - mat.set_base_color(hex2rgba("#DBB539")) - mat.metallic = 1.0 - mat.roughness = 0.0 - mat.specular = 1.0 - pose = Pose([-receptacle_size[0], -(gap * 0.5 + peg_size[1]), 0]) - half_size = [receptacle_size[0], peg_size[1], peg_size[2]] - builder.add_box_visual(pose, half_size, material=mat) - pose = Pose([-receptacle_size[0], gap * 0.5 + peg_size[1], 0]) - builder.add_box_visual(pose, half_size, material=mat) - - return builder.build_static(name="receptacle") - - def _load_actors(self): - TableSceneBuilder().build(self._scene) - self.charger = self._build_charger( - self._peg_size, - self._base_size, - self._peg_gap, - ) - self.receptacle = self._build_receptacle( - [ - self._peg_size[0], - self._peg_size[1] + self._clearance, - self._peg_size[2] + self._clearance, - ], - self._receptacle_size, - self._peg_gap, - ) - - def _initialize_actors(self): - # Initialize charger - xy = self._episode_rng.uniform( - [-0.1, -0.2], [-0.01 - self._peg_size[0] * 2, 0.2] - ) - # xy = [-0.05, 0] - pos = np.hstack([xy, self._base_size[2]]) - ori = self._episode_rng.uniform(-np.pi / 3, np.pi / 3) - # ori = 0 - quat = euler2quat(0, 0, ori) - self.charger.set_pose(Pose(pos, quat)) - - # Initialize receptacle - xy = self._episode_rng.uniform([0.01, -0.1], [0.1, 0.1]) - # xy = [0.05, 0] - pos = np.hstack([xy, 0.1]) - ori = np.pi + self._episode_rng.uniform(-np.pi / 8, np.pi / 8) - # ori = np.pi - quat = euler2quat(0, 0, ori) - self.receptacle.set_pose(Pose(pos, quat)) - - # Adjust render camera - if "render_camera" in self._human_render_cameras: - self._human_render_cameras[ - "render_camera" - ].camera.local_pose = self.receptacle.pose * look_at( - [0.3, 0.4, 0.1], [0, 0, 0] - ) - - def _initialize_task(self): - self.goal_pose = self.receptacle.pose * (Pose(q=euler2quat(0, 0, np.pi))) - # NOTE(jigu): clearance need to be set to 1e-3 so that the charger will not fall off - # self.charger.set_pose(self.goal_pose) - - @property - def charger_base_pose(self): - return self.charger.pose * (Pose([-self._base_size[0], 0, 0])) - - def _get_obs_extra(self) -> OrderedDict: - obs = OrderedDict(tcp_pose=vectorize_pose(self.agent.tcp.pose)) - if self._obs_mode in ["state", "state_dict"]: - obs.update( - charger_pose=vectorize_pose(self.charger.pose), - receptacle_pose=vectorize_pose(self.receptacle.pose), - goal_pose=vectorize_pose(self.goal_pose), - ) - return obs - - def evaluate(self, **kwargs): - obj_to_goal_dist, obj_to_goal_angle = self._compute_distance() - success = obj_to_goal_dist <= 5e-3 and obj_to_goal_angle <= 0.2 - return dict( - obj_to_goal_dist=obj_to_goal_dist, - obj_to_goal_angle=obj_to_goal_angle, - success=success, - ) - - def _compute_distance(self): - obj_pose = self.charger.pose - obj_to_goal_pos = self.goal_pose.p - obj_pose.p - obj_to_goal_dist = np.linalg.norm(obj_to_goal_pos) - - obj_to_goal_quat = qmult(qinverse(self.goal_pose.q), obj_pose.q) - _, obj_to_goal_angle = quat2axangle(obj_to_goal_quat) - obj_to_goal_angle = min(obj_to_goal_angle, np.pi * 2 - obj_to_goal_angle) - assert obj_to_goal_angle >= 0.0, obj_to_goal_angle - - return obj_to_goal_dist, obj_to_goal_angle - - def compute_dense_reward(self, info, **kwargs): - reward = 0.0 - - if info["success"]: - return 50.0 - - cmass_pose = self.charger.pose * ( - self.charger.find_component_by_type( - physx.PhysxRigidDynamicComponent - ).cmass_local_pose - ) - # grasp pose rotation reward - tcp_pose_wrt_charger = cmass_pose.inv() * self.agent.tcp.pose - tcp_rot_wrt_charger = tcp_pose_wrt_charger.to_transformation_matrix()[:3, :3] - gt_rot_1 = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]]) - gt_rot_2 = np.array([[-1, 0, 0], [0, 1, 0], [0, 0, -1]]) - grasp_rot_loss_fxn = lambda A: np.tanh( - 1 / 4 * np.trace(A.T @ A) - ) # trace(A.T @ A) has range [0,8] for A being difference of rotation matrices - grasp_rot_loss = np.minimum( - grasp_rot_loss_fxn(gt_rot_1 - tcp_rot_wrt_charger), - grasp_rot_loss_fxn(gt_rot_2 - tcp_rot_wrt_charger), - ) - reward += 2 * (1 - grasp_rot_loss) - rotated_properly = grasp_rot_loss < 0.1 - - if rotated_properly: - # reaching reward - gripper_to_obj_pos = cmass_pose.p - self.agent.tcp.pose.p - gripper_to_obj_dist = np.linalg.norm(gripper_to_obj_pos) - reaching_reward = 1 - np.tanh(5.0 * gripper_to_obj_dist) - reward += 2 * reaching_reward - - # grasp reward - is_grasped = self.agent.is_grasping( - self.charger, max_angle=20 - ) # max_angle ensures that the gripper grasps the charger appropriately, not in a strange pose - if is_grasped: - reward += 2.0 - - # pre-insertion and insertion award - if is_grasped: - pre_inserted = False - charger_cmass_wrt_goal = self.goal_pose.inv() * cmass_pose - charger_cmass_wrt_goal_yz_dist = np.linalg.norm( - charger_cmass_wrt_goal.p[1:] - ) - charger_cmass_wrt_goal_rot = ( - charger_cmass_wrt_goal.to_transformation_matrix()[:3, :3] - ) - charger_wrt_goal = self.goal_pose.inv() * self.charger.pose - charger_wrt_goal_yz_dist = np.linalg.norm(charger_wrt_goal.p[1:]) - charger_wrt_goal_dist = np.linalg.norm(charger_wrt_goal.p) - charger_wrt_goal_rot = ( - charger_cmass_wrt_goal.to_transformation_matrix()[:3, :3] - ) - - gt_rot = np.eye(3) - rot_loss_fxn = lambda A: np.tanh(1 / 2 * np.trace(A.T @ A)) - rot_loss = np.maximum( - rot_loss_fxn(charger_cmass_wrt_goal_rot - gt_rot), - rot_loss_fxn(charger_wrt_goal_rot - gt_rot), - ) - - pre_insertion_reward = 3 * ( - 1 - - np.tanh( - 1.0 - * (charger_cmass_wrt_goal_yz_dist + charger_wrt_goal_yz_dist) - + 9.0 - * np.maximum( - charger_cmass_wrt_goal_yz_dist, charger_wrt_goal_yz_dist - ) - ) - ) - pre_insertion_reward += 3 * (1 - np.tanh(3 * charger_wrt_goal_dist)) - pre_insertion_reward += 3 * (1 - rot_loss) - reward += pre_insertion_reward - - if ( - charger_cmass_wrt_goal_yz_dist < 0.01 - and charger_wrt_goal_yz_dist < 0.01 - and charger_wrt_goal_dist < 0.02 - and rot_loss < 0.15 - ): - pre_inserted = True - reward += 2.0 - - if pre_inserted: - insertion_reward = 2 * (1 - np.tanh(25.0 * charger_wrt_goal_dist)) - insertion_reward += 5 * ( - 1 - np.tanh(2.0 * np.abs(info["obj_to_goal_angle"])) - ) - insertion_reward += 5 * (1 - rot_loss) - reward += insertion_reward - else: - reward = reward - 10 * np.maximum( - self.charger.pose.p[2] - + self._base_size[2] / 2 - + 0.015 - - self.agent.tcp.pose.p[2], - 0.0, - ) - reward = reward - 10 * np.linalg.norm( - self.charger.pose.p[:2] - self.agent.tcp.pose.p[:2] - ) - - return reward - - def compute_normalized_dense_reward(self, **kwargs): - return self.compute_dense_reward(**kwargs) / 50.0 - - def _register_sensors(self): - cam_cfg = super()._register_sensors() - cam_cfg.pose = look_at([-0.3, 0, 0.1], [0, 0, 0.1]) - return cam_cfg - - def _register_human_render_cameras(self): - cam_cfg = super()._register_render_cameras() - cam_cfg.pose = look_at([-0.3, -0.4, 0.2], [0, 0, 0.1]) - return cam_cfg - - def _setup_lighting(self): - super()._setup_lighting() - if self.enable_shadow: - self._scene.add_point_light( - [-0.2, -0.5, 1], [2, 2, 2], shadow=self.enable_shadow - ) - - def set_state_dict(self, state): - super().set_state_dict(state) - self.goal_pose = self.receptacle.pose * Pose(q=euler2quat(0, 0, np.pi)) diff --git a/mani_skill2/envs/ms2/misc/__init__.py b/mani_skill2/envs/ms2/misc/__init__.py deleted file mode 100644 index 6a2efa9be..000000000 --- a/mani_skill2/envs/ms2/misc/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from . import avoid_obstacles, turn_faucet diff --git a/mani_skill2/envs/ms2/misc/assets/Sink_19.glb b/mani_skill2/envs/ms2/misc/assets/Sink_19.glb deleted file mode 100644 index 8035f974f..000000000 Binary files a/mani_skill2/envs/ms2/misc/assets/Sink_19.glb and /dev/null differ diff --git a/mani_skill2/envs/ms2/misc/avoid_obstacles.py b/mani_skill2/envs/ms2/misc/avoid_obstacles.py deleted file mode 100644 index b78950f40..000000000 --- a/mani_skill2/envs/ms2/misc/avoid_obstacles.py +++ /dev/null @@ -1,264 +0,0 @@ -from collections import OrderedDict -from pathlib import Path - -import numpy as np -import sapien -import sapien.physx as physx -from sapien import Pose - -from mani_skill2 import format_path -from mani_skill2.agents.robots.panda import PandaRealSensed435 -from mani_skill2.envs.sapien_env import BaseEnv -from mani_skill2.sensors.camera import CameraConfig -from mani_skill2.utils.building.ground import build_ground -from mani_skill2.utils.io_utils import load_json -from mani_skill2.utils.registration import register_env -from mani_skill2.utils.sapien_utils import ( - get_articulation_max_impulse_norm, - get_obj_by_name, - hide_entity, - look_at, - set_articulation_render_material, - show_entity, -) -from mani_skill2.utils.structs.pose import vectorize_pose - - -class AvoidObstaclesBaseEnv(BaseEnv): - DEFAULT_EPISODE_JSON: str - ASSET_UID: str - - tcp: physx.PhysxArticulationLinkComponent # Tool Center Point of the robot - - def __init__(self, episode_json=None, **kwargs): - if episode_json is None: - episode_json = self.DEFAULT_EPISODE_JSON - episode_json = format_path(episode_json) - if not Path(episode_json).exists(): - raise FileNotFoundError( - f"Episode json ({episode_json}) is not found." - "To download default json:" - "`python -m mani_skill2.utils.download_asset {}`.".format( - self.ASSET_UID - ) - ) - self.episodes = load_json(episode_json) - self.episode_idx = None - self.episode_config = None - super().__init__(**kwargs) - - def _set_scene_config(self): - super()._set_scene_config() - cfg = physx.get_scene_config() - cfg.contact_offset = 0.01 - physx.set_scene_config(cfg) - - def reset(self, *args, seed=None, options=None): - if options is None: - options = dict() - self._set_episode_rng(seed) - episode_idx = options.pop("episode_idx", None) - reconfigure = options.pop("reconfigure", False) - if episode_idx is None: - episode_idx = self._episode_rng.choice(len(self.episodes)) - if episode_idx != self.episode_idx: - reconfigure = True - self.episode_idx = episode_idx - self.episode_config = self.episodes[episode_idx] - options["reconfigure"] = reconfigure - return super().reset(*args, seed=self._episode_seed, options=options) - - def _build_cube( - self, - half_size, - color=(1, 0, 0), - name="cube", - static=True, - render_material: sapien.render.RenderMaterial = None, - ): - if render_material is None: - render_material = self._renderer.create_material() - render_material.set_base_color(np.hstack([color, 1.0])) - - builder = self._scene.create_actor_builder() - builder.add_box_collision(half_size=half_size) - builder.add_box_visual(half_size=half_size, material=render_material) - if static: - return builder.build_static(name) - else: - return builder.build(name) - - def _build_coord_frame_site(self, scale=0.1, name="coord_frame"): - builder = self._scene.create_actor_builder() - radius = scale * 0.05 - half_length = scale * 0.5 - builder.add_capsule_visual( - sapien.Pose(p=[scale * 0.5, 0, 0], q=[1, 0, 0, 0]), - radius=radius, - half_length=half_length, - material=sapien.render.RenderMaterial(base_color=[1, 0, 0, 1]), - name="x", - ) - builder.add_capsule_visual( - sapien.Pose(p=[0, scale * 0.5, 0], q=[0.707, 0, 0, 0.707]), - radius=radius, - half_length=half_length, - material=sapien.render.RenderMaterial(base_color=[0, 1, 0, 1]), - name="y", - ) - builder.add_capsule_visual( - sapien.Pose(p=[0, 0, scale * 0.5], q=[0.707, 0, -0.707, 0]), - radius=radius, - half_length=half_length, - material=sapien.render.RenderMaterial(base_color=[0, 0, 1, 1]), - name="z", - ) - actor = builder.build_static(name) - # NOTE(jigu): Must hide upon creation to avoid pollute observations! - hide_entity(actor) - return actor - - def _load_actors(self): - build_ground(self._scene) - - # Add a wall - if "wall" in self.episode_config: - cfg = self.episode_config["wall"] - self.wall = self._build_cube(cfg["half_size"], color=(1, 1, 1), name="wall") - self.wall.set_pose(Pose(cfg["pose"][:3], cfg["pose"][3:])) - - self.obstacles = [] - for i, cfg in enumerate(self.episode_config["obstacles"]): - actor = self._build_cube( - cfg["half_size"], self._episode_rng.rand(3), name=f"obstacle_{i}" - ) - actor.set_pose(Pose(cfg["pose"][:3], cfg["pose"][3:])) - self.obstacles.append(actor) - - self.goal_site = self._build_coord_frame_site(scale=0.05) - - def _initialize_agent(self): - qpos = self.episode_config["start_qpos"] - # qpos = self.episode_config["end_qpos"] - self.agent.reset(qpos) - self.agent.robot.set_pose(Pose([0, 0, 0])) - - def _update_goal_to_obstacle_dist(self): - obstacle_pos = [actor.pose.p for actor in self.obstacles] - goal_pos = self.goal_pose.p - goal_to_obstacle_dist = [np.linalg.norm(goal_pos - x) for x in obstacle_pos] - self.goal_to_obstacle_dist = np.sort(goal_to_obstacle_dist) - - def _initialize_task(self): - end_pose = self.episode_config["end_pose"] - self.goal_pose = Pose(end_pose[:3], end_pose[3:]) - self.goal_site.set_pose(self.goal_pose) - self._update_goal_to_obstacle_dist() - - def _get_obs_agent(self): - obs = self.agent.get_proprioception() - obs["base_pose"] = vectorize_pose(self.agent.robot.pose) - return obs - - def _get_obs_extra(self) -> OrderedDict: - tcp_pose = self.tcp.pose - goal_pose = self.goal_pose - return OrderedDict( - tcp_pose=vectorize_pose(tcp_pose), - goal_pose=vectorize_pose(goal_pose), - ) - - def evaluate(self, **kwargs) -> dict: - tcp_pose_at_goal = self.goal_pose.inv() * self.tcp.pose - pos_dist = np.linalg.norm(tcp_pose_at_goal.p) - ang_dist = np.arccos(np.clip(tcp_pose_at_goal.q[0], -1, 1)) * 2 - if ang_dist > np.pi: # [0, 2 * pi] -> [-pi, pi] - ang_dist = ang_dist - 2 * np.pi - ang_dist = np.abs(ang_dist) - ang_dist = np.rad2deg(ang_dist) - success = pos_dist <= 0.025 and ang_dist <= 15 - return dict(pos_dist=pos_dist, ang_dist=ang_dist, success=success) - - def compute_dense_reward(self, info, **kwargs): - if info["success"]: - return 10.0 - - pos_threshold = 0.025 - ang_threshold = 15 - reward = 0.0 - pos_dist, ang_dist = info["pos_dist"], info["ang_dist"] - num_obstacles = len(self.obstacles) - - close_to_goal_reward = ( - 4.0 * np.sum(pos_dist < self.goal_to_obstacle_dist) / num_obstacles - ) - # close_to_goal_reward += 1 - np.tanh(pos_dist) - angular_reward = 0.0 - - smallest_g2o_dist = self.goal_to_obstacle_dist[0] - if pos_dist < smallest_g2o_dist: - angular_reward = 3.0 * ( - 1 - np.tanh(np.maximum(ang_dist - ang_threshold, 0.0) / 180) - ) - if ang_dist <= 25: - close_to_goal_reward += 2.0 * ( - 1 - - np.tanh( - np.maximum(pos_dist - pos_threshold, 0.0) / smallest_g2o_dist - ) - ) - - contacts = self._scene.get_contacts() - max_impulse_norm = np.minimum( - get_articulation_max_impulse_norm(contacts, self.agent.robot), 2.0 - ) - reward = close_to_goal_reward + angular_reward - 50.0 * max_impulse_norm - return reward - - def compute_normalized_dense_reward(self, **kwargs): - return self.compute_dense_reward(**kwargs) / 10.0 - - def _register_sensors(self): - pose = look_at([-0.25, 0, 1.2], [0.6, 0, 0.6]) - 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([1.5, 0, 1.5], [0.0, 0.0, 0.5]) - return CameraConfig("render_camera", pose.p, pose.q, 512, 512, 1, 0.01, 10) - - def _setup_viewer(self): - super()._setup_viewer() - self._viewer.set_camera_xyz(1.5, 0.0, 1.5) - self._viewer.set_camera_rpy(0, -0.6, 3.14) - - def render_human(self): - show_entity(self.goal_site) - ret = super().render_human() - hide_entity(self.goal_site) - return ret - - def render_rgb_array(self): - show_entity(self.goal_site) - ret = super().render_rgb_array() - hide_entity(self.goal_site) - return ret - - -@register_env("PandaAvoidObstacles-v0", max_episode_steps=500) -class PandaAvoidObstaclesEnv(AvoidObstaclesBaseEnv): - DEFAULT_EPISODE_JSON = "{ASSET_DIR}/avoid_obstacles/panda_train_2k.json.gz" - ASSET_UID = "panda_avoid_obstacles" - - def _load_agent(self): - self.robot_uids = "panda_realsensed435" # TODO (stao): This old code here is not good practice to override user argument like this - self.agent = PandaRealSensed435( - self._scene, - self._control_freq, - self._control_mode, - ) - self.tcp: sapien.Entity = get_obj_by_name( - self.agent.robot.get_links(), self.agent.ee_link_name - ).entity - set_articulation_render_material(self.agent.robot, specular=0.9, roughness=0.3) diff --git a/mani_skill2/envs/ms2/misc/turn_faucet.py b/mani_skill2/envs/ms2/misc/turn_faucet.py deleted file mode 100644 index aad882000..000000000 --- a/mani_skill2/envs/ms2/misc/turn_faucet.py +++ /dev/null @@ -1,428 +0,0 @@ -from collections import OrderedDict -from pathlib import Path -from typing import Dict, List - -import numpy as np -import sapien -import sapien.physx as physx -import sapien.render -import trimesh -import trimesh.sample -from sapien import Pose -from scipy.spatial.distance import cdist -from transforms3d.euler import euler2quat - -from mani_skill2 import format_path -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.building.ground import build_ground -from mani_skill2.utils.common import np_random, random_choice -from mani_skill2.utils.geometry import transform_points -from mani_skill2.utils.geometry.trimesh_utils import get_component_mesh -from mani_skill2.utils.io_utils import load_json -from mani_skill2.utils.registration import register_env -from mani_skill2.utils.sapien_utils import ( - get_obj_by_name, - hex2rgba, - look_at, - set_articulation_render_material, -) -from mani_skill2.utils.structs.pose import vectorize_pose - - -class TurnFaucetBaseEnv(BaseEnv): - agent: Panda - - def __init__( - self, - *args, - robot_uids="panda", - robot_init_qpos_noise=0.02, - **kwargs, - ): - self.robot_init_qpos_noise = robot_init_qpos_noise - super().__init__(*args, robot_uids=robot_uids, **kwargs) - - def _load_actors(self): - # builder = self._scene.create_actor_builder() - # model_dir = Path(osp.dirname(__file__)) / "assets" - # scale = 1 - # collision_file = str(model_dir / "Sink_19.glb") # a metal table - # sink_pose = sapien.Pose(q=euler2quat(np.pi / 2, 0, 0)) - # builder.add_nonconvex_collision_from_file( - # filename=collision_file, scale=[scale] * 3, material=None, pose=sink_pose - # ) - # visual_file = str(model_dir / "Sink_19.glb") - # builder.add_visual_from_file( - # filename=visual_file, scale=[scale] * 3, pose=sink_pose - # ) - # self.sink = builder.build_static(name="sink") - # aabb = self.sink.find_component_by_type( - # sapien.render.RenderBodyComponent - # ).compute_global_aabb_tight() - # sink_height = aabb[1, 2] - aabb[0, 2] - - # self.sink.set_pose( - # Pose(p=[-0.24, 0, -sink_height], q=euler2quat(0, 0, -np.pi / 2)) - # ) - - build_ground(self._scene) - - # # add wall - # wall_mtl = sapien.render.RenderMaterial( - # base_color=[32 / 255, 67 / 255, 80 / 255, 1], - # metallic=0, - # roughness=0.9, - # specular=0.8, - # ) - # wall = self._scene.create_actor_builder() - # half_size = (0.02, 6, 2.1) - # wall.add_box_collision(half_size=half_size) - # wall.add_box_visual(half_size=half_size, material=wall_mtl) - # self.wall = wall.build_static("wall") - # self.wall.set_pose(Pose(p=[0.25, 0, 1])) - - def _initialize_agent(self): - if self.robot_uids == "panda": - # fmt: off - qpos = np.array([0, -0.785, 0, -2.356, 0, 1.57, 0.785, 0, 0]) - # fmt: on - qpos[:-2] += self._episode_rng.normal( - 0, self.robot_init_qpos_noise, len(qpos) - 2 - ) - self.agent.reset(qpos) - self.agent.robot.set_pose(Pose([-0.56, 0, 0])) - else: - raise NotImplementedError(self.robot_uids) - - def _get_obs_agent(self): - obs = self.agent.get_proprioception() - obs["base_pose"] = vectorize_pose(self.agent.robot.pose) - return obs - - def _register_sensors(self): - pose = look_at([-0.4, 0, 0.3], [0, 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([-1.3, 0.6, 0.6], [0.0, 0.0, 0.4]) - return CameraConfig("render_camera", pose.p, pose.q, 512, 512, 1, 0.01, 10) - - def _setup_viewer(self): - super()._setup_viewer() - self._viewer.set_camera_pose(look_at([-1.3, 0.6, 0.6], [0.0, 0.0, 0.4])) - - -@register_env("TurnFaucet-v0", max_episode_steps=200) -class TurnFaucetEnv(TurnFaucetBaseEnv): - target_link: physx.PhysxArticulationLinkComponent - target_joint: physx.PhysxJointComponent - - def __init__( - self, - asset_root: str = "{ASSET_DIR}/partnet_mobility/dataset", - model_json: str = "{PACKAGE_ASSET_DIR}/partnet_mobility/meta/info_faucet_train.json", - model_ids: List[str] = (), - **kwargs, - ): - self.asset_root = Path(format_path(asset_root)) - - model_json = format_path(model_json) - self.model_db: Dict[str, Dict] = load_json(model_json) - - if isinstance(model_ids, str): - model_ids = [model_ids] - if len(model_ids) == 0: - model_ids = sorted(self.model_db.keys()) - assert len(model_ids) > 0, model_json - - # model_ids = list(map(str, model_ids)) - self.model_ids = model_ids - self.model_id = None - self.model_scale = None - - # Find and check urdf paths - self.model_urdf_paths = {} - for model_id in self.model_ids: - self.model_urdf_paths[model_id] = self.find_urdf_path(model_id) - - super().__init__(**kwargs) - - def find_urdf_path(self, model_id): - model_dir = self.asset_root / str(model_id) - - urdf_names = ["mobility_cvx.urdf"] - for urdf_name in urdf_names: - urdf_path = model_dir / urdf_name - if urdf_path.exists(): - return urdf_path - - raise FileNotFoundError( - f"No valid URDF is found for {model_id}." - "Please download Partnet-Mobility (ManiSkill2):" - "`python -m mani_skill2.utils.download_asset partnet_mobility_faucet`." - ) - - def reset(self, seed=None, options=None): - if options is None: - options = dict() - self._set_episode_rng(seed) - model_id = options.pop("model_id", None) - model_scale = options.pop("model_scale", None) - reconfigure = options.pop("reconfigure", False) - - _reconfigure = self._set_model(model_id, model_scale) - reconfigure = _reconfigure or reconfigure - options["reconfigure"] = reconfigure - return super().reset(seed=self._episode_seed, options=options) - - def _set_model(self, model_id, model_scale): - """Set the model id and scale. If not provided, choose one randomly.""" - reconfigure = False - - # Model ID - if model_id is None: - model_id = random_choice(self.model_ids, self._episode_rng) - if model_id != self.model_id: - reconfigure = True - self.model_id = model_id - self.model_info = self.model_db[self.model_id] - - # Scale - if model_scale is None: - model_scale = self.model_info.get("scale") - if model_scale is None: - bbox = self.model_info["bbox"] - bbox_size = np.float32(bbox["max"]) - np.float32(bbox["min"]) - model_scale = 0.3 / max(bbox_size) # hardcode - if model_scale != self.model_scale: - reconfigure = True - self.model_scale = model_scale - - if "offset" in self.model_info: - self.model_offset = np.float32(self.model_info["offset"]) - else: - self.model_offset = -np.float32(bbox["min"]) * model_scale - # Add a small clearance - self.model_offset[2] += 0.01 - - return reconfigure - - def _load_articulations(self): - self.faucet: physx.PhysxArticulation = self._load_faucet() - # Cache qpos to restore - self._faucet_init_qpos = self.faucet.get_qpos() - - # Set friction and damping for all joints - for joint in self.faucet.get_active_joints(): - joint.set_friction(1.0) - joint.set_drive_property(0.0, 10.0) - - self._set_switch_links() - - def _load_faucet(self): - loader = self._scene.create_urdf_loader() - loader.scale = self.model_scale - loader.fix_root_link = True - - model_dir = self.asset_root / str(self.model_id) - urdf_path = model_dir / "mobility_cvx.urdf" - - density = self.model_info.get("density", 8e3) - articulation: physx.PhysxArticulation = loader.load(str(urdf_path)) - loader.set_density(density) - articulation.set_name("faucet") - - # TODO (stao): find out why we did this before - set_articulation_render_material( - articulation, color=hex2rgba("#AAAAAA"), metallic=1, roughness=0.4 - ) - - return articulation - - def _set_switch_links(self): - switch_link_names = [] - for semantic in self.model_info["semantics"]: - if semantic[2] == "switch": - switch_link_names.append(semantic[0]) - if len(switch_link_names) == 0: - raise RuntimeError(self.model_id) - self.switch_link_names = switch_link_names - - self.switch_links = [] - self.switch_links_mesh: List[trimesh.Trimesh] = [] - self.switch_joints = [] - all_links = self.faucet.get_links() - all_joints = self.faucet.get_joints() - - for name in self.switch_link_names: - link = get_obj_by_name(all_links, name) - self.switch_links.append(link) - - # cache mesh - link_mesh = get_component_mesh(link, False) - self.switch_links_mesh.append(link_mesh) - - # hardcode - joint = all_joints[link.get_index()] - joint.set_friction(0.1) - joint.set_drive_property(0.0, 2.0) - self.switch_joints.append(joint) - - def _load_agent(self): - super()._load_agent() - - links = self.agent.robot.get_links() - self.lfinger = get_obj_by_name(links, "panda_leftfinger") - self.rfinger = get_obj_by_name(links, "panda_rightfinger") - self.lfinger_mesh = get_component_mesh(self.lfinger, False) - self.rfinger_mesh = get_component_mesh(self.rfinger, False) - - def _initialize_articulations(self): - p = np.zeros(3) - p[:2] = self._episode_rng.uniform(-0.05, 0.05, [2]) - p[2] = self.model_offset[2] - ori = self._episode_rng.uniform(-np.pi / 12, np.pi / 12) - q = euler2quat(0, 0, ori) - self.faucet.set_pose(Pose(p, q)) - # self.sink.set_pose(self.sink.pose * Pose(p=[p[0], p[1], 0.011], q=q)) - # self.wall.set_pose(self.wall.pose * Pose(p=[p[0], p[1], 0], q=q)) - - def _initialize_task(self): - self._set_target_link() - self._set_init_and_target_angle() - - qpos = self._faucet_init_qpos.copy() - qpos[self.target_joint_idx] = self.init_angle - self.faucet.set_qpos(qpos) - - # -------------------------------------------------------------------------- # - # For dense reward - # -------------------------------------------------------------------------- # - # NOTE(jigu): trimesh uses np.random to sample - self.lfinger_pcd = trimesh.sample.sample_surface( - self.lfinger_mesh, 256, seed=self._episode_seed - )[0] - self.rfinger_pcd = trimesh.sample.sample_surface( - self.rfinger_mesh, 256, seed=self._episode_seed - )[0] - - self.last_angle_diff = self.target_angle - self.current_angle - - def _set_target_link(self): - n_switch_links = len(self.switch_link_names) - idx = random_choice(np.arange(n_switch_links), self._episode_rng) - - self.target_link_name = self.switch_link_names[idx] - self.target_link: physx.PhysxArticulationLinkComponent = self.switch_links[idx] - self.target_joint: physx.PhysxArticulationJoint = self.switch_joints[idx] - self.target_joint_idx = self.faucet.get_active_joints().index(self.target_joint) - - # x-axis is the revolute joint direction - - assert ( - self.target_joint.type == "revolute_unwrapped" - or self.target_joint.type == "revolute" - ), self.target_joint.type - joint_pose = self.target_joint.get_global_pose().to_transformation_matrix() - self.target_joint_axis = joint_pose[:3, 0] - - self.target_link_mesh: trimesh.Trimesh = self.switch_links_mesh[idx] - - # NOTE(jigu): trimesh uses np.random to sample - self.target_link_pcd = trimesh.sample.sample_surface( - self.target_link_mesh, 256, seed=self._episode_seed - )[0] - - # NOTE(jigu): joint origin can be anywhere on the joint axis. - # Thus, I use the center of mass at the beginning instead - cmass_pose = self.target_link.pose * self.target_link.cmass_local_pose - self.target_link_pos = cmass_pose.p - - def _set_init_and_target_angle(self): - qmin, qmax = self.target_joint.get_limit()[0] - if np.isinf(qmin): - self.init_angle = 0 - else: - self.init_angle = qmin - if np.isinf(qmax): - # maybe we can count statics - self.target_angle = np.pi / 2 - else: - self.target_angle = qmin + (qmax - qmin) * 0.9 - - # The angle to go - self.target_angle_diff = self.target_angle - self.init_angle - - def _get_obs_extra(self) -> OrderedDict: - obs = OrderedDict( - tcp_pose=vectorize_pose(self.agent.tcp.pose), - target_angle_diff=np.array(self.target_angle_diff), - target_joint_axis=self.target_joint_axis, - target_link_pos=self.target_link_pos, - ) - if self._obs_mode in ["state", "state_dict"]: - angle_dist = self.target_angle - self.current_angle - obs["angle_dist"] = np.array(angle_dist) - return obs - - @property - def current_angle(self): - return self.faucet.get_qpos()[self.target_joint_idx] - - def evaluate(self, **kwargs): - angle_dist = self.target_angle - self.current_angle - return dict(success=angle_dist < 0, angle_dist=angle_dist) - - def _compute_distance(self): - """Compute the distance between the tap and robot fingers.""" - T = self.target_link.pose.to_transformation_matrix() - pcd = transform_points(T, self.target_link_pcd) - T1 = self.lfinger.pose.to_transformation_matrix() - T2 = self.rfinger.pose.to_transformation_matrix() - pcd1 = transform_points(T1, self.lfinger_pcd) - pcd2 = transform_points(T2, self.rfinger_pcd) - # trimesh.PointCloud(np.vstack([pcd, pcd1, pcd2])).show() - distance1 = cdist(pcd, pcd1) - distance2 = cdist(pcd, pcd2) - - return min(distance1.min(), distance2.min()) - - def compute_dense_reward(self, info, **kwargs): - reward = 0.0 - - if info["success"]: - return 10.0 - - distance = self._compute_distance() - reward += 1 - np.tanh(distance * 5.0) - - angle_diff = self.target_angle - self.current_angle - turn_reward_1 = 3 * (1 - np.tanh(max(angle_diff, 0) * 2.0)) - reward += turn_reward_1 - - delta_angle = angle_diff - self.last_angle_diff - if angle_diff > 0: - turn_reward_2 = -np.tanh(delta_angle * 2) - else: - turn_reward_2 = np.tanh(delta_angle * 2) - turn_reward_2 *= 5 - reward += turn_reward_2 - self.last_angle_diff = angle_diff - - return reward - - def compute_normalized_dense_reward(self, **kwargs): - return self.compute_dense_reward(**kwargs) / 10.0 - - def get_state_dict(self) -> np.ndarray: - state = super().get_state_dict() - return np.hstack([state, self.target_angle]) - - def set_state_dict(self, state): - self.target_angle = state[-1] - super().set_state_dict(state[:-1]) - self.last_angle_diff = self.target_angle - self.current_angle diff --git a/mani_skill2/envs/ms2/mpm/.gitignore b/mani_skill2/envs/ms2/mpm/.gitignore deleted file mode 100644 index cf011ea61..000000000 --- a/mani_skill2/envs/ms2/mpm/.gitignore +++ /dev/null @@ -1 +0,0 @@ -*.sdf \ No newline at end of file diff --git a/mani_skill2/envs/ms2/mpm/RopeInit.pkl b/mani_skill2/envs/ms2/mpm/RopeInit.pkl deleted file mode 100644 index 4b9f5d5ca..000000000 Binary files a/mani_skill2/envs/ms2/mpm/RopeInit.pkl and /dev/null differ diff --git a/mani_skill2/envs/ms2/mpm/__init__.py b/mani_skill2/envs/ms2/mpm/__init__.py deleted file mode 100644 index 2a7eb87e3..000000000 --- a/mani_skill2/envs/ms2/mpm/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# isort: off -from . import excavate_env -from . import fill_env -from . import pour_env -from . import hang_env -from . import write_env -from . import pinch_env diff --git a/mani_skill2/envs/ms2/mpm/base_env.py b/mani_skill2/envs/ms2/mpm/base_env.py deleted file mode 100644 index 991342fe2..000000000 --- a/mani_skill2/envs/ms2/mpm/base_env.py +++ /dev/null @@ -1,673 +0,0 @@ -import ctypes -import os -import sys -from typing import Dict, Union - -import mani_skill2 -from mani_skill2 import PACKAGE_ASSET_DIR - -warp_path = os.path.join(os.path.dirname(mani_skill2.__file__), "..", "warp_maniskill") -warp_path = os.path.normpath(warp_path) -if warp_path not in sys.path: - sys.path.append(warp_path) - -# build warp if warp.so does not exist -from warp_maniskill.build_lib import build, build_path - -dll = os.path.join(build_path, "bin/warp.so") - -import hashlib -from collections import OrderedDict - -import numpy as np -import sapien -import trimesh -import warp as wp -from transforms3d.euler import euler2quat -from transforms3d.quaternions import quat2mat - -from mani_skill2.envs.ms2.mpm.utils import actor2meshes, trimesh2sdf -from mani_skill2.envs.sapien_env import BaseEnv -from mani_skill2.utils.logging_utils import logger -from mani_skill2.utils.structs.pose import vectorize_pose -from warp_maniskill.mpm.mpm_model import MPMModelBuilder -from warp_maniskill.mpm.mpm_simulator import DenseVolume as MPMVolume -from warp_maniskill.mpm.mpm_simulator import Mesh as MPMMesh -from warp_maniskill.mpm.mpm_simulator import Simulator as MPMSimulator - - -def task(meshes): - if meshes is None or len(meshes) == 0: - return None - bbox = trimesh.util.concatenate(meshes).bounds - length = np.max(bbox[1] - bbox[0]) - dx = min(0.01, length / 40) # dx should at most be 0.01 - margin = max( - dx * 3, 0.01 - ) # margin should be greater than 3 dx and at least be 0.01 - return trimesh2sdf(meshes, margin, dx) - - -class MPMBaseEnv(BaseEnv): - # fmt: off - SUPPORTED_OBS_MODES = ("none", "image") - # fmt: on - - def __init__( - self, - *args, - sim_freq=500, - mpm_freq=2000, - max_particles=65536, - **kwargs, - ): - if not os.path.isfile(dll): - - class ARGS: - verbose = False - mode = "release" - cuda_path = None - verify_fp = False - - build(ARGS) - - wp.init() - - if "shader_dir" in kwargs: - logger.warning("`shader_dir` is ignored for soft-body environments.") - kwargs.pop("shader_dir") - shader_dir = os.path.join(os.path.dirname(__file__), "shader", "point") - - self.sim_crashed = False - - self._mpm_step_per_sapien_step = mpm_freq // sim_freq - self._mpm_dt = 1 / mpm_freq - self.max_particles = max_particles - - self.sdf_cache = os.path.join( - os.path.dirname(__file__), self.__class__.__name__ + ".sdf" - ) - - super().__init__(*args, shader_dir=shader_dir, **kwargs) - - # ---------------------------------------------------------------------------- # - # Setup - # ---------------------------------------------------------------------------- # - - def reconfigure(self): - self._clear() - - self._setup_scene() - self._load_agent() - self._load_actors() - self._load_articulations() - self._setup_sensors() - self._setup_lighting() - - self._setup_mpm() - self._setup_render_particles() - - if self._viewer is not None: - self._setup_viewer() - - # Cache actors and articulations - self._actors = self.get_actors() - self._articulations = self.get_articulations() - - def _load_actors(self): - self._scene.add_ground(altitude=0.0, render=False) - if self.bg_name is None: - b = self._scene.create_actor_builder() - b.add_visual_from_file(str(PACKAGE_ASSET_DIR / "maniskill2-scene-2.glb")) - b.build_kinematic() - - def _get_coupling_actors( - self, - ): - return self.agent.robot.get_links() - - def _setup_mpm(self): - self.model_builder = MPMModelBuilder() - self.model_builder.set_mpm_domain( - domain_size=[0.5, 0.5, 0.5], grid_length=0.005 - ) - self.model_builder.reserve_mpm_particles(count=self.max_particles) - - self._setup_mpm_bodies() - - self.mpm_simulator = MPMSimulator(device="cuda") - self.mpm_model = self.model_builder.finalize(device="cuda") - self.mpm_model.gravity = np.array((0.0, 0.0, -9.81), dtype=np.float32) - self.mpm_model.struct.ground_normal = wp.vec3(0.0, 0.0, 1.0) - self.mpm_model.struct.particle_radius = 0.005 - self.mpm_states = [ - self.mpm_model.state() for _ in range(self._mpm_step_per_sapien_step + 1) - ] - - def _setup_mpm_bodies(self): - # convert actors to meshes - self._coupled_actors = [] - actor_meshes = [] - actor_primitives = [] - for actor in self._get_coupling_actors(): - visual = False - if isinstance(actor, tuple) or isinstance(actor, list): - type = actor[1] - actor = actor[0] - if type == "visual": - visual = True - - meshes, primitives = actor2meshes( - actor, visual=visual, return_primitives=True - ) - if meshes or primitives: - self._coupled_actors.append(actor) - actor_meshes.append(meshes) - actor_primitives.append(primitives) - - if not self._coupled_actors: - return - - # compute signature for current meshes - hash = hashlib.sha256() - hash.update(bytes("v2", encoding="utf-8")) - for meshes in actor_meshes: - for m in meshes: - hash.update(m.vertices.tobytes()) - hash.update(m.faces.tobytes()) - signature = hash.digest() - sdfs = None - - # load cached sdfs or recompute sdfs - import pickle - - if os.path.isfile(self.sdf_cache): - with open(self.sdf_cache, "rb") as f: - cache = pickle.load(f) - if "signature" in cache and cache["signature"] == signature: - logger.info("load sdfs from file") - sdfs = cache["sdfs"] - if len(sdfs) != len(actor_meshes): - sdfs = None - - if sdfs is None: - from multiprocessing import Pool - - import tqdm - - print("generating cached SDF volumes") - with Pool(8) as p: - sdfs = list( - tqdm.tqdm( - p.imap(task, actor_meshes), - total=len(actor_meshes), - ) - ) - - with open(self.sdf_cache, "wb") as f: - meshes = [ - [(np.array(m.vertices), np.array(m.faces)) for m in ms] - for ms in actor_meshes - ] - pickle.dump({"signature": signature, "sdfs": sdfs, "meshes": meshes}, f) - - # convert sdfs to dense volumes - for actor, sdf, meshes, primitives in zip( - self._coupled_actors, sdfs, actor_meshes, actor_primitives - ): - id = self.model_builder.add_body( - origin=wp.transform((0.0, 0.0, 0.0), wp.quat_identity()) - ) - if sdf is not None: - data = np.concatenate( - [sdf["normal"], sdf["sdf"].reshape(sdf["sdf"].shape + (1,))], -1 - ) - volume = MPMVolume( - data, - sdf["position"], - sdf["scale"], - mass=1, - I=np.eye(3), - com=np.zeros(3), - ) - self.model_builder.add_shape_dense_volume(id, volume=volume) - for type, params, pose in primitives: - if type == "box": - self.model_builder.add_shape_box( - id, - pos=tuple(pose.p), - rot=(pose.q[3], *pose.q[:3]), - hx=params[0], - hy=params[1], - hz=params[2], - ) - if type == "capsule": - self.model_builder.add_shape_capsule( - id, - pos=tuple(pose.p), - rot=(pose.q[3], *pose.q[:3]), - radius=params[0], - half_width=params[1], - ) - - cmass = actor.cmass_local_pose - R = quat2mat(cmass.q) - self.model_builder.set_body_mass( - id, actor.mass, R @ np.diag(actor.inertia) @ R.T, cmass.p - ) - - def _setup_render_particles(self): - N = self.mpm_model.max_n_particles - scales = np.ones(N, dtype=np.float32) * 0.005 - colors = np.zeros((N, 4)) - colors[:, :3] = self.mpm_model.mpm_particle_colors - - self.pcd = self._scene.add_particle_entity(np.zeros((N, 3), dtype=np.float32)) - self.pcd.visual_body.set_attribute("color", colors.astype(np.float32)) - self.pcd.visual_body.set_attribute("scale", scales) - - self.vertices_ptr = sapien.pysapien.dlpack.dl_ptr( - self.pcd.visual_body.dl_vertices - ) - self.vertices_shape = sapien.pysapien.dlpack.dl_shape( - self.pcd.visual_body.dl_vertices - ) - - def _setup_viewer(self): - super()._setup_viewer() - self._viewer.set_camera_xyz(1.0, 0.0, 1.2) - self._viewer.set_camera_rpy(0, -0.5, 3.14) - - def _clear(self): - super()._clear() - - self.vertices_ptr = None - self.vertices_shape = None - self.mpm_model = None - self.mpm_states = None - self.pcd = None - - # ---------------------------------------------------------------------------- # - # Initialization - # ---------------------------------------------------------------------------- # - - def reset(self, *args, **kwargs): - self.sim_crashed = False - return super().reset(*args, **kwargs) - - def initialize_episode(self): - super().initialize_episode() - self._initialize_mpm() - self._initialize_render_particles() - - def _initialize_mpm(self): - self.model_builder.clear_particles() - - E = 1e5 - nu = 0.3 - mu, lam = E / (2 * (1 + nu)), E * nu / ((1 + nu) * (1 - 2 * nu)) - - # 0 for von-mises, 1 for drucker-prager - type = 1 - - # von-mises - ys = 1e4 - - # drucker-prager - friction_angle = 1.0 - cohesion = 0.1 - - x = 0.05 - y = 0.05 - z = 0.05 - cell_x = 0.005 - self.model_builder.add_mpm_grid( - pos=(-0.0, 0.0, 0.05), - vel=(0.0, 0.0, 0.0), - dim_x=int(x // cell_x), - dim_y=int(y // cell_x), - dim_z=int(z // cell_x), - cell_x=cell_x, - cell_y=cell_x, - cell_z=cell_x, - density=3.0e3, - mu_lambda_ys=(mu, lam, ys), - friction_cohesion=(friction_angle, cohesion, 0.0), - type=type, - jitter=True, - placement_x="center", - placement_y="center", - placement_z="start", - color=(0.65237011, 0.14198029, 0.02201299), - random_state=self._episode_rng, - ) - self.model_builder.init_model_state(self.mpm_model, self.mpm_states) - self.mpm_model.struct.static_ke = 100.0 - self.mpm_model.struct.static_kd = 0.0 - self.mpm_model.struct.static_mu = 1.0 - self.mpm_model.struct.static_ka = 0.0 - - self.mpm_model.struct.body_ke = 100.0 - self.mpm_model.struct.body_kd = 0.0 - self.mpm_model.struct.body_mu = 1.0 - self.mpm_model.struct.body_ka = 0.0 - - self.mpm_model.adaptive_grid = False - - self.mpm_model.struct.body_sticky = 1 - self.mpm_model.struct.ground_sticky = 1 - self.mpm_model.particle_contact = True - self.mpm_model.grid_contact = True - - def _initialize_render_particles(self): - n = self.mpm_model.struct.n_particles - - self.pcd.visual_body.set_rendered_point_count(n) - - buffer = np.empty((n, 4), dtype=np.float32) - buffer[:, 0] = self.mpm_model.struct.particle_radius * 1.0 - buffer[:, 1:] = self.mpm_model.mpm_particle_colors - - # transfer - scale_colors = wp.array(buffer, dtype=float, device="cuda") - width = spitch = 16 - dpitch = self.vertices_shape[1] * 4 - height = n - wp.context.runtime.core.memcpy2d_d2d( - ctypes.c_void_p(self.vertices_ptr + 12), - ctypes.c_size_t(dpitch), - ctypes.c_void_p(scale_colors.ptr), - ctypes.c_size_t(spitch), - ctypes.c_size_t(width), - ctypes.c_size_t(height), - ) - wp.synchronize() - - # ---------------------------------------------------------------------------- # - # Observation - # ---------------------------------------------------------------------------- # - def get_obs(self): - mpm_state = self.get_mpm_state() - - if self._obs_mode == "particles": - obs = OrderedDict( - particles=mpm_state, - agent=self._get_obs_agent(), - extra=self._get_obs_extra(), - ) - else: - obs = super().get_obs() - - self._last_obs = obs - return obs - - def _get_obs_agent(self): - obs = self.agent.get_proprioception() - obs["base_pose"] = vectorize_pose(self.agent.robot.pose) - return obs - - def copy_array_to_numpy(self, array, length): - dtype = np.dtype(array.__cuda_array_interface__["typestr"]) - shape = array.__cuda_array_interface__["shape"] - assert shape[0] >= length - result = np.empty((length,) + shape[1:], dtype=dtype) - assert result.strides == array.__cuda_array_interface__["strides"] - - wp.context.runtime.core.memcpy_d2h( - ctypes.c_void_p(result.__array_interface__["data"][0]), - ctypes.c_void_p(array.ptr), - ctypes.c_size_t(result.strides[0] * length), - ) - - return result - - # -------------------------------------------------------------------------- # - # Simulation state - # -------------------------------------------------------------------------- # - def get_mpm_state(self): - n = self.mpm_model.struct.n_particles - - return OrderedDict( - x=self.copy_array_to_numpy(self.mpm_states[0].struct.particle_q, n), - v=self.copy_array_to_numpy(self.mpm_states[0].struct.particle_qd, n), - F=self.copy_array_to_numpy(self.mpm_states[0].struct.particle_F, n), - C=self.copy_array_to_numpy(self.mpm_states[0].struct.particle_C, n), - vc=self.copy_array_to_numpy( - self.mpm_states[0].struct.particle_volume_correction, n - ), - ) - - def set_mpm_state(self, state): - self.mpm_states[0].struct.particle_q.assign(state["x"]) - self.mpm_states[0].struct.particle_qd.assign(state["v"]) - self.mpm_states[0].struct.particle_F.assign(state["F"]) - self.mpm_states[0].struct.particle_C.assign(state["C"]) - self.mpm_states[0].struct.particle_volume_correction.assign(state["vc"]) - - def get_sim_state(self): - sapien_state = super().get_sim_state() - mpm_state = self.get_mpm_state() - return OrderedDict(sapien=sapien_state, mpm=mpm_state) - - def set_sim_state(self, state): - super().set_sim_state(state["sapien"]) - self.set_mpm_state(state["mpm"]) - - def get_state_dict(self): - """Get environment state. Override to include task information (e.g., goal)""" - # with sapien.ProfilerBlock("get_state"): - ret_state = [] - sim_state = self.get_sim_state() - for key, value in sim_state.items(): - if key == "mpm": - for key, value in value.items(): - ret_state.append(value.reshape((1, -1)).squeeze()) - else: - ret_state.append(value.reshape((1, -1)).squeeze()) - - ret_state = np.concatenate(ret_state) - return ret_state - - @property - def n_particles(self): - return self.mpm_model.struct.n_particles - - def set_state_dict(self, state: np.ndarray): - """Set environment state. Override to include task information (e.g., goal)""" - sim_state = OrderedDict() - mpm_state = OrderedDict() - n = self.mpm_model.struct.n_particles - - sim_state["sapien"] = state[: -n * 25] - mpm_state["x"] = state[-n * 25 : -n * 22].reshape((n, 3)) - mpm_state["v"] = state[-n * 22 : -n * 19].reshape((n, 3)) - mpm_state["F"] = state[-n * 19 : -n * 10].reshape((n, 3, 3)) - mpm_state["C"] = state[-n * 10 : -n].reshape((n, 3, 3)) - mpm_state["vc"] = state[-n:].reshape((n,)) - sim_state["mpm"] = mpm_state - - return self.set_sim_state(sim_state) - - # -------------------------------------------------------------------------- # - # Visualization - # -------------------------------------------------------------------------- # - def update_render(self): - width = spitch = 12 - dpitch = self.vertices_shape[1] * 4 - height = self.mpm_model.struct.n_particles - wp.context.runtime.core.memcpy2d_d2d( - ctypes.c_void_p(self.vertices_ptr), - ctypes.c_size_t(dpitch), - ctypes.c_void_p(self.mpm_states[0].struct.particle_q.ptr), - ctypes.c_size_t(spitch), - ctypes.c_size_t(width), - ctypes.c_size_t(height), - ) - wp.synchronize() - super().update_render() - - def _setup_lighting(self): - self._scene.set_ambient_light([0.3, 0.3, 0.3]) - self._scene.add_directional_light( - [1, 1, -1], [1, 1, 1], shadow=True, scale=5, shadow_map_size=2048 - ) - self._scene.add_directional_light([0, 0, -1], [1, 1, 1]) - - # -------------------------------------------------------------------------- # - # Step - # -------------------------------------------------------------------------- # - - def sync_actors(self): - if not self._coupled_actors: - return - - body_q = np.empty((len(self._coupled_actors), 7), dtype=np.float32) - body_qd = np.empty((len(self._coupled_actors), 6), dtype=np.float32) - for i, a in enumerate(self._coupled_actors): - pose = a.pose - p = pose.p - q = pose.q - body_q[i, :3] = p - # different quaternion convention - body_q[i, 3:6] = q[1:] - body_q[i, 6] = q[0] - - body_qd[i, :3] = a.angular_velocity - body_qd[i, 3:] = a.velocity - - for s in self.mpm_states: - s.body_q.assign(body_q) - s.body_qd.assign(body_qd) - - def sync_actors_prepare(self): - if not self._coupled_actors: - return - - self.body_q = np.empty((len(self._coupled_actors), 7), dtype=np.float32) - self.body_qd = np.empty((len(self._coupled_actors), 6), dtype=np.float32) - for i, a in enumerate(self._coupled_actors): - pose = a.pose - p = pose.p - q = pose.q - self.body_q[i, :3] = p - # different quaternion convention - self.body_q[i, 3:6] = q[1:] - self.body_q[i, 6] = q[0] - - self.body_qd[i, :3] = a.angular_velocity - self.body_qd[i, 3:] = a.velocity - - def sync_actors_state(self, state): - state.body_q.assign(self.body_q) - state.body_qd.assign(self.body_qd) - - def step(self, action: Union[None, np.ndarray, Dict]): - if not self.sim_crashed: - obs, rew, terminated, truncated, info = super().step(action) - info["crashed"] = False - return obs, rew, terminated, truncated, info - - logger.warn("simulation has crashed!") - info = self.get_info(obs=self._last_obs) - info["crashed"] = True - return self._last_obs, -10, True, False, info - - def step_action(self, action: np.ndarray): - if action is None: - pass - elif isinstance(action, np.ndarray): - self.agent.set_action(action) - elif isinstance(action, dict): - if action["control_mode"] != self.agent.control_mode: - self.agent.set_control_mode(action["control_mode"]) - self.agent.set_action(action["action"]) - else: - raise TypeError(type(action)) - - for _ in range(self._sim_steps_per_control): - self.sync_actors() - for mpm_step in range(self._mpm_step_per_sapien_step): - self.mpm_simulator.simulate( - self.mpm_model, - self.mpm_states[mpm_step], - self.mpm_states[mpm_step + 1], - self._mpm_dt, - ) - - self.agent.before_simulation_step() - - # apply wrench - tfs = [s.ext_body_f.numpy() for s in self.mpm_states[:-1]] - tfs = np.mean(tfs, 0) - - if np.isnan(tfs).any(): - self.sim_crashed = True - return - - if self.mpm_states[-1].struct.error.numpy()[0] == 1: - self.sim_crashed = True - return - - for actor, tf in zip(self._coupled_actors, tfs): - if actor.type not in ["kinematic", "static"]: - actor.add_force_torque(tf[3:], tf[:3]) - - self._scene.step() - self.mpm_states = [self.mpm_states[-1]] + self.mpm_states[ - :-1 - ] # rotate states - - # -------------------------------------------------------------------------- # - # Utils - # -------------------------------------------------------------------------- # - def _get_bbox(self, percentile=None): - x = self.mpm_states[0].struct.particle_q.numpy() - if percentile is None: - return np.array([x.min(0), x.max(0)]) - assert percentile < 50 - return np.array( - [np.percentile(x, percentile, 0), np.percentile(x, 100 - percentile, 0)] - ) - - def _add_draw_box(self, bbox): - from sapien import renderer as R - - context: R.Context = self._renderer._internal_context - # fmt:off - lines = context.create_line_set( - [ - -1, -1, -1, -1, -1, 1, - -1, -1, -1, -1, 1, -1, - -1, 1, 1, -1, -1, 1, - -1, 1, 1, -1, 1, -1, - 1, -1, -1, 1, -1, 1, - 1, -1, -1, 1, 1, -1, - 1, 1, 1, 1, -1, 1, - 1, 1, 1, 1, 1, -1, - 1, 1, 1, -1, 1, 1, - 1, -1, 1, -1, -1, 1, - 1, 1, -1, -1, 1, -1, - 1, -1, -1, -1, -1, -1, - ], - [1, 1, 1, 1] * 24, - ) - # fmt:on - lineset: R.LineSetObject = ( - self._scene.get_renderer_scene()._internal_scene.add_line_set(lines) - ) - - lineset.set_scale((bbox[1] - bbox[0]) / 2) - lineset.set_position((bbox[1] + bbox[0]) / 2) - return lineset - - def _remove_draw_box(self, lineset): - self._scene.get_renderer_scene()._internal_scene.remove_node(lineset) - - def render(self, draw_box=False): - if draw_box: - bbox = self._get_bbox(5) - box = self._add_draw_box(bbox) - img = super().render() - if draw_box: - self._remove_draw_box(box) - return img diff --git a/mani_skill2/envs/ms2/mpm/excavate_env.py b/mani_skill2/envs/ms2/mpm/excavate_env.py deleted file mode 100644 index bc56be320..000000000 --- a/mani_skill2/envs/ms2/mpm/excavate_env.py +++ /dev/null @@ -1,449 +0,0 @@ -from collections import OrderedDict - -import numpy as np -import sapien -from transforms3d.euler import euler2quat - -from mani_skill2.agents.robots.panda.variants import PandaBucket -from mani_skill2.envs.ms2.mpm import perlin -from mani_skill2.envs.ms2.mpm.base_env import MPMBaseEnv -from mani_skill2.envs.ms2.mpm.utils import actor2meshes -from mani_skill2.sensors.camera import CameraConfig -from mani_skill2.utils.registration import register_env -from mani_skill2.utils.sapien_utils import get_obj_by_name -from mani_skill2.utils.structs.pose import vectorize_pose - - -@register_env("Excavate-v0", max_episode_steps=250) -class ExcavateEnv(MPMBaseEnv): - def __init__( - self, - *args, - sim_freq=500, - mpm_freq=2000, - **kwargs, - ): - super().__init__( - *args, - sim_freq=sim_freq, - mpm_freq=mpm_freq, - **kwargs, - ) - - def _initialize_mpm(self): - self.model_builder.clear_particles() - - E = 1e4 - nu = 0.3 - mu, lam = E / (2 * (1 + nu)), E * nu / ((1 + nu) * (1 - 2 * nu)) - - # 0 for von-mises, 1 for drucker-prager - type = 1 - - # von-mises - ys = 1e4 - - # drucker-prager - friction_angle = 0.6 - cohesion = 0.05 - - height_map = 0.06 + perlin.added_perlin( - [0.03, 0.02, 0.02], - [1, 2, 4], - phases=[(0, 0), (0, 0), (0, 0)], - shape=(30, 30), - random_state=self._episode_rng, - ) - - count = self.model_builder.add_mpm_from_height_map( - pos=(0.0, 0.0, -0.0), - vel=(0.0, 0.0, 0.0), - dx=0.005, - height_map=height_map, - density=3.0e3, - mu_lambda_ys=(mu, lam, ys), - friction_cohesion=(friction_angle, cohesion, 0.0), - type=type, - jitter=True, - color=(1, 1, 0.5), - random_state=self._episode_rng, - ) - - self.model_builder.init_model_state(self.mpm_model, self.mpm_states) - self.mpm_model.struct.static_ke = 100.0 - self.mpm_model.struct.static_kd = 0.0 - self.mpm_model.struct.static_mu = 1.0 - self.mpm_model.struct.static_ka = 0.0 - - self.mpm_model.struct.body_ke = 100.0 - self.mpm_model.struct.body_kd = 0.0 - self.mpm_model.struct.body_mu = 1.0 - self.mpm_model.struct.body_ka = 0.0 - - self.mpm_model.adaptive_grid = True - - self.mpm_model.grid_contact = True - self.mpm_model.particle_contact = True - self.mpm_model.struct.body_sticky = 1 - self.mpm_model.struct.ground_sticky = 1 - - self.mpm_model.struct.particle_radius = 0.0025 - - def _load_agent(self): - self.agent = PandaBucket( - self._scene, - self._control_freq, - control_mode=self._control_mode, - ) - - self.grasp_site: sapien.Link = get_obj_by_name( - self.agent.robot.get_links(), "bucket" - ) - - def _initialize_agent(self): - qpos = np.array([-0.174, 0.457, 0.203, -1.864, -0.093, 2.025, 1.588]) - qpos[:-1] += self._episode_rng.normal(0, 0.02, len(qpos) - 1) - qpos[-1] += self._episode_rng.normal(0, 0.2, 1) - self.agent.reset(qpos) - self.agent.robot.set_pose(sapien.Pose([-0.56, 0, 0])) - self.home_ee = np.array([0.0, 0.0, self.target_height]) - - def _initialize_actors(self): - super()._initialize_actors() - self.target_height = 0.2 - self.target_num = self._episode_rng.choice(range(250, 1150), 1)[0] - self.mpm_model.struct.n_particles = len(self.model_builder.mpm_particle_q) - - bucket_mesh = actor2meshes(self.grasp_site)[0] - vertices = bucket_mesh.vertices - ones = np.ones(len(vertices)) - self.vertices_mat = np.column_stack((vertices, ones.T)) - - def _load_actors(self): - super()._load_actors() - b = self._scene.create_actor_builder() - b.add_box_collision(half_size=[0.12, 0.02, 0.03]) - b.add_box_visual(half_size=[0.12, 0.02, 0.03]) - w0 = b.build_kinematic("wall") - w1 = b.build_kinematic("wall") - w2 = b.build_kinematic("wall") - w3 = b.build_kinematic("wall") - - w0.set_pose(sapien.Pose([0, -0.1, 0.03])) - w1.set_pose(sapien.Pose([0, 0.1, 0.03])) - w2.set_pose(sapien.Pose([-0.1, 0, 0.03], [0.7071068, 0, 0, 0.7071068])) - w3.set_pose(sapien.Pose([0.1, 0, 0.03], [0.7071068, 0, 0, 0.7071068])) - self.walls = [w0, w1, w2, w3] - - def _get_coupling_actors( - self, - ): - return [ - (l, "visual") for l in self.agent.robot.get_links() if l.name == "bucket" - ] + self.walls - - def _register_sensors(self): - p, q = [-0.2, -0, 0.4], euler2quat(0, np.pi / 6, 0) - return CameraConfig("base_camera", p, q, 128, 128, np.pi / 2, 0.001, 10) - - def _register_render_cameras(self): - p, q = [-0.35, -0, 0.4], euler2quat(0, np.pi / 6, 0) - return CameraConfig("render_camera", p, q, 512, 512, 1, 0.001, 10) - - def _get_obs_extra(self) -> OrderedDict: - return OrderedDict( - tcp_pose=vectorize_pose(self.grasp_site.get_pose()), - target=np.array([self.target_num]), - ) - - def reset(self, seed=None, options=None): - if options is None: - options = dict() - target_num = options.pop("target_num", None) - ret = super().reset(seed=seed, options=options) - if target_num is not None: - self.target_num = int(target_num) - return ret - - def evaluate(self, **kwargs): - particles_x = self.get_mpm_state()["x"] - particles_v = self.get_mpm_state()["v"] - lift_num = len(np.where(particles_x[:, 2] > self.target_height)[0]) - spill_num = self.mpm_model.struct.n_particles - len( - np.where( - (particles_x[:, 0] > -0.12) - & (particles_x[:, 0] < 0.12) - & (particles_x[:, 1] > -0.12) - & (particles_x[:, 1] < 0.12) - )[0] - ) - return dict( - success=( - lift_num > self.target_num - 100 - and lift_num < self.target_num + 150 - and spill_num < 20 - and len(np.where((particles_v < 0.05) & (particles_v > -0.05))[0]) - / (self.mpm_model.struct.n_particles * 3) - > 0.99 - ) - ) - - def _in_bbox_ids(self, particles_x, bbox): - return np.where( - (particles_x[:, 0] >= np.min(bbox[:, 0])) - & (particles_x[:, 0] <= np.max(bbox[:, 0])) - & (particles_x[:, 1] >= np.min(bbox[:, 1])) - & (particles_x[:, 1] <= np.max(bbox[:, 1])) - & (particles_x[:, 2] >= np.min(bbox[:, 2])) - & (particles_x[:, 2] <= np.max(bbox[:, 2])) - )[0] - - def _in_bucket_ids(self, particles_x, bbox, top_signs, bot_signs): - return np.where( - (particles_x[:, 0] >= np.min(bbox[:, 0])) - & (particles_x[:, 0] <= np.max(bbox[:, 0])) - & (particles_x[:, 1] >= np.min(bbox[:, 1])) - & (particles_x[:, 1] <= np.max(bbox[:, 1])) - & (particles_x[:, 2] >= np.min(bbox[:, 2])) - & (particles_x[:, 2] <= np.max(bbox[:, 2])) - & (top_signs > 0) - & (bot_signs > 0) - )[0] - - def _bucket_keypoints(self): - gripper_mat = self.grasp_site.get_pose().to_transformation_matrix() - bucket_base_mat = np.array( - [[1, 0, 0, 0], [0, 1, 0, -0.01], [0, 0, 1, 0.045], [0, 0, 0, 1]] - ) - bucket_tlmat = np.array( - [[1, 0, 0, -0.03], [0, 1, 0, -0.01], [0, 0, 1, 0.01], [0, 0, 0, 1]] - ) - bucket_trmat = np.array( - [[1, 0, 0, 0.03], [0, 1, 0, -0.01], [0, 0, 1, 0.01], [0, 0, 0, 1]] - ) - bucket_blmat = np.array( - [[1, 0, 0, -0.03], [0, 1, 0, 0.02], [0, 0, 1, 0.08], [0, 0, 0, 1]] - ) - bucket_brmat = np.array( - [[1, 0, 0, 0.03], [0, 1, 0, 0.02], [0, 0, 1, 0.08], [0, 0, 0, 1]] - ) - bucket_base_pos = sapien.Pose.from_transformation_matrix( - gripper_mat @ bucket_base_mat - ).p - bucket_tlpos = sapien.Pose.from_transformation_matrix( - gripper_mat @ bucket_tlmat - ).p - bucket_trpos = sapien.Pose.from_transformation_matrix( - gripper_mat @ bucket_trmat - ).p - bucket_blpos = sapien.Pose.from_transformation_matrix( - gripper_mat @ bucket_blmat - ).p - bucket_brpos = sapien.Pose.from_transformation_matrix( - gripper_mat @ bucket_brmat - ).p - return ( - bucket_base_pos, - bucket_tlpos, - bucket_trpos, - bucket_blpos, - bucket_brpos, - gripper_mat, - ) - - def _get_bbox(self, points): - return np.array( - [ - [np.min(points[:, 0]), np.min(points[:, 1]), np.min(points[:, 2])], - [np.max(points[:, 0]), np.max(points[:, 1]), np.max(points[:, 2])], - ] - ) - - def bucket_top_normal(self): - ( - bucket_base_pos, - bucket_tlpos, - bucket_trpos, - bucket_blpos, - bucket_brpos, - gripper_mat, - ) = self._bucket_keypoints() - bucket_top_normal = np.cross( - bucket_base_pos - bucket_trpos, bucket_trpos - bucket_tlpos - ) - bucket_top_normal /= np.linalg.norm(bucket_top_normal) - return bucket_top_normal - - def particles_inside_bucket(self): - # bucket boundary - ( - bucket_base_pos, - bucket_tlpos, - bucket_trpos, - bucket_blpos, - bucket_brpos, - gripper_mat, - ) = self._bucket_keypoints() - - bucket_top_normal = np.cross( - bucket_base_pos - bucket_trpos, bucket_trpos - bucket_tlpos - ) - bucket_top_normal /= np.linalg.norm(bucket_top_normal) - top_d = -np.sum(bucket_top_normal * bucket_base_pos) - top_vec = np.array(list(bucket_top_normal) + [top_d]) - bucket_bot_normal = np.cross( - bucket_base_pos - bucket_blpos, bucket_blpos - bucket_brpos - ) - bucket_bot_normal /= np.linalg.norm(bucket_bot_normal) - bot_d = -np.sum(bucket_bot_normal * bucket_base_pos) - bot_vec = np.array(list(bucket_bot_normal) + [bot_d]) - - vertices = (gripper_mat @ self.vertices_mat.T).T[:, :3] - bbox = self._get_bbox(vertices) - - # pick particles - particles_x = self.get_mpm_state()["x"] - ones = np.ones(len(particles_x)) - particles = np.column_stack((particles_x, ones.T)) - top_signs = particles @ top_vec.T - bot_signs = particles @ bot_vec.T - lifted_particles = particles_x[ - self._in_bucket_ids(particles_x, bbox, top_signs, bot_signs) - ] - return lifted_particles - - def compute_dense_reward(self, reward_info=False, **kwargs): - if self.evaluate()["success"]: - if reward_info: - return {"reward": 6.0} - return 6.0 - particles_x = self.get_mpm_state()["x"] - - stage = 0 - - # spill reward - spill_num = self.n_particles - len( - np.where( - (particles_x[:, 0] > -0.12) - & (particles_x[:, 0] < 0.12) - & (particles_x[:, 1] > -0.12) - & (particles_x[:, 1] < 0.12) - )[0] - ) - spill_reward = -spill_num / 100 - - ( - bucket_base_pos, - bucket_tlpos, - bucket_trpos, - bucket_blpos, - bucket_brpos, - gripper_mat, - ) = self._bucket_keypoints() - - lifted_particles = self.particles_inside_bucket() - lift_num = len(lifted_particles) - lift_reward = ( - min(lift_num / self.target_num, 1) - - max(0, lift_num - self.target_num - 500) * 0.001 - ) - - gripper_pos = self.grasp_site.get_pose().p - height_dist = ( - max(self.target_height + 0.05 - np.mean(lifted_particles[:, 2]), 0) - if len(lifted_particles) > 0 - else 1 - ) - # reaching reward & height reward & flat reward - if height_dist > 0.1 and lift_num > self.target_num + 300: - reaching_reward = 1 - height_reward = 1 - np.tanh(3 * height_dist) - flat_dist = 0.5 * ( - max(bucket_base_pos[2] + 0.01 - bucket_blpos[2], 0) - + max(bucket_blpos[2] - bucket_brpos[2], 0) - ) - flat_reward = 1 - np.tanh(50 * flat_dist) - stage = 1 - elif height_dist <= 0.1: - lift_reward = ( - 1 - + min(lift_num / self.target_num, 1) - - max(0, lift_num - self.target_num - 100) * 0.001 - ) - reaching_reward = 1 - height_reward = 1 - np.tanh(3 * height_dist) - flat_dist = 0.5 * ( - max(bucket_base_pos[2] - 0.01 - bucket_blpos[2], 0) - + max(bucket_blpos[2] - bucket_brpos[2], 0) - ) - flat_reward = 1 - np.tanh(50 * flat_dist) - stage = 2 - else: - if ( - gripper_pos[0] > -0.1 - and gripper_pos[0] < 0.1 - and gripper_pos[1] > -0.1 - and gripper_pos[1] < 0.1 - ): - dist = gripper_pos[2] + max(0.04 - gripper_pos[0], 0) - reaching_reward = 1 - np.tanh(10 * dist) - else: - reaching_reward = 0 - height_reward = 0 - flat_reward = 0 - - reward = ( - reaching_reward * 0.5 - + lift_reward - + height_reward - + spill_reward - + flat_reward - ) - if reward_info: - return { - "reward": reward, - "reaching_reward": reaching_reward, - "lift_reward": lift_reward, - "lift_num": lift_num, - "target_num": self.target_num, - "flat_reward": flat_reward, - "height_reward": height_reward, - "spill_reward": spill_reward, - "stage": stage, - "height_dist": height_dist, - } - return reward - - def compute_normalized_dense_reward(self, **kwargs): - return self.compute_dense_reward(**kwargs) / 6.0 - - def render(self, draw_box=False, draw_target=False): - if draw_target: - bbox = self.target_box - box = self._add_draw_box(bbox) - - img = super().render(draw_box) - if draw_target: - self._remove_draw_box(box) - return img - - def get_state_dict(self) -> np.ndarray: - state = super().get_state_dict() - return np.hstack([state, self.target_num]) - - def set_state_dict(self, state): - self.target_num = state[-1] - super().set_state_dict(state[:-1]) - - -if __name__ == "__main__": - env = ExcavateEnv(reward_mode="dense") - env.reset() - env.agent.set_control_mode("pd_ee_delta_pose") - - a = env.get_state_dict() - env.set_state_dict(a) - - for i in range(100): - env.step(None) - env.render() diff --git a/mani_skill2/envs/ms2/mpm/fill_env.py b/mani_skill2/envs/ms2/mpm/fill_env.py deleted file mode 100644 index 36096dd66..000000000 --- a/mani_skill2/envs/ms2/mpm/fill_env.py +++ /dev/null @@ -1,333 +0,0 @@ -import os -from collections import OrderedDict - -import numpy as np -import sapien -import warp as wp -from transforms3d.euler import euler2quat - -from mani_skill2 import PACKAGE_ASSET_DIR -from mani_skill2.agents.robots.panda.variants import PandaBucket -from mani_skill2.envs.ms2.mpm.base_env import MPMBaseEnv -from mani_skill2.sensors.camera import CameraConfig -from mani_skill2.utils.registration import register_env -from mani_skill2.utils.sapien_utils import get_obj_by_name -from mani_skill2.utils.structs.pose import vectorize_pose - - -@wp.kernel -def success_kernel( - particle_q: wp.array(dtype=wp.vec3), - center: wp.vec2, - radius: float, - height: float, - h1: float, - h2: float, - output: wp.array(dtype=int), -): - tid = wp.tid() - x = particle_q[tid] - a = x[0] - center[0] - b = x[1] - center[1] - z = x[2] - if a * a + b * b < radius * radius and z < height: - if z > h1: - wp.atomic_add(output, 0, 1) - if z > h2: - wp.atomic_add(output, 1, 1) - else: - # spill - if z < 0.005: - wp.atomic_add(output, 2, 1) - - -@register_env("Fill-v0", max_episode_steps=250) -class FillEnv(MPMBaseEnv): - def _setup_mpm(self): - super()._setup_mpm() - self._success_helper = wp.zeros(3, dtype=int, device=self.mpm_model.device) - - def _initialize_mpm(self): - self.model_builder.clear_particles() - - E = 1e4 - nu = 0.3 - mu, lam = E / (2 * (1 + nu)), E * nu / ((1 + nu) * (1 - 2 * nu)) - - # 0 for von-mises, 1 for drucker-prager - type = 1 - - # von-mises - ys = 1e4 - - # drucker-prager - friction_angle = 0.6 - cohesion = 0.05 - - x = 0.03 - y = 0.04 - z = 0.03 - cell_x = 0.004 - self.model_builder.add_mpm_grid( - pos=(-0.2, -0.01, 0.27), - vel=(0.0, 0.0, 0.0), - dim_x=int(x // cell_x), - dim_y=int(y // cell_x), - dim_z=int(z // cell_x), - cell_x=cell_x, - cell_y=cell_x, - cell_z=cell_x, - density=3e3, - mu_lambda_ys=(mu, lam, ys), - friction_cohesion=(friction_angle, cohesion, 0.0), - type=type, - jitter=True, - placement_x="center", - placement_y="center", - placement_z="start", - color=(1, 1, 0.5), - random_state=self._episode_rng, - ) - self.model_builder.init_model_state(self.mpm_model, self.mpm_states) - self.mpm_model.struct.static_ke = 100.0 - self.mpm_model.struct.static_kd = 0.0 - self.mpm_model.struct.static_mu = 1.0 - self.mpm_model.struct.static_ka = 0.0 - - self.mpm_model.struct.body_ke = 100.0 - self.mpm_model.struct.body_kd = 0.0 - self.mpm_model.struct.body_mu = 1.0 - self.mpm_model.struct.body_ka = 0.0 - - self.mpm_model.adaptive_grid = True - - self.mpm_model.grid_contact = True - self.mpm_model.particle_contact = True - self.mpm_model.struct.body_sticky = 1 - self.mpm_model.struct.ground_sticky = 1 - self.mpm_model.struct.particle_radius = 0.0025 - - def _load_agent(self): - self.agent = PandaBucket( - self._scene, - self._control_freq, - control_mode=self._control_mode, - ) - - self.grasp_site: sapien.Link = get_obj_by_name( - self.agent.robot.get_links(), "bucket" - ) - - def _initialize_agent(self): - qpos = np.array([-0.188, 0.234, 0.201, -2.114, -0.088, 1.35, 1.571]) - qpos[-2] += self._episode_rng.normal(0, 0.03, 1) - self.agent.reset(qpos) - self.agent.robot.set_pose(sapien.Pose([-0.6, 0, 0])) - - def _register_sensors(self): - p, q = [-0.4, -0.0, 0.4], euler2quat(0, np.pi / 6, 0) - return CameraConfig("base_camera", p, q, 128, 128, np.pi / 2, 0.001, 10) - - def _register_render_cameras(self): - p, q = [-0.5, -0.4, 0.6], euler2quat(0, np.pi / 6, np.pi / 2 - np.pi / 5) - return CameraConfig("render_camera", p, q, 512, 512, 1, 0.001, 10) - - def _load_actors(self): - super()._load_actors() - beaker_file = os.path.join( - PACKAGE_ASSET_DIR, "deformable_manipulation", "beaker.glb" - ) - - b = self._scene.create_actor_builder() - b.add_visual_from_file(beaker_file, scale=[0.04] * 3) - b.add_collision_from_file(beaker_file, scale=[0.04] * 3, density=300) - self.target_beaker = b.build_kinematic("target_beaker") - - def _initialize_actors(self): - super()._initialize_actors() - random_x = self._episode_rng.rand(1)[0] * 2 - 1 - random_y = self._episode_rng.rand(1)[0] * 2 - 1 - self.beaker_x = -0.16 + random_x * 0.1 - self.beaker_y = random_y * 0.1 - self.target_beaker.set_pose(sapien.Pose([self.beaker_x, self.beaker_y, 0])) - - vs = self.target_beaker.get_visual_bodies() - assert len(vs) == 1 - v = vs[0] - vertices = np.concatenate([s.mesh.vertices for s in v.get_render_shapes()], 0) - self._target_height = vertices[:, 2].max() * v.scale[2] - self._target_radius = v.scale[0] - - def _get_coupling_actors( - self, - ): - return [ - (l, "visual") for l in self.agent.robot.get_links() if l.name == "bucket" - ] + [(self.target_beaker, "visual")] - - def _get_obs_extra(self) -> OrderedDict: - return OrderedDict( - tcp_pose=vectorize_pose(self.grasp_site.get_pose()), - target=np.array([self.beaker_x, self.beaker_y]), - ) - - def evaluate(self, **kwargs): - particles_v = self.get_mpm_state()["v"] - self._success_helper.zero_() - wp.launch( - success_kernel, - dim=self.mpm_model.struct.n_particles, - inputs=[ - self.mpm_states[0].struct.particle_q, - self.target_beaker.pose.p[:2], - self._target_radius, - self._target_height, - 0, - self._target_height, - self._success_helper, - ], - device=self.mpm_model.device, - ) - above_start, above_end, spill = self._success_helper.numpy() - - if ( - above_start / self.mpm_model.struct.n_particles > 0.9 - and len(np.where((particles_v < 0.05) & (particles_v > -0.05))[0]) - / (self.mpm_model.struct.n_particles * 3) - > 0.99 - ): - return dict(success=True) - return dict(success=False) - - def _bucket_keypoints(self): - gripper_mat = self.grasp_site.get_pose().to_transformation_matrix() - bucket_base_mat = np.array( - [[1, 0, 0, 0], [0, 1, 0, -0.01], [0, 0, 1, 0.045], [0, 0, 0, 1]] - ) - bucket_tlmat = np.array( - [[1, 0, 0, -0.03], [0, 1, 0, -0.01], [0, 0, 1, 0.01], [0, 0, 0, 1]] - ) - bucket_trmat = np.array( - [[1, 0, 0, 0.03], [0, 1, 0, -0.01], [0, 0, 1, 0.01], [0, 0, 0, 1]] - ) - bucket_blmat = np.array( - [[1, 0, 0, -0.03], [0, 1, 0, 0.02], [0, 0, 1, 0.08], [0, 0, 0, 1]] - ) - bucket_brmat = np.array( - [[1, 0, 0, 0.03], [0, 1, 0, 0.02], [0, 0, 1, 0.08], [0, 0, 0, 1]] - ) - bucket_base_pos = sapien.Pose.from_transformation_matrix( - gripper_mat @ bucket_base_mat - ).p - bucket_tlpos = sapien.Pose.from_transformation_matrix( - gripper_mat @ bucket_tlmat - ).p - bucket_trpos = sapien.Pose.from_transformation_matrix( - gripper_mat @ bucket_trmat - ).p - bucket_blpos = sapien.Pose.from_transformation_matrix( - gripper_mat @ bucket_blmat - ).p - bucket_brpos = sapien.Pose.from_transformation_matrix( - gripper_mat @ bucket_brmat - ).p - return ( - bucket_base_pos, - bucket_tlpos, - bucket_trpos, - bucket_blpos, - bucket_brpos, - gripper_mat, - ) - - def compute_dense_reward(self, reward_info=False, **kwargs): - if self.evaluate()["success"]: - if reward_info: - return {"reward": 2.5} - return 2.5 - # bucket frame - gripper_mat = self.grasp_site.get_pose().to_transformation_matrix() - gripper_bucket_mat = np.array( - [[1, 0, 0, 0], [0, 1, 0, 0.02], [0, 0, 1, 0.08], [0, 0, 0, 1]] - ) - bucket_mat = gripper_mat @ gripper_bucket_mat - bucket_pos = sapien.Pose.from_transformation_matrix(bucket_mat).p - - beaker_pos = np.array([self.beaker_x, self.beaker_y]) - dist = np.linalg.norm(bucket_pos[:2] - beaker_pos) - reaching_reward = 1 - np.tanh(10.0 * dist) - - self._success_helper.zero_() - wp.launch( - success_kernel, - dim=self.mpm_model.struct.n_particles, - inputs=[ - self.mpm_states[0].struct.particle_q, - self.target_beaker.pose.p[:2], - self._target_radius, - self._target_height, - 0, - self._target_height, - self._success_helper, - ], - device=self.mpm_model.device, - ) - above_start, above_end, spill = self._success_helper.numpy() - inside_reward = above_start / self.n_particles - spill_reward = -spill / 100 - - if reaching_reward > 0.9: - ( - bucket_base_pos, - bucket_tlpos, - bucket_trpos, - bucket_blpos, - bucket_brpos, - gripper_mat, - ) = self._bucket_keypoints() - bucket_bottom = (bucket_blpos + bucket_brpos) / 2 - tilt_reward = 1 - np.tanh(100.0 * (bucket_bottom[2] - bucket_base_pos[2])) - else: - tilt_reward = 0.4 - - if reward_info: - return { - "reaching_reward": reaching_reward, - "inside_reward": inside_reward, - "spill_reward": spill_reward, - "tilt_reward": tilt_reward, - } - return reaching_reward * 0.1 + inside_reward + spill_reward + tilt_reward * 0.5 - - def compute_normalized_dense_reward(self, **kwargs): - return self.compute_dense_reward(**kwargs) / 2.5 - - def render(self, draw_box=False, draw_target=False): - if draw_target: - bbox = self.target_box - box = self._add_draw_box(bbox) - img = super().render(draw_box) - if draw_target: - self._remove_draw_box(box) - return img - - def get_state_dict(self) -> np.ndarray: - state = super().get_state_dict() - return np.hstack([state, self.beaker_x, self.beaker_y]) - - def set_state_dict(self, state): - self.beaker_x = state[-2] - self.beaker_y = state[-1] - super().set_state_dict(state[:-2]) - - -if __name__ == "__main__": - env = FillEnv() - env.reset() - - a = env.get_state_dict() - env.set_state_dict(a) - - for _ in range(100): - env.step(None) - env.render(draw_box=True) diff --git a/mani_skill2/envs/ms2/mpm/hang_env.py b/mani_skill2/envs/ms2/mpm/hang_env.py deleted file mode 100644 index ac9de103b..000000000 --- a/mani_skill2/envs/ms2/mpm/hang_env.py +++ /dev/null @@ -1,394 +0,0 @@ -import os -import pickle -from collections import OrderedDict - -import numpy as np -import sapien -import sapien.physx as physx -import warp as wp -from transforms3d.euler import euler2quat -from transforms3d.quaternions import axangle2quat - -from mani_skill2.agents.robots.panda import Panda -from mani_skill2.envs.ms2.mpm.base_env import MPMBaseEnv, MPMModelBuilder -from mani_skill2.sensors.camera import CameraConfig -from mani_skill2.utils.registration import register_env -from mani_skill2.utils.sapien_utils import get_obj_by_name -from mani_skill2.utils.structs.pose import vectorize_pose -from warp_maniskill.mpm.mpm_simulator import Simulator as MPMSimulator - - -@register_env("Hang-v0", max_episode_steps=350) -class HangEnv(MPMBaseEnv): - def __init__( - self, - *args, - **kwargs, - ): - with open(os.path.join(os.path.dirname(__file__), "RopeInit.pkl"), "rb") as f: - self.rope_init = pickle.load(f) - super().__init__(*args, **kwargs) - - def _setup_mpm(self): - self.model_builder = MPMModelBuilder() - self.model_builder.set_mpm_domain( - domain_size=[0.5, 0.5, 0.5], grid_length=0.015 - ) - self.model_builder.reserve_mpm_particles(count=self.max_particles) - - self._setup_mpm_bodies() - - self.mpm_simulator = MPMSimulator(device="cuda") - self.mpm_model = self.model_builder.finalize(device="cuda") - self.mpm_model.gravity = np.array((0.0, 0.0, -9.81), dtype=np.float32) - self.mpm_model.struct.ground_normal = wp.vec3(0.0, 0.0, 1.0) - self.mpm_model.struct.particle_radius = 0.005 - self.mpm_states = [ - self.mpm_model.state() for _ in range(self._mpm_step_per_sapien_step + 1) - ] - - def _initialize_mpm(self): - self.model_builder.clear_particles() - - E = 1e4 - nu = 0.3 - mu, lam = E / (2 * (1 + nu)), E * nu / ((1 + nu) * (1 - 2 * nu)) - - # 0 for von-mises, 1 for drucker-prager - type = 0 - - # von-mises - ys = 1e4 - - # drucker-prager - friction_angle = 0.6 - cohesion = 0.05 - - x = 0.4 - y = 0.022 - z = 0.022 - cell_x = 0.004 - - self.model_builder.add_mpm_grid( - pos=(0.1, 0.0, 0.05), - vel=(0.0, 0.0, 0.0), - dim_x=int(x // cell_x), - dim_y=int(y // cell_x), - dim_z=int(z // cell_x), - cell_x=cell_x, - cell_y=cell_x, - cell_z=cell_x, - density=300, - mu_lambda_ys=(mu, lam, ys), - friction_cohesion=(friction_angle, cohesion, 0.0), - type=type, - jitter=True, - placement_x="center", - placement_y="center", - placement_z="start", - color=(1, 1, 0.5), - random_state=self._episode_rng, - ) - self.model_builder.init_model_state(self.mpm_model, self.mpm_states) - self.mpm_model.struct.static_ke = 100.0 - self.mpm_model.struct.static_kd = 0.0 - self.mpm_model.struct.static_mu = 1.0 - self.mpm_model.struct.static_ka = 0.0 - - self.mpm_model.struct.body_ke = 100.0 - self.mpm_model.struct.body_kd = 0.0 - self.mpm_model.struct.body_mu = 1.0 - self.mpm_model.struct.body_ka = 0.0 - - self.mpm_model.adaptive_grid = True - - self.mpm_model.struct.body_sticky = 1 - self.mpm_model.struct.ground_sticky = 1 - self.mpm_model.particle_contact = True - self.mpm_model.grid_contact = True - - self.mpm_model.struct.particle_radius = 0.005 - - x = self.get_mpm_state()["x"] - lower = x.min(0) - upper = x.max(0) - dim = np.argmax(upper - lower) - length = upper[dim] - lower[dim] - - # on one side - mask1 = (lower[dim] + length * 0.05 < x[..., dim]) * ( - x[..., dim] < lower[dim] + length * 0.15 - ) - - index1 = np.where(mask1)[0][0] - mask2 = (lower[dim] + length * 0.25 < x[..., dim]) * ( - x[..., dim] < lower[dim] + length * 0.35 - ) - index2 = np.where(mask2)[0][0] - - # on top - mask3 = (lower[dim] + length * 0.48 < x[..., dim]) * ( - x[..., dim] < lower[dim] + length * 0.52 - ) - index3 = np.where(mask3)[0][0] - - # on the other side - mask4 = (lower[dim] + length * 0.65 < x[..., dim]) * ( - x[..., dim] < lower[dim] + length * 0.75 - ) - index4 = np.where(mask4)[0][0] - mask5 = (lower[dim] + length * 0.85 < x[..., dim]) * ( - x[..., dim] < lower[dim] + length * 0.95 - ) - index5 = np.where(mask5)[0][0] - - self.selected_indices = [index1, index2, index3, index4, index5] - - def _load_agent(self): - self.agent = Panda( - self._scene, - self._control_freq, - control_mode=self._control_mode, - ) - self.grasp_site: physx.PhysxArticulationLinkComponent = get_obj_by_name( - self.agent.robot.get_links(), "panda_hand" - ) - - self.lfinger = get_obj_by_name(self.agent.robot.get_links(), "panda_leftfinger") - self.rfinger = get_obj_by_name( - self.agent.robot.get_links(), "panda_rightfinger" - ) - - def _load_actors(self): - super()._load_actors() - builder = self._scene.create_actor_builder() - builder.add_box_collision(half_size=[0.3, 0.01, 0.01]) - builder.add_box_visual(half_size=[0.3, 0.01, 0.01], color=[1, 0, 0, 1]) - self.rod = builder.build_kinematic("rod") - - def _register_sensors(self): - p, q = [0.45, -0.0, 0.5], euler2quat(0, np.pi / 5, np.pi) - return CameraConfig("base_camera", p, q, 128, 128, np.pi / 2, 0.001, 10) - - def _register_render_cameras(self): - p, q = [0.2, 1.0, 0.5], euler2quat(0, 0.2, 4.4) - return CameraConfig("render_camera", p, q, 512, 512, 1, 0.001, 10) - - def initialize_episode(self): - super().initialize_episode() - idx = self._episode_rng.randint(len(self.rope_init["mpm_states"])) - self.agent.set_state( - self.rope_init["agent_states"][idx], ignore_controller=True - ) - self.set_mpm_state(self.rope_init["mpm_states"][idx]) - - def _initialize_actors(self): - # default pose - pose = sapien.Pose([0.1, 0, 0.21], axangle2quat([0, 0, 1], np.pi / 2)) - - # random pose - r = 0.2 + self._episode_rng.rand() * 0.03 - a = np.pi / 4 + self._episode_rng.rand() * np.pi / 2 - x = np.sin(a) * r - y = np.cos(a) * r - z = 0.2 + self._episode_rng.rand() * 0.1 - pose = sapien.Pose([x, y, z], axangle2quat([0, 0, 1], -a)) - self.rod.set_pose(pose) - - def _get_coupling_actors(self): - return self.agent.robot.get_links() + [self.rod] - - def _initialize_agent(self): - qpos = np.array( - [0, np.pi / 16, 0, -np.pi * 5 / 6, 0, np.pi - 0.2, np.pi / 4, 0, 0] - ) - self.agent.reset(qpos) - self.agent.robot.set_pose(sapien.Pose([-0.46, 0, 0])) - - def _get_obs_extra(self) -> OrderedDict: - return OrderedDict( - tcp_pose=vectorize_pose(self.grasp_site.get_pose()), - target=np.hstack([self.rod.get_pose().p, self.rod.get_pose().q]), - ) - - def evaluate(self, **kwargs): - particles_x = self.get_mpm_state()["x"] - particles_v = self.get_mpm_state()["v"] - - lf_pos = self.lfinger.get_pose().p - rf_pos = self.rfinger.get_pose().p - finger_dist = np.linalg.norm(lf_pos - rf_pos) - - pose = self.rod.pose - center = pose.p - normal = pose.to_transformation_matrix()[:3, :3] @ np.array([0, 1, 0]) - - x = particles_x[self.selected_indices] - dirs = x - center - - signs = np.sign(dirs @ normal) - - side = signs[0] == signs[1] and signs[3] == signs[4] and signs[0] != signs[3] - top = ( - np.max(particles_x[:, 2]) > center[2] - and np.max(particles_x[:, 2]) < center[2] + 0.05 - ) - down = dirs[0, 2] < 0 and dirs[4, 2] < 0 - bottom = np.min(particles_x[:, 2]) > 0.03 - - high_ind = np.where(particles_x[:, 2] > center[2] - 0.03) - particles_v = particles_v[high_ind] - return dict( - success=( - side - and top - and down - and bottom - and len(np.where((particles_v < 0.05) & (particles_v > -0.05))[0]) - / (len(particles_v) * 3 + 0.001) - > 0.99 - and finger_dist > 0.07 - ) - ) - - def compute_dense_reward(self, reward_info=False, **kwargs): - gripper_width = ( - self.agent.robot.get_qlimits()[-1, 1] * 2 - ) # NOTE: hard-coded with panda - if self.evaluate()["success"]: - reward = 6 - reaching_reward = 1 - center_reward = 1 - side_reward = 1 - top_reward = 1 - release_reward = 1 - else: - # reaching reward - gripper_pos = self.grasp_site.get_pose().p - particles_x = self.get_mpm_state()["x"] - distance = np.min(np.linalg.norm(particles_x - gripper_pos, axis=-1)) - reaching_reward = 1 - np.tanh(10.0 * distance) - - # rope center reward - center_reward = 0.0 - pose = self.rod.pose - rod_center = pose.p - rope_center = particles_x[self.selected_indices[2]] - - distance = np.linalg.norm(rod_center[:2] - rope_center[:2]) - center_reward += 0.5 * (1 - np.tanh(10.0 * distance)) - - if rod_center[2] >= rope_center[2]: - distance = rod_center[2] - rope_center[2] - center_reward += 0.5 * (1 - np.tanh(10.0 * distance)) - else: - center_reward += 0.5 - - # rope pos reward - top_reward = 0 - bottom_reward = 0 - release_reward = 0 - if rod_center[2] >= rope_center[2]: - side_reward = 0 - else: - bottom_reward = 1 - np.tanh( - 10.0 * max(0, 0.04 - np.min(particles_x[:, 2])) - ) - rod_normal = pose.to_transformation_matrix()[:3, :3] @ np.array( - [0, 1, 0] - ) - x = particles_x[self.selected_indices] - dirs = x - rod_center - signs = np.sign(dirs @ rod_normal) - - side_reward = 0.5 * ( - int(signs[0] != signs[3]) - + int( - signs[0] == signs[1] - and signs[3] == signs[4] - and signs[0] != signs[3] - ) - ) - - if ( - signs[0] == signs[1] - and signs[3] == signs[4] - and signs[0] != signs[3] - ): - top_reward = 0.25 * ( - int(dirs[0, 2] < 0) - + int(dirs[1, 2] < 0) - + int(dirs[3, 2] < 0) - + int(dirs[4, 2] < 0) - ) - reaching_reward = 1 - if top_reward > 0.9: - release_reward = ( - np.sum(self.agent.robot.get_qpos()[-2:]) / gripper_width - ) - reward = ( - reaching_reward - + center_reward - + side_reward - + top_reward - + release_reward - + bottom_reward * 0.2 - ) - - if reward_info: - return { - "reward": reward, - "reaching_reward": reaching_reward, - "center_reward": center_reward, - "side_reward": side_reward, - "top_reward": top_reward, - "release_reward": release_reward, - "bottom_reward": bottom_reward, - "bottom_pos": np.min(particles_x[:, 2]), - } - return reward - - def compute_normalized_dense_reward(self, **kwargs): - return self.compute_dense_reward(**kwargs) / 6.0 - - def get_state_dict(self) -> np.ndarray: - state = super().get_state_dict() - return np.hstack([state, self.rod.get_pose().p, self.rod.get_pose().q]) - - def set_state_dict(self, state): - pose = sapien.Pose(state[-7:-4], state[-4:]) - self.rod.set_pose(pose) - super().set_state_dict(state[:-7]) - - -if __name__ == "__main__": - env = HangEnv(reward_mode="dense", obs_mode="pointcloud") - agent_states = [] - soft_states = [] - - a = env.get_state_dict() - env.set_state_dict(a) - - for X in np.linspace(-0.081, -0.021, 10): - print(X) - env.reset() - env.agent.set_control_mode("pd_ee_delta_pos") - env.step(np.array([X, 0, 1, 1])) - for _ in range(10): - env.step(None) - env.render() - - env.step(np.array([0, 0, 0, -1])) - for _ in range(20): - env.step(None) - env.render() - - for i in range(50): - env.step(np.array([0, 0, -0.2 * i / 100, -1])) - env.render() - agent_states.append(env.agent.get_state()) - print(env.agent.get_state().keys()) - soft_states.append(env.get_mpm_state()) - - # with open("RopeInit.pkl", "wb") as f: - # pickle.dump({"mpm_states": soft_states, "agent_states": agent_states}, f) diff --git a/mani_skill2/envs/ms2/mpm/perlin.py b/mani_skill2/envs/ms2/mpm/perlin.py deleted file mode 100644 index 8ade51f6f..000000000 --- a/mani_skill2/envs/ms2/mpm/perlin.py +++ /dev/null @@ -1,214 +0,0 @@ -"""Perlin noise implementation.""" -# Licensed under ISC -# adapted from https://gist.github.com/eevee/26f547457522755cb1fb8739d0ea89a1 -import math -from itertools import product - -import numpy as np - - -def smoothstep(t): - """Smooth curve with a zero derivative at 0 and 1, making it useful for - interpolating. - """ - return t * t * (3.0 - 2.0 * t) - - -def lerp(t, a, b): - """Linear interpolation between a and b, given a fraction t.""" - return a + t * (b - a) - - -class PerlinNoiseFactory(object): - """Callable that produces Perlin noise for an arbitrary point in an - arbitrary number of dimensions. The underlying grid is aligned with the - integers. - There is no limit to the coordinates used; new gradients are generated on - the fly as necessary. - """ - - def __init__( - self, - dimension, - octaves=1, - tile=(), - unbias=False, - random_state=np.random.RandomState(), - ): - """Create a new Perlin noise factory in the given number of dimensions, - which should be an integer and at least 1. - More octaves create a foggier and more-detailed noise pattern. More - than 4 octaves is rather excessive. - ``tile`` can be used to make a seamlessly tiling pattern. For example: - pnf = PerlinNoiseFactory(2, tile=(0, 3)) - This will produce noise that tiles every 3 units vertically, but never - tiles horizontally. - If ``unbias`` is true, the smoothstep function will be applied to the - output before returning it, to counteract some of Perlin noise's - significant bias towards the center of its output range. - """ - self.dimension = dimension - self.octaves = octaves - self.tile = tile + (0,) * dimension - self.unbias = unbias - - # For n dimensions, the range of Perlin noise is ±sqrt(n)/2; multiply - # by this to scale to ±1 - self.scale_factor = 2 * dimension**-0.5 - - self.gradient = {} - - self.random_state = random_state - - def _generate_gradient(self): - # Generate a random unit vector at each grid point -- this is the - # "gradient" vector, in that the grid tile slopes towards it - - # 1 dimension is special, since the only unit vector is trivial; - # instead, use a slope between -1 and 1 - if self.dimension == 1: - return (self.random_state.uniform(-1, 1),) - - # Generate a random point on the surface of the unit n-hypersphere; - # this is the same as a random unit vector in n dimensions. Thanks - # to: http://mathworld.wolfram.com/SpherePointPicking.html - # Pick n normal random variables with stddev 1 - - random_point = self.random_state.normal(0, 1, size=self.dimension) - # random_point = [random.gauss(0, 1) for _ in range(self.dimension)] - - # Then scale the result to a unit vector - scale = sum(n * n for n in random_point) ** -0.5 - return tuple(coord * scale for coord in random_point) - - def get_plain_noise(self, *point): - """Get plain noise for a single point, without taking into account - either octaves or tiling. - """ - if len(point) != self.dimension: - raise ValueError( - "Expected {} values, got {}".format(self.dimension, len(point)) - ) - - # Build a list of the (min, max) bounds in each dimension - grid_coords = [] - for coord in point: - min_coord = math.floor(coord) - max_coord = min_coord + 1 - grid_coords.append((min_coord, max_coord)) - - # Compute the dot product of each gradient vector and the point's - # distance from the corresponding grid point. This gives you each - # gradient's "influence" on the chosen point. - dots = [] - for grid_point in product(*grid_coords): - if grid_point not in self.gradient: - self.gradient[grid_point] = self._generate_gradient() - gradient = self.gradient[grid_point] - - dot = 0 - for i in range(self.dimension): - dot += gradient[i] * (point[i] - grid_point[i]) - dots.append(dot) - - # Interpolate all those dot products together. The interpolation is - # done with smoothstep to smooth out the slope as you pass from one - # grid cell into the next. - # Due to the way product() works, dot products are ordered such that - # the last dimension alternates: (..., min), (..., max), etc. So we - # can interpolate adjacent pairs to "collapse" that last dimension. Then - # the results will alternate in their second-to-last dimension, and so - # forth, until we only have a single value left. - dim = self.dimension - while len(dots) > 1: - dim -= 1 - s = smoothstep(point[dim] - grid_coords[dim][0]) - - next_dots = [] - while dots: - next_dots.append(lerp(s, dots.pop(0), dots.pop(0))) - - dots = next_dots - - return dots[0] * self.scale_factor - - def __call__(self, *point): - """Get the value of this Perlin noise function at the given point. The - number of values given should match the number of dimensions. - """ - ret = 0 - for o in range(self.octaves): - o2 = 1 << o - new_point = [] - for i, coord in enumerate(point): - coord *= o2 - if self.tile[i]: - coord %= self.tile[i] * o2 - new_point.append(coord) - ret += self.get_plain_noise(*new_point) / o2 - - # Need to scale n back down since adding all those extra octaves has - # probably expanded it beyond ±1 - # 1 octave: ±1 - # 2 octaves: ±1½ - # 3 octaves: ±1¾ - ret /= 2 - 2 ** (1 - self.octaves) - - if self.unbias: - # The output of the plain Perlin noise algorithm has a fairly - # strong bias towards the center due to the central limit theorem - # -- in fact the top and bottom 1/8 virtually never happen. That's - # a quarter of our entire output range! If only we had a function - # in [0..1] that could introduce a bias towards the endpoints... - r = (ret + 1) / 2 - # Doing it this many times is a completely made-up heuristic. - for _ in range(int(self.octaves / 2 + 0.5)): - r = smoothstep(r) - ret = r * 2 - 1 - - return ret - - -def perlin( - freq=1, shape=(128, 128), phase=(0, 0), random_state=np.random.RandomState() -): - h, w = shape - freq_h = freq - freq_w = freq_h / h * w - ph, pw = phase - factory = PerlinNoiseFactory(2, random_state=random_state) - grid = np.stack( - np.meshgrid(np.linspace(0, freq_w, w), np.linspace(0, freq_h, h)), -1 - ) - m = np.zeros(grid.shape[:2]) - - for i in range(len(grid)): - for j in range(len(grid[0])): - x, y = grid[i, j] - m[i, j] = factory(x + ph, y + pw) - return m - - -def added_perlin( - amps=[1], - freqs=[1], - phases=[(0, 0)], - shape=(128, 128), - random_state=np.random.RandomState(), -): - m = np.zeros(shape) - for amp, freq, phase in zip(amps, freqs, phases): - m += amp * perlin(freq, shape, phase, random_state) - return m - - -# import matplotlib.pyplot as plt -# m = added_perlin( -# amps=[1, 0.5, 0.3], -# freqs=[1, 2, 3], -# phases=[(0, 0), (0.25, 0.5), (0.4, 0.6)], -# shape=(128, 256), -# ) -# print(m.min(), m.max()) -# plt.imshow(m) -# plt.show() diff --git a/mani_skill2/envs/ms2/mpm/pinch_env.py b/mani_skill2/envs/ms2/mpm/pinch_env.py deleted file mode 100644 index 61c9a2075..000000000 --- a/mani_skill2/envs/ms2/mpm/pinch_env.py +++ /dev/null @@ -1,254 +0,0 @@ -from collections import OrderedDict - -import h5py -import numpy as np -import sapien -import sapien.physx as physx -import warp as wp -from transforms3d.euler import euler2quat -from warp.distance import compute_chamfer_distance - -from mani_skill2 import ASSET_DIR -from mani_skill2.agents.robots.panda.variants import PandaPinch -from mani_skill2.envs.ms2.mpm.base_env import MPMBaseEnv, MPMModelBuilder, MPMSimulator -from mani_skill2.envs.ms2.mpm.utils import load_h5_as_dict -from mani_skill2.sensors.camera import CameraConfig -from mani_skill2.utils.registration import register_env -from mani_skill2.utils.sapien_utils import get_obj_by_name -from mani_skill2.utils.structs.pose import vectorize_pose - - -@register_env("Pinch-v0", max_episode_steps=300) -class PinchEnv(MPMBaseEnv): - def __init__( - self, - *args, - level_dir="pinch/levels", - **kwargs, - ): - self.level_dir = ASSET_DIR / level_dir - self.all_filepaths = sorted(self.level_dir.glob("*.h5")) - if len(self.all_filepaths) == 0: - raise FileNotFoundError( - "Please download required assets for Pinch:" - "`python -m mani_skill2.utils.download_asset pinch`" - ) - - super().__init__(*args, **kwargs) - - def reset(self, seed=None, options=None): - if options is None: - options = dict() - self.level_file = options.pop("level_file", None) - return super().reset(seed=seed, options=options) - - def _setup_mpm(self): - self.model_builder = MPMModelBuilder() - self.model_builder.set_mpm_domain(domain_size=[0.5, 0.5, 0.5], grid_length=0.01) - self.model_builder.reserve_mpm_particles(count=self.max_particles) - - self._setup_mpm_bodies() - - self.mpm_simulator = MPMSimulator(device="cuda") - self.mpm_model = self.model_builder.finalize(device="cuda") - self.mpm_model.gravity = np.array((0.0, 0.0, -1), dtype=np.float32) - self.mpm_model.struct.ground_normal = wp.vec3(0.0, 0.0, 1.0) - self.mpm_model.struct.particle_radius = 0.005 - self.mpm_states = [ - self.mpm_model.state() for _ in range(self._mpm_step_per_sapien_step + 1) - ] - - def _initialize_mpm(self): - if self.level_file is not None: - filepath = self.level_dir / self.level_file - assert filepath is not None - else: - filepath = self._episode_rng.choice(self.all_filepaths) - - # reset rng after choosing file - self._episode_rng = np.random.RandomState(self._episode_seed) - - self.info = load_h5_as_dict(h5py.File(str(filepath), "r")) - self.goal_depths = self.info["goal_depths"].reshape((4, 128, 128)) - self.goal_rgbs = self.info["goal_rgbs"].reshape((4, 128, 128, 3)) - self.goal_points = self.get_goal_points() - - n = len(self.info["init_state"]["mpm"]["x"]) - self.mpm_model.struct.n_particles = n - self.set_sim_state(self.info["init_state"]) - self.mpm_model.mpm_particle_colors = ( - np.ones((n, 3)) * np.array([0.65237011, 0.14198029, 0.02201299]) - ).astype(np.float32) - - self.goal_array = wp.array( - self.info["end_state"]["mpm"]["x"], dtype=wp.vec3, device="cuda" - ) - self.dist1 = wp.empty(n, dtype=float, device="cuda") - self.dist2 = wp.empty(n, dtype=float, device="cuda") - self.index1 = wp.empty(n, dtype=wp.int64, device="cuda") - self.index2 = wp.empty(n, dtype=wp.int64, device="cuda") - - self.total_deformed_distance = dist1, dist2 = self.info["deformed_distance"] - self.target_dist = (dist1 + dist2) * 0.3 - - self.mpm_model.struct.particle_type.assign(np.zeros(n, dtype=np.int32)) - self.mpm_model.struct.particle_mass.assign(self.info["model"]["particle_mass"]) - self.mpm_model.struct.particle_vol.assign(self.info["model"]["particle_vol"]) - self.mpm_model.struct.particle_mu_lam_ys.assign( - self.info["model"]["particle_mu_lam_ys"] - ) - - self.mpm_model.struct.static_ke = 100.0 - self.mpm_model.struct.static_kd = 0.5 - self.mpm_model.struct.static_mu = 0.9 - self.mpm_model.struct.static_ka = 0.0 - - self.mpm_model.struct.body_ke = 100.0 - self.mpm_model.struct.body_kd = 0.2 - self.mpm_model.struct.body_mu = 0.5 - self.mpm_model.struct.body_ka = 0.0 - - self.mpm_model.adaptive_grid = False - self.mpm_model.particle_contact = True - self.mpm_model.grid_contact = False - self.mpm_model.struct.ground_sticky = True - self.mpm_model.struct.body_sticky = False - - def _get_coupling_actors(self): - return [ - l - for l in self.agent.robot.get_links() - if l.name in ["panda_hand", "panda_leftfinger", "panda_rightfinger"] - ] - - def _load_agent(self): - self.agent = PandaPinch( - self._scene, - self.control_freq, - control_mode=self._control_mode, - ) - self.grasp_site: physx.PhysxArticulationLinkComponent = get_obj_by_name( - self.agent.robot.get_links(), "panda_hand_tcp" - ) - - def _initialize_agent(self): - noise = self._episode_rng.uniform([-0.1] * 7 + [0, 0], [0.1] * 7 + [0, 0]) - qpos = np.array([0, 0.01, 0, -1.96, 0.0, 1.98, 0.0, 0.06, 0.06]) + noise - - self.agent.reset(qpos) - self.agent.robot.set_pose(sapien.Pose([-0.56, 0, 0])) - - def step(self, *args, **kwargs): - self._chamfer_dist = None - return super().step(*args, **kwargs) - - def _register_sensors(self): - p, q = [0.4, 0, 0.3], euler2quat(0, np.pi / 10, -np.pi) - return CameraConfig("base_camera", p, q, 128, 128, np.pi / 2, 0.001, 10) - - def _register_render_cameras(self): - p, q = [-0.05, 0.7, 0.3], euler2quat(0, np.pi / 10, -np.pi / 2) - return CameraConfig("render_camera", p, q, 512, 512, 1, 0.001, 10) - - def compute_dense_reward(self, **kwargs): - # reaching reward - gripper_mat = self.grasp_site.get_pose().to_transformation_matrix() - gripper_bottom_mat = np.array( - [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0.02], [0, 0, 0, 1]] - ) - bottom_mat = gripper_mat @ gripper_bottom_mat - bottom_pos = sapien.Pose.from_transformation_matrix(bottom_mat).p - particles_x = self.get_mpm_state()["x"] - distance = np.min(np.linalg.norm(particles_x - bottom_pos, axis=-1)) - reaching_reward = 1 - np.tanh(10.0 * distance) - - # vertical reward - gripper_mat = self.grasp_site.pose.to_transformation_matrix() - z = gripper_mat[:3, 2] - angle = np.arcsin(np.clip(np.linalg.norm(np.cross(z, [0, 0, -1])), -1, 1)) - reward_orientation = 1 - angle - - # chamfer dist - if self._chamfer_dist is None: - n = self.n_particles - compute_chamfer_distance( - self.mpm_states[0].struct.particle_q, - n, - self.goal_array, - n, - self.dist1, - self.dist2, - self.index1, - self.index2, - ) - wp.synchronize() - chamfer4_1 = (self.dist1.numpy() ** 2).mean() ** 0.25 - chamfer4_2 = (self.dist2.numpy() ** 2).mean() ** 0.25 - - self._chamfer_dist = [chamfer4_1, chamfer4_2] - - return ( - -sum(self._chamfer_dist) * 100.0 - + 0.1 * reaching_reward - + 0.1 * reward_orientation - ) - - def compute_normalized_dense_reward(self, **kwargs): - return self.compute_dense_reward( - **kwargs - ) # original reward scale is low, so we don't normalize for now - - def check_success(self, **kwargs): - self.compute_dense_reward() - return sum(self._chamfer_dist) < self.target_dist - - def _get_obs_extra(self): - target_rgb = self.goal_rgbs - target_depth = self.goal_depths - target_points = self.goal_points - return OrderedDict( - tcp_pose=vectorize_pose(self.grasp_site.get_pose()), - target_rgb=target_rgb, - target_depth=target_depth, - target_points=target_points, - ) - - def evaluate(self, **kwargs): - self.compute_dense_reward() - progress = 1 - sum(self._chamfer_dist) / sum(self.total_deformed_distance) - # progerss: 0 for initial shape, 1 for target shape, measured with Chamfer - return {"success": self.check_success(), "progress": progress} - - def get_goal_points(self): - cam_pos = self.info["goal_cam_pos"] - cam_rot = self.info["goal_cam_rot"] - cam_intrinsic = self.info["goal_cam_intrinsic"] - - world_xyzws = [] - total_size = 0 - for depth, rgb, pos, rot in zip( - self.goal_depths, self.goal_rgbs, cam_pos, cam_rot - ): - H, W = depth.shape[:2] - total_size += H * W - T = sapien.Pose(pos, rot).to_transformation_matrix() - xyz = np.stack( - np.meshgrid(np.arange(0.5, W + 0.5), np.arange(0.5, H + 0.5)) - + [np.ones((H, W))], - -1, - ) - cam_xyz = np.linalg.solve( - cam_intrinsic, (xyz * depth[..., None]).reshape(-1, 3).T - ).T.reshape(H, W, 3) - cam_xyz = cam_xyz * [1, -1, -1] # opencv to opengl - cam_xyz = ( - cam_xyz @ np.array([[0, 0, -1], [-1, 0, 0], [0, 1, 0]]).T - ) # opengl to ros/sapien - - world_xyz = cam_xyz[depth > 0] @ T[:3, :3].T + T[:3, 3] - w = np.ones(world_xyz.shape[0]) - world_xyzw = np.c_[world_xyz, w] - world_xyzws.append(world_xyzw) - - points = np.concatenate(world_xyzws, axis=0) - return np.pad(points, [[0, total_size - len(points)], [0, 0]]) diff --git a/mani_skill2/envs/ms2/mpm/pour_env.py b/mani_skill2/envs/ms2/mpm/pour_env.py deleted file mode 100644 index 5594acc9c..000000000 --- a/mani_skill2/envs/ms2/mpm/pour_env.py +++ /dev/null @@ -1,599 +0,0 @@ -import os -from collections import OrderedDict - -import numpy as np -import sapien -import warp as wp -from transforms3d.euler import euler2quat - -from mani_skill2 import PACKAGE_ASSET_DIR -from mani_skill2.agents.robots.panda.variants import PandaPour -from mani_skill2.envs.ms2.mpm.base_env import MPMBaseEnv, MPMModelBuilder, MPMSimulator -from mani_skill2.sensors.camera import CameraConfig -from mani_skill2.utils.geometry import ( - get_local_aabc_for_actor, - get_local_axis_aligned_bbox_for_link, -) -from mani_skill2.utils.registration import register_env -from mani_skill2.utils.sapien_utils import get_obj_by_name -from mani_skill2.utils.structs.pose import vectorize_pose - - -@wp.kernel -def success_kernel( - particle_q: wp.array(dtype=wp.vec3), - center: wp.vec2, - radius: float, - height: float, - h1: float, - h2: float, - output: wp.array(dtype=int), -): - tid = wp.tid() - x = particle_q[tid] - a = x[0] - center[0] - b = x[1] - center[1] - z = x[2] - if a * a + b * b < radius * radius and z < height: - wp.atomic_add(output, 3, 1) - if z > h1: - wp.atomic_add(output, 0, 1) - if z > h2: - wp.atomic_add(output, 1, 1) - else: - # spill - if z < 0.001: - wp.atomic_add(output, 2, 1) - - -def create_ring(): - segs = 16 - angles = np.linspace(0, 2 * np.pi, segs, endpoint=False) - xs = np.cos(angles) - ys = np.sin(angles) - vs = np.zeros((segs, 3)) - vs[:, 0] = xs - vs[:, 1] = ys - - vs2 = vs.copy() - vs2[:, 2] = 1 - - vertices = np.concatenate([vs, vs2], 0) - indices = [] - for i in range(segs): - a = i - b = (i + 1) % segs - c = b + segs - d = a + segs - indices.append(a) - indices.append(b) - indices.append(c) - indices.append(a) - indices.append(c) - indices.append(d) - - indices.append(a) - indices.append(c) - indices.append(b) - indices.append(a) - indices.append(d) - indices.append(c) - - return vertices, np.array(indices) - - -@register_env("Pour-v0", max_episode_steps=350) -class PourEnv(MPMBaseEnv): - def __init__( - self, - *args, - **kwargs, - ): - self.robot_uid = "panda" - self._ring = None - super().__init__(*args, **kwargs) - - def _load_actors(self): - super()._load_actors() - bottle_file = os.path.join( - PACKAGE_ASSET_DIR, "deformable_manipulation", "bottle.glb" - ) - beaker_file = os.path.join( - PACKAGE_ASSET_DIR, "deformable_manipulation", "beaker.glb" - ) - - b = self._scene.create_actor_builder() - b.add_visual_from_file(bottle_file, scale=[0.025] * 3) - b.add_convex_collision_from_file(bottle_file, scale=[0.025] * 3, density=300) - self.source_container = b.build("bottle") - self.source_aabb = get_local_axis_aligned_bbox_for_link(self.source_container) - - target_radius = 0.04 - b = self._scene.create_actor_builder() - b.add_visual_from_file(beaker_file, scale=[target_radius] * 3) - b.add_convex_collision_from_file( - beaker_file, scale=[target_radius] * 3, density=300 - ) - self.target_beaker = b.build_kinematic("target_beaker") - self.target_aabb = get_local_axis_aligned_bbox_for_link(self.target_beaker) - self.target_aabc = get_local_aabc_for_actor(self.target_beaker) - - def _get_coupling_actors( - self, - ): - return [ - (self.source_container, "visual"), - (self.target_beaker, "visual"), - ] - - def _load_agent(self): - self.agent = PandaPour( - self._scene, - self._control_freq, - control_mode=self._control_mode, - ) - self.grasp_site: sapien.Link = get_obj_by_name( - self.agent.robot.get_links(), "panda_hand_tcp" - ) - self.lfinger = get_obj_by_name(self.agent.robot.get_links(), "panda_leftfinger") - self.rfinger = get_obj_by_name( - self.agent.robot.get_links(), "panda_rightfinger" - ) - - def _setup_mpm(self): - self.model_builder = MPMModelBuilder() - self.model_builder.set_mpm_domain( - domain_size=[0.8, 0.8, 0.8], grid_length=0.005 - ) - self.model_builder.reserve_mpm_particles(count=self.max_particles) - - self._setup_mpm_bodies() - - self.mpm_simulator = MPMSimulator(device="cuda") - self.mpm_model = self.model_builder.finalize(device="cuda") - self.mpm_model.gravity = np.array((0.0, 0.0, -9.81), dtype=np.float32) - self.mpm_model.struct.ground_normal = wp.vec3(0.0, 0.0, 1.0) - self.mpm_model.struct.particle_radius = 0.005 - self.mpm_states = [ - self.mpm_model.state() for _ in range(self._mpm_step_per_sapien_step + 1) - ] - self._success_helper = wp.zeros(4, dtype=int, device=self.mpm_model.device) - - def _initialize_mpm(self): - self.model_builder.clear_particles() - - E = 3e5 - nu = 0.1 - mu, lam = E / (2 * (1 + nu)), E * nu / ((1 + nu) * (1 - 2 * nu)) - ys = 0.0 - type = 2 - - viscosity = self._episode_rng.uniform(0.0, 3.0) - density = self._episode_rng.uniform(0.8e3, 2.0e3) - - count = self.model_builder.add_mpm_cylinder( - pos=(*self._source_pos.p[:2], 0.01), - vel=(0.0, 0.0, 0.0), - radius=0.024, - height=self._episode_rng.uniform(0.07, 0.09), - dx=0.0025, - density=density, - mu_lambda_ys=(mu, lam, ys), - friction_cohesion=(0.0, 0.0, viscosity), - type=type, - jitter=False, - random_state=self._episode_rng, - color=[0.0, 0.5, 0.8], - ) - - for i in range(len(self.model_builder.mpm_particle_volume)): - self.model_builder.mpm_particle_volume[i] *= 1.2 - - self.model_builder.init_model_state(self.mpm_model, self.mpm_states) - self.mpm_model.struct.static_ke = 100.0 - self.mpm_model.struct.static_kd = 0.0 - self.mpm_model.struct.static_mu = 1.0 - self.mpm_model.struct.static_ka = 0.0 - - self.mpm_model.struct.body_ke = 100.0 - self.mpm_model.struct.body_kd = 0.0 - self.mpm_model.struct.body_mu = 1.0 - self.mpm_model.struct.body_ka = 0.0 - - self.mpm_model.adaptive_grid = False - - self.mpm_model.struct.body_sticky = 0 - self.mpm_model.struct.ground_sticky = 1 - self.mpm_model.particle_contact = True - self.mpm_model.grid_contact = False - self.mpm_model.struct.particle_radius = 0.0025 - - def _initialize_actors(self): - super()._initialize_actors() - - self._determine_target_pos() - - self.target_beaker.set_pose(self._target_pos) - self.source_container.set_pose(self._source_pos) - - self.source_container.set_velocity([0, 0, 0]) - self.source_container.set_angular_velocity([0, 0, 0]) - - vs = self.target_beaker.get_visual_bodies() - assert len(vs) == 1 - v = vs[0] - vertices = np.concatenate([s.mesh.vertices for s in v.get_render_shapes()], 0) - self._target_height = vertices[:, 2].max() * v.scale[2] - self._target_radius = v.scale[0] - - def _initialize_agent(self): - self.agent.reset(self._init_qpos) - self.agent.robot.set_pose(sapien.Pose([-0.55, 0, 0])) - - def _register_sensors(self): - p, q = [0.4, 0, 0.3], euler2quat(0, np.pi / 10, -np.pi) - return CameraConfig("base_camera", p, q, 128, 128, np.pi / 2, 0.001, 10) - - def _register_render_cameras(self): - p, q = [-0.05, 0.7, 0.3], euler2quat(0, np.pi / 10, -np.pi / 2) - return CameraConfig("render_camera", p, q, 512, 512, 1, 0.001, 10) - - def initialize_episode(self): - super().initialize_episode() - - self.h1 = self._episode_rng.uniform(0.01, 0.02) - self.h2 = self.h1 + 0.004 - - if self._ring is not None: - self._scene.remove_actor(self._ring) - - vertices, indices = create_ring() - b = self._scene.create_actor_builder() - mesh = self._renderer.create_mesh( - vertices.reshape((-1, 3)), indices.reshape((-1, 3)) - ) - mat = self._renderer.create_material() - mat.set_base_color([1, 0, 0, 1]) - b.add_visual_from_mesh( - mesh, - scale=[ - self._target_radius * 1.02, - self._target_radius * 1.02, - self.h2 - self.h1, - ], - material=mat, - ) - ring = b.build_kinematic("ring") - ring.set_pose(sapien.Pose([*self.target_beaker.pose.p[:2], self.h1])) - self._ring = ring - - def _clear(self): - if self._ring is not None: - self._scene.remove_actor(self._ring) - self._ring = None - super()._clear() - - def _setup_viewer(self): - super()._setup_viewer() - self._viewer.set_camera_xyz(-0.05, 0.3, 0.3) - self._viewer.set_camera_rpy(0.0, -0.7, 1.57) - - def _determine_target_pos(self): - pmodel = self.agent.robot.create_pinocchio_model() - hand = next( - l for l in self.agent.robot.get_links() if l.name == "panda_hand_tcp" - ) - while True: - r = self._episode_rng.uniform(0.2, 0.25) - t = self._episode_rng.uniform(0, np.pi) - self._target_pos = sapien.Pose([r * np.cos(t), r * np.sin(t), 0.0]) - - r = self._episode_rng.uniform(0.05, 0.1) - t = self._episode_rng.uniform(np.pi, np.pi * 2) - self._source_pos = sapien.Pose([r * np.cos(t), r * np.sin(t), 0.0]) - - from transforms3d.quaternions import axangle2quat, qmult - - q = qmult( - axangle2quat( - [0, 0, 1], self._episode_rng.uniform(-np.pi / 8, np.pi / 8) - ), - [0.5, -0.5, -0.5, -0.5], - ) - - result, success, error = pmodel.compute_inverse_kinematics( - hand.get_index(), - sapien.Pose( - [ - self._source_pos.p[0] + 0.55, - self._source_pos.p[1], - self._episode_rng.uniform(0.04, 0.06), - ], - q, - ), - [-0.555, 0.646, 0.181, -1.892, 1.171, 1.423, -1.75, 0.04, 0.04], - active_qmask=[1] * 7 + [0] * 2, - ) - if not success: - continue - - result[-2:] = 0.04 - self._init_qpos = result - return - - def _get_obs_extra(self) -> OrderedDict: - return OrderedDict( - tcp_pose=vectorize_pose(self.grasp_site.get_pose()), - target=np.array([self.h1]), - ) - - def in_beaker_num(self): - self._success_helper.zero_() - wp.launch( - success_kernel, - dim=self.mpm_model.struct.n_particles, - inputs=[ - self.mpm_states[0].struct.particle_q, - self.target_beaker.pose.p[:2], - self._target_radius, - self._target_height, - self.h1, - self.h2, - self._success_helper, - ], - device=self.mpm_model.device, - ) - above_start, above_end, spill, in_beaker = self._success_helper.numpy() - return above_start, above_end - - def evaluate(self, **kwargs): - self.get_mpm_state()["v"] - self._success_helper.zero_() - wp.launch( - success_kernel, - dim=self.mpm_model.struct.n_particles, - inputs=[ - self.mpm_states[0].struct.particle_q, - self.target_beaker.pose.p[:2], - self._target_radius, - self._target_height, - self.h1, - self.h2, - self._success_helper, - ], - device=self.mpm_model.device, - ) - above_start, above_end, spill, in_beaker = self._success_helper.numpy() - - upright = ( - self.source_container.pose.to_transformation_matrix()[:3, 2] - @ np.array([0, 0, 1]) - ) >= 0.866 # 30 degrees - - if ( - above_start > 100 - and above_end < 10 - and spill < 100 - and upright - and np.max(self.agent.robot.get_qvel()) < 0.05 - and np.min(self.agent.robot.get_qvel()) > -0.05 - ): - return dict(success=True) - return dict(success=False) - - def compute_dense_reward(self, reward_info=False, **kwargs): - if self.evaluate()["success"]: - if reward_info: - return {"reward": 15} - return 15 - self._success_helper.zero_() - wp.launch( - success_kernel, - dim=self.mpm_model.struct.n_particles, - inputs=[ - self.mpm_states[0].struct.particle_q, - self.target_beaker.pose.p[:2], - self._target_radius, - self._target_height, - self.h1, - self.h2, - self._success_helper, - ], - device=self.mpm_model.device, - ) - above_start, above_end, spill, in_beaker = self._success_helper.numpy() - - source_mat = self.source_container.pose.to_transformation_matrix() - target_mat = self.target_beaker.pose.to_transformation_matrix() - - a, b = self.source_aabb - t = np.array([0.5, 0.5, 0.33]) - grasp_site_target = a * (1 - t) + b * t - grasp_site_target = source_mat[:3, :3] @ grasp_site_target + source_mat[:3, 3] - grasp_site_dist = np.linalg.norm(self.grasp_site.pose.p - grasp_site_target) - reward_grasp_site = -grasp_site_dist - - if grasp_site_dist < 0.05: - reward_grasp_site = 0 - is_grasping = self.agent.is_grasping(self.source_container) - reward_grasp = float(is_grasping) - else: - is_grasping = False - reward_grasp = 0 - - top_center = np.zeros(3) - top_center[:2] = (a[:2] + b[:2]) * 0.5 - top_center[2] = b[2] - top_center = source_mat[:3, :3] @ top_center + source_mat[:3, 3] - - bottom_center = np.zeros(3) - bottom_center[:2] = (a[:2] + b[:2]) * 0.5 - bottom_center[2] = a[2] - bottom_center = source_mat[:3, :3] @ bottom_center + source_mat[:3, 3] - - tx, ty, tr, tzmin, tzmax = self.target_aabc - target_top_center = ( - target_mat[:3, :3] @ np.array([tx, ty, tzmax]) + target_mat[:3, 3] - ) - - dist_lf = np.linalg.norm(self.lfinger.pose.p[:2] - target_top_center[:2]) - dist_rf = np.linalg.norm(self.rfinger.pose.p[:2] - target_top_center[:2]) - reward_finger = 10 * (dist_rf - dist_lf) - - hdist = np.linalg.norm(target_top_center[:2] - self.grasp_site.pose.p[:2]) - vdist = target_top_center[2] - self.grasp_site.pose.p[2] - - reward_in_beaker = 0 - if above_start < 100 or (above_start > 100 and above_start - above_end < 100): - done = False - reward_in_beaker = in_beaker - else: - done = True - reward_in_beaker = in_beaker - max(0, above_start - 500) * 2 - - reward_spill = -spill - - stage = 0 - z = source_mat[:3, 2] - edist = 0 - if not is_grasping: - # not grasping the bottle - reward_orientation = 0 - reward_dist = 0 - stage = 0 - elif done: - # finish pouring - angle = np.arcsin(np.clip(np.linalg.norm(np.cross(z, [0, 0, 1])), -1, 1)) - reward_orientation = 3.5 - angle - reward_dist = 5.5 - np.tanh(10.0 * (bottom_center[2] - top_center[2])) - reward_finger = 1 - stage = 3 - elif vdist > -0.06: - # looking for the right range - angle = np.arcsin(np.clip(np.linalg.norm(np.cross(z, [0, 0, 1])), -1, 1)) - reward_orientation = 1 - angle - reward_dist = 1 - np.tanh(10.0 * (max(0, vdist + 0.06))) - stage = 1 - else: - # within the right range to pour - - top_dist = np.linalg.norm(top_center[:2] - target_top_center[:2]) - bottom_dist = np.linalg.norm(bottom_center[:2] - target_top_center[:2]) - - reward_dist = 2 - np.tanh( - 10.0 * np.linalg.norm(top_center - target_top_center) - ) - if dist_rf - dist_lf > 0: - reward_finger = max( - 10 * (dist_rf - dist_lf), 10 * (bottom_dist - top_dist) - ) - - reward_orientation = 2 - np.tanh( - 10.0 * max(0, top_center[2] - bottom_center[2]) - ) - stage = 2 - - if reward_info: - return { - "reward": reward_grasp_site - + reward_grasp - + reward_in_beaker * 0.001 - + reward_spill * 0.01 - + reward_orientation - + reward_dist - + reward_finger, - "reward_grasp_site": reward_grasp_site, - "reward_grasp": reward_grasp, - "reward_in_beaker": reward_in_beaker, - "reward_spill": reward_spill, - "reward_orientation": reward_orientation, - "reward_dist": reward_dist, - "stage": stage, - "hdist": hdist, - "vdist": vdist, - "tilt": top_center[2] - bottom_center[2], - "edist": edist, - "above_start": above_start, - "in_beaker": in_beaker, - "above_end": above_end, - "tr": tr, - "reward_finger": reward_finger, - } - return ( - reward_grasp_site - + reward_grasp - + reward_in_beaker * 0.001 - + reward_spill * 0.01 - + reward_orientation - + reward_dist - + reward_finger - ) - - def compute_normalized_dense_reward(self, **kwargs): - return self.compute_dense_reward(**kwargs) / 15.0 - - def get_mpm_state(self): - n = self.mpm_model.struct.n_particles - - return OrderedDict( - x=self.copy_array_to_numpy(self.mpm_states[0].struct.particle_q, n), - v=self.copy_array_to_numpy(self.mpm_states[0].struct.particle_qd, n), - F=self.copy_array_to_numpy(self.mpm_states[0].struct.particle_F, n), - C=self.copy_array_to_numpy(self.mpm_states[0].struct.particle_C, n), - vol=self.copy_array_to_numpy(self.mpm_states[0].struct.particle_vol, n), - ) - - def set_mpm_state(self, state): - self.mpm_states[0].struct.particle_q.assign(state["x"]) - self.mpm_states[0].struct.particle_qd.assign(state["v"]) - self.mpm_states[0].struct.particle_F.assign(state["F"]) - self.mpm_states[0].struct.particle_C.assign(state["C"]) - self.mpm_states[0].struct.particle_vol.assign(state["vol"]) - - def get_state_dict(self) -> np.ndarray: - state = super().get_state_dict() - return np.hstack( - [ - state, - self.h1, - self.target_beaker.pose.p, - self.target_beaker.pose.q, - self.source_container.pose.p, - self.source_container.pose.q, - ] - ) - - def set_state_dict(self, state): - source_pos = sapien.Pose(state[-7:-4], state[-4:]) - target_pos = sapien.Pose(state[-14:-11], state[-11:-7]) - self.source_container.set_pose(source_pos) - self.target_beaker.set_pose(target_pos) - self.h1 = state[-15] - - state = state[:-15] - sim_state = OrderedDict() - mpm_state = OrderedDict() - n = self.mpm_model.struct.n_particles - - sim_state["sapien"] = state[: -n * 25] - mpm_state["x"] = state[-n * 25 : -n * 22].reshape((n, 3)) - mpm_state["v"] = state[-n * 22 : -n * 19].reshape((n, 3)) - mpm_state["F"] = state[-n * 19 : -n * 10].reshape((n, 3, 3)) - mpm_state["C"] = state[-n * 10 : -n].reshape((n, 3, 3)) - mpm_state["vol"] = state[-n:].reshape((n,)) - sim_state["mpm"] = mpm_state - - return self.set_sim_state(sim_state) - - -if __name__ == "__main__": - env = PourEnv(mpm_freq=2000) - env.reset() - env.agent.set_control_mode("pd_ee_delta_pose") - - a = env.get_state_dict() - env.set_state_dict(a) - - for i in range(100): - env.step(None) - env.render() diff --git a/mani_skill2/envs/ms2/mpm/shader/common/lights.glsl b/mani_skill2/envs/ms2/mpm/shader/common/lights.glsl deleted file mode 100644 index 504c4db1f..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/common/lights.glsl +++ /dev/null @@ -1,128 +0,0 @@ -struct PointLight { - vec4 position; - vec4 emission; -}; - -struct DirectionalLight { - vec4 direction; - vec4 emission; -}; - -struct SpotLight { - vec4 position; - vec4 direction; - vec4 emission; -}; - - -float diffuse(float NoL) { - return NoL / 3.141592653589793f; -} - -vec3 ggx(float NoL, float NoV, float NoH, float VoH, float roughness, vec3 fresnel) { - float alpha = roughness * roughness; - float alpha2 = alpha * alpha; - - float k = (alpha + 2 * roughness + 1.0) / 8.0; - - float FMi = ((-5.55473) * VoH - 6.98316) * VoH; - vec3 frac = (fresnel + (1 - fresnel) * pow(2.0, FMi)) * alpha2; - float nom0 = NoH * NoH * (alpha2 - 1) + 1; - float nom1 = NoV * (1 - k) + k; - float nom2 = NoL * (1 - k) + k; - float nom = clamp((4 * 3.141592653589793f * nom0 * nom0 * nom1 * nom2), 1e-6, 4 * 3.141592653589793f); - vec3 spec = frac / nom; - - return spec * NoL; -} - -vec3 computeDirectionalLight(vec3 direction, vec3 emission, vec3 normal, vec3 camDir, vec3 diffuseAlbedo, float roughness, vec3 fresnel) { - vec3 lightDir = -direction; - - vec3 H = lightDir + camDir; - float H2 = dot(H, H); - H = H2 < 1e-6 ? vec3(0) : normalize(H); - float NoH = clamp(dot(normal, H), 1e-6, 1); - float VoH = clamp(dot(camDir, H), 1e-6, 1); - float NoL = clamp(dot(normal, lightDir), 0, 1); - float NoV = clamp(dot(normal, camDir), 1e-6, 1); - - vec3 color = diffuseAlbedo * emission * diffuse(NoL); - color += emission * ggx(NoL, NoV, NoH, VoH, roughness, fresnel); - return color; -} - -vec3 computePointLight(vec3 emission, vec3 l, vec3 normal, vec3 camDir, vec3 diffuseAlbedo, float roughness, vec3 fresnel) { - float d = max(length(l), 0.0001); - - if (length(l) == 0) { - return vec3(0.f); - } - - vec3 lightDir = normalize(l); - - vec3 H = lightDir + camDir; - float H2 = dot(H, H); - H = H2 < 1e-6 ? vec3(0) : normalize(H); - float NoH = clamp(dot(normal, H), 1e-6, 1); - float VoH = clamp(dot(camDir, H), 1e-6, 1); - float NoL = clamp(dot(normal, lightDir), 0, 1); - float NoV = clamp(dot(normal, camDir), 1e-6, 1); - - vec3 color = diffuseAlbedo * emission * diffuse(NoL) / d / d; - color += emission * ggx(NoL, NoV, NoH, VoH, roughness, fresnel) / d / d; - return color; -} - -vec3 computeSpotLight(float fov1, float fov2, vec3 centerDir, vec3 emission, vec3 l, vec3 normal, vec3 camDir, vec3 diffuseAlbedo, float roughness, vec3 fresnel) { - float d = max(length(l), 0.0001); - - if (length(l) == 0) { - return vec3(0.f); - } - - vec3 lightDir = normalize(l); - - float cf1 = cos(fov1/2) + 1e-6; - float cf2 = cos(fov2/2); - - float visibility = clamp((dot(-lightDir, centerDir) - cf2) / (cf1 - cf2), 0, 1); - - vec3 H = lightDir + camDir; - float H2 = dot(H, H); - H = H2 < 1e-6 ? vec3(0) : normalize(H); - float NoH = clamp(dot(normal, H), 1e-6, 1); - float VoH = clamp(dot(camDir, H), 1e-6, 1); - float NoL = clamp(dot(normal, lightDir), 0, 1); - float NoV = clamp(dot(normal, camDir), 1e-6, 1); - - vec3 color = diffuseAlbedo * emission * diffuse(NoL) / d / d; - color += emission * ggx(NoL, NoV, NoH, VoH, roughness, fresnel) / d / d; - return visibility * color; -} - -vec3 computeSpotLight2(float fov, vec3 centerDir, vec3 emission, vec3 l, vec3 normal, vec3 camDir, vec3 diffuseAlbedo, float roughness, vec3 fresnel) { - float d = max(length(l), 0.0001); - - if (length(l) == 0) { - return vec3(0.f); - } - - vec3 lightDir = normalize(l); - - if (dot(-lightDir, centerDir) <= 0) { - return vec3(0.f); - } - - vec3 H = lightDir + camDir; - float H2 = dot(H, H); - H = H2 < 1e-6 ? vec3(0) : normalize(H); - float NoH = clamp(dot(normal, H), 1e-6, 1); - float VoH = clamp(dot(camDir, H), 1e-6, 1); - float NoL = clamp(dot(normal, lightDir), 0, 1); - float NoV = clamp(dot(normal, camDir), 1e-6, 1); - - vec3 color = diffuseAlbedo * emission * diffuse(NoL) / d / d; - color += emission * ggx(NoL, NoV, NoH, VoH, roughness, fresnel) / d / d; - return color; -} diff --git a/mani_skill2/envs/ms2/mpm/shader/common/shadow.glsl b/mani_skill2/envs/ms2/mpm/shader/common/shadow.glsl deleted file mode 100644 index cabf3a8c5..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/common/shadow.glsl +++ /dev/null @@ -1,31 +0,0 @@ -vec3 project(mat4 proj, vec3 point) { - vec4 v = proj * vec4(point, 1); - return v.xyz / v.w; -} - -const int PCF_SampleCount = 25; -vec2 PCF_Samples[PCF_SampleCount] = { - {-2, -2}, {-1, -2}, {0, -2}, {1, -2}, {2, -2}, - {-2, -1}, {-1, -1}, {0, -1}, {1, -1}, {2, -1}, - {-2, 0}, {-1, 0}, {0, 0}, {1, 0}, {2, 0}, - {-2, 1}, {-1, 1}, {0, 1}, {1, 1}, {2, 1}, - {-2, 2}, {-1, 2}, {0, 2}, {1, 2}, {2, 2} -}; - -float ShadowMapPCF( - sampler2D shadowTex, vec3 projCoord, float resolution, float searchUV, float filterSize) -{ - float shadow = 0.0f; - vec2 grad = fract(projCoord.xy * resolution + 0.5f); - - for (int i = 0; i < PCF_SampleCount; i++) - { - vec4 tmp = textureGather(shadowTex, projCoord.xy + filterSize * PCF_Samples[i] * searchUV); - tmp.x = tmp.x < projCoord.z ? 0.0f : 1.0f; - tmp.y = tmp.y < projCoord.z ? 0.0f : 1.0f; - tmp.z = tmp.z < projCoord.z ? 0.0f : 1.0f; - tmp.w = tmp.w < projCoord.z ? 0.0f : 1.0f; - shadow += mix(mix(tmp.w, tmp.z, grad.x), mix(tmp.x, tmp.y, grad.x), grad.y); - } - return shadow / PCF_SampleCount; -} diff --git a/mani_skill2/envs/ms2/mpm/shader/common/view.glsl b/mani_skill2/envs/ms2/mpm/shader/common/view.glsl deleted file mode 100644 index e05b6a36b..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/common/view.glsl +++ /dev/null @@ -1,7 +0,0 @@ -float ndc_depth_to_z(float z, float proj22, float proj32) { - return -proj32 / (z + proj22); -} - -float z_to_ndc_depth(float z, float proj22, float proj32) { - return -proj22 - proj32 / z; -} diff --git a/mani_skill2/envs/ms2/mpm/shader/point/composite.vert b/mani_skill2/envs/ms2/mpm/shader/point/composite.vert deleted file mode 100644 index 520c95688..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/point/composite.vert +++ /dev/null @@ -1,9 +0,0 @@ -#version 450 - -layout (location = 0) out vec2 outUV; - -void main() -{ - outUV = vec2((gl_VertexIndex << 1) & 2, gl_VertexIndex & 2); - gl_Position = vec4(outUV * 2.0f + -1.0f, 0.0f, 1.0f); -} diff --git a/mani_skill2/envs/ms2/mpm/shader/point/composite0.frag b/mani_skill2/envs/ms2/mpm/shader/point/composite0.frag deleted file mode 100644 index d95434a8c..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/point/composite0.frag +++ /dev/null @@ -1,83 +0,0 @@ -#version 450 - -layout(set = 0, binding = 0) uniform sampler2D samplerPointDepthLinear; - -layout(location = 0) in vec2 inUV; -layout(location = 0) out vec4 outSmoothedDepthLinear; - -layout(set = 1, binding = 0) uniform CameraBuffer { - mat4 viewMatrix; - mat4 projectionMatrix; - mat4 viewMatrixInverse; - mat4 projectionMatrixInverse; - mat4 prevViewMatrix; - mat4 prevViewMatrixInverse; - float width; - float height; -} cameraBuffer; - -float sqr(float x) { - return x * x; -} - -// void main() { -// outSmoothedDepthLinear.x = texture(samplerPointDepthLinear, inUV).x; -// } - -void main() -{ - vec2 res = vec2(cameraBuffer.width, cameraBuffer.height); - vec2 coord = vec2(gl_FragCoord.xy); - float depth = texture(samplerPointDepthLinear, coord / res).x; - - if (depth == 0.0) - { - outSmoothedDepthLinear.x = 0.0; - return; - } - - float blurDepthFalloff = 5.5; - // float radius = 10.0; - float radius = min(3.0 / (-depth), 16); - float radiusInv = 1.0/radius; - float taps = ceil(radius); - float frac = taps - radius; - - float sum = 0.0; - float wsum = 0.0; - float count = 0.0; - - for(float y=-taps; y <= taps; y += 1) { - for(float x=-taps; x <= taps; x += 1) { - vec2 offset = vec2(x, y); - float samp = texture(samplerPointDepthLinear, (coord + offset) / res).x; - - if (samp == 0.0) { - continue; - } - - // spatial domain - float r1 = length(vec2(x, y))*radiusInv; - float w = exp(-(r1*r1)); - - // range domain (based on depth difference) - float r2 = (samp - depth) * blurDepthFalloff; - float g = exp(-(r2*r2)); - - //fractional radius contributions - float wBoundary = step(radius, max(abs(x), abs(y))); - float wFrac = 1.0 - wBoundary*frac; - - sum += samp * w * g * wFrac; - wsum += w * g * wFrac; - count += g * wFrac; - } - } - - if (wsum > 0.0) { - sum /= wsum; - } - - float blend = count/sqr(2.0*radius+1.0); - outSmoothedDepthLinear.x = mix(depth, sum, blend); -} diff --git a/mani_skill2/envs/ms2/mpm/shader/point/composite1.frag b/mani_skill2/envs/ms2/mpm/shader/point/composite1.frag deleted file mode 100644 index 68323cd4c..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/point/composite1.frag +++ /dev/null @@ -1,273 +0,0 @@ -#version 450 - -layout(location = 0) in vec2 inUV; -layout(location = 0) out vec4 outPoint; - -layout(set = 0, binding = 0) uniform sampler2D samplerSmoothedDepthLinear; -layout(set = 0, binding = 1) uniform sampler2D samplerPointColor; - -// IBL -layout(set = 0, binding = 2) uniform samplerCube samplerEnvironment; -layout(set = 0, binding = 3) uniform sampler2D samplerBRDFLUT; - -// layout (constant_id = 0) const int NUM_DIRECTIONAL_LIGHTS = 3; -// layout (constant_id = 1) const int NUM_POINT_LIGHTS = 10; -// layout (constant_id = 2) const int NUM_SPOT_LIGHTS = 10; - -layout (constant_id = 0) const int NUM_DIRECTIONAL_LIGHTS = 3; -layout (constant_id = 1) const int NUM_POINT_LIGHTS = 10; -layout (constant_id = 2) const int NUM_DIRECTIONAL_LIGHT_SHADOWS = 1; -layout (constant_id = 3) const int NUM_POINT_LIGHT_SHADOWS = 3; -layout (constant_id = 4) const int NUM_TEXTURED_LIGHT_SHADOWS = 1; -layout (constant_id = 5) const int NUM_SPOT_LIGHT_SHADOWS = 10; -layout (constant_id = 6) const int NUM_SPOT_LIGHTS = 10; - - -layout(set = 1, binding = 0) uniform CameraBuffer { - mat4 viewMatrix; - mat4 projectionMatrix; - mat4 viewMatrixInverse; - mat4 projectionMatrixInverse; - mat4 prevViewMatrix; - mat4 prevViewMatrixInverse; - float width; - float height; -} cameraBuffer; - -#include "../common/lights.glsl" -#include "../common/shadow.glsl" -layout(set = 2, binding = 0) uniform SceneBuffer { - vec4 ambientLight; - DirectionalLight directionalLights[3]; - SpotLight spotLights[10]; - PointLight pointLights[10]; - SpotLight texturedLights[1]; -} sceneBuffer; - -struct LightBuffer { - mat4 viewMatrix; - mat4 viewMatrixInverse; - mat4 projectionMatrix; - mat4 projectionMatrixInverse; - int width; - int height; -}; - -layout(set = 2, binding = 1) uniform ShadowBuffer { - LightBuffer directionalLightBuffers[3]; - LightBuffer spotLightBuffers[10]; - LightBuffer pointLightBuffers[60]; - LightBuffer texturedLightBuffers[1]; -} shadowBuffer; - -layout(set = 2, binding = 2) uniform samplerCube samplerPointLightDepths[3]; -layout(set = 2, binding = 3) uniform sampler2D samplerDirectionalLightDepths[1]; -layout(set = 2, binding = 4) uniform sampler2D samplerTexturedLightDepths[1]; -layout(set = 2, binding = 5) uniform sampler2D samplerSpotLightDepths[10]; -layout(set = 2, binding = 6) uniform sampler2D samplerTexturedLightTextures[1]; - -vec3 getCameraSpacePosition(vec2 uv, float z) { - float depth = -cameraBuffer.projectionMatrix[2][2] - cameraBuffer.projectionMatrix[3][2] / z; - vec4 ndc = vec4(uv * 2.0 - 1.0, depth, 1.0); - vec4 csPosition = cameraBuffer.projectionMatrixInverse * ndc; - return vec3(csPosition.xy / csPosition.w, z); -} - -vec4 world2camera(vec4 pos) { - return cameraBuffer.viewMatrix * pos; -} - - -vec3 diffuseIBL(vec3 albedo, vec3 N) { - N = vec3(-N.y, N.z, -N.x); - vec3 color = textureLod(samplerEnvironment, N, 5).rgb; - return color * albedo; -} - -vec3 specularIBL(vec3 fresnel, float roughness, vec3 N, vec3 V) { - float dotNV = max(dot(N, V), 0); - vec3 R = 2 * dot(N, V) * N - V; - R = vec3(-R.y, R.z, -R.x); - vec3 color = textureLod(samplerEnvironment, R, roughness * 5).rgb; - vec2 envBRDF = texture(samplerBRDFLUT, vec2(roughness, dotNV)).xy; - return color * (fresnel * envBRDF.x + envBRDF.y); -} - -void main() { - vec2 res = vec2(cameraBuffer.width, cameraBuffer.height); - - vec2 uv = inUV; - float z = texture(samplerSmoothedDepthLinear, uv).x; - if (z >= 0.0) { - discard; - } - vec3 csp = getCameraSpacePosition(uv, z); - - uv = inUV + vec2(-1 / cameraBuffer.width, 0); - vec3 cspnx = getCameraSpacePosition(uv, texture(samplerSmoothedDepthLinear, uv).x); - - uv = inUV + vec2(1 / cameraBuffer.width, 0); - vec3 csppx = getCameraSpacePosition(uv, texture(samplerSmoothedDepthLinear, uv).x); - - uv = inUV + vec2(0, -1 / cameraBuffer.height); - vec3 cspny = getCameraSpacePosition(uv, texture(samplerSmoothedDepthLinear, uv).x); - - uv = inUV + vec2(0, 1 / cameraBuffer.height); - vec3 csppy = getCameraSpacePosition(uv, texture(samplerSmoothedDepthLinear, uv).x); - - vec3 pdx = csppx - csp; - vec3 ndx = csp - cspnx; - - vec3 pdy = csppy - csp; - vec3 ndy = csp - cspny; - - vec3 dx = abs(pdx.z) < abs(ndx.z) ? pdx: ndx; - vec3 dy = abs(pdy.z) < abs(ndy.z) ? pdy: ndy; - - vec3 normal = normalize(cross(dy, dx)); - vec3 albedo = texture(samplerPointColor, inUV).rgb; - - // lighting - vec4 csPosition = vec4(csp, 1.0); - vec3 camDir = -normalize(csPosition.xyz); - vec4 ndc = cameraBuffer.projectionMatrix * csPosition; - ndc /= ndc.w; - - float roughness = 0.5; - vec3 fresnel = vec3(0.2); - - vec3 color = vec3(0.0); - - - // point light - for (int i = 0; i < NUM_POINT_LIGHT_SHADOWS; ++i) { - vec3 pos = world2camera(vec4(sceneBuffer.pointLights[i].position.xyz, 1.f)).xyz; - mat4 shadowProj = shadowBuffer.pointLightBuffers[6 * i].projectionMatrix; - - vec3 l = pos - csPosition.xyz; - vec3 wsl = vec3(cameraBuffer.viewMatrixInverse * vec4(l, 0)); - - vec3 v = abs(wsl); - vec4 p = shadowProj * vec4(0, 0, -max(max(v.x, v.y), v.z), 1); - float pixelDepth = p.z / p.w; - float shadowDepth = texture(samplerPointLightDepths[i], wsl).x; - - float visibility = step(pixelDepth - shadowDepth, 0); - color += visibility * computePointLight( - sceneBuffer.pointLights[i].emission.rgb, - l, normal, camDir, albedo, roughness, fresnel); - } - - for (int i = NUM_POINT_LIGHT_SHADOWS; i < NUM_POINT_LIGHTS; i++) { - vec3 pos = world2camera(vec4(sceneBuffer.pointLights[i].position.xyz, 1.f)).xyz; - vec3 l = pos - csPosition.xyz; - color += computePointLight( - sceneBuffer.pointLights[i].emission.rgb, - l, normal, camDir, albedo, roughness, fresnel); - } - - // directional light - for (int i = 0; i < NUM_DIRECTIONAL_LIGHT_SHADOWS; ++i) { - mat4 shadowView = shadowBuffer.directionalLightBuffers[i].viewMatrix; - mat4 shadowProj = shadowBuffer.directionalLightBuffers[i].projectionMatrix; - - vec3 lightDir = mat3(cameraBuffer.viewMatrix) * sceneBuffer.directionalLights[i].direction.xyz; - - vec4 ssPosition = shadowView * cameraBuffer.viewMatrixInverse * vec4((csPosition.xyz), 1); - vec4 shadowMapCoord = shadowProj * ssPosition; - shadowMapCoord /= shadowMapCoord.w; - shadowMapCoord.xy = shadowMapCoord.xy * 0.5 + 0.5; - - float resolution = textureSize(samplerDirectionalLightDepths[i], 0).x; - float visibility = ShadowMapPCF( - samplerDirectionalLightDepths[i], shadowMapCoord.xyz, resolution, 1 / resolution, 1); - - color += visibility * computeDirectionalLight( - lightDir, - sceneBuffer.directionalLights[i].emission.rgb, - normal, camDir, albedo, roughness, fresnel); - } - - for (int i = NUM_DIRECTIONAL_LIGHT_SHADOWS; i < NUM_DIRECTIONAL_LIGHTS; ++i) { - color += computeDirectionalLight( - mat3(cameraBuffer.viewMatrix) * sceneBuffer.directionalLights[i].direction.xyz, - sceneBuffer.directionalLights[i].emission.rgb, - normal, camDir, albedo, roughness, fresnel); - } - - // spot light - for (int i = 0; i < NUM_SPOT_LIGHT_SHADOWS; ++i) { - mat4 shadowView = shadowBuffer.spotLightBuffers[i].viewMatrix; - mat4 shadowProj = shadowBuffer.spotLightBuffers[i].projectionMatrix; - - vec3 pos = world2camera(vec4(sceneBuffer.spotLights[i].position.xyz, 1.f)).xyz; - vec3 centerDir = mat3(cameraBuffer.viewMatrix) * sceneBuffer.spotLights[i].direction.xyz; - vec3 l = pos - csPosition.xyz; - - vec4 ssPosition = shadowView * cameraBuffer.viewMatrixInverse * vec4((csPosition.xyz), 1); - vec4 shadowMapCoord = shadowProj * ssPosition; - shadowMapCoord /= shadowMapCoord.w; - shadowMapCoord.xy = shadowMapCoord.xy * 0.5 + 0.5; - - float resolution = textureSize(samplerSpotLightDepths[i], 0).x; - float visibility = ShadowMapPCF( - samplerSpotLightDepths[i], shadowMapCoord.xyz, resolution, 1 / resolution, 1); - - color += visibility * computeSpotLight( - sceneBuffer.spotLights[i].emission.a, - sceneBuffer.spotLights[i].direction.a, - centerDir, - sceneBuffer.spotLights[i].emission.rgb, - l, normal, camDir, albedo, roughness, fresnel); - } - - for (int i = NUM_SPOT_LIGHT_SHADOWS; i < NUM_SPOT_LIGHTS; ++i) { - vec3 pos = world2camera(vec4(sceneBuffer.spotLights[i].position.xyz, 1.f)).xyz; - vec3 l = pos - csPosition.xyz; - vec3 centerDir = mat3(cameraBuffer.viewMatrix) * sceneBuffer.spotLights[i].direction.xyz; - color += computeSpotLight( - sceneBuffer.spotLights[i].emission.a, - sceneBuffer.spotLights[i].direction.a, - centerDir, - sceneBuffer.spotLights[i].emission.rgb, - l, normal, camDir, albedo, roughness, fresnel); - } - - // textured light - for (int i = 0; i < NUM_TEXTURED_LIGHT_SHADOWS; ++i) { - mat4 shadowView = shadowBuffer.texturedLightBuffers[i].viewMatrix; - mat4 shadowProj = shadowBuffer.texturedLightBuffers[i].projectionMatrix; - - vec3 pos = world2camera(vec4(sceneBuffer.texturedLights[i].position.xyz, 1.f)).xyz; - vec3 centerDir = mat3(cameraBuffer.viewMatrix) * sceneBuffer.texturedLights[i].direction.xyz; - vec3 l = pos - csPosition.xyz; - - vec4 ssPosition = shadowView * cameraBuffer.viewMatrixInverse * vec4((csPosition.xyz), 1); - vec4 shadowMapCoord = shadowProj * ssPosition; - shadowMapCoord /= shadowMapCoord.w; - shadowMapCoord.xy = shadowMapCoord.xy * 0.5 + 0.5; - - float resolution = textureSize(samplerTexturedLightDepths[i], 0).x; - float visibility = ShadowMapPCF( - samplerTexturedLightDepths[i], shadowMapCoord.xyz, resolution, 1 / resolution, 1); - visibility *= texture(samplerTexturedLightTextures[i], shadowMapCoord.xy).x; - - color += visibility * computeSpotLight( - sceneBuffer.texturedLights[i].emission.a, - sceneBuffer.texturedLights[i].direction.a, - centerDir, - sceneBuffer.texturedLights[i].emission.rgb, - l, normal, camDir, albedo, roughness, fresnel); - } - - // environmental light - vec3 wnormal = mat3(cameraBuffer.viewMatrixInverse) * normal; - color += diffuseIBL(albedo, wnormal); - color += specularIBL(fresnel, roughness, - wnormal, - mat3(cameraBuffer.viewMatrixInverse) * camDir); - - color += sceneBuffer.ambientLight.rgb * albedo; - - outPoint.xyz = color; -} diff --git a/mani_skill2/envs/ms2/mpm/shader/point/composite2.frag b/mani_skill2/envs/ms2/mpm/shader/point/composite2.frag deleted file mode 100644 index fe6bba2c9..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/point/composite2.frag +++ /dev/null @@ -1,176 +0,0 @@ -#version 450 - -layout(set = 0, binding = 0) uniform sampler2D samplerLighting; -layout(set = 0, binding = 1) uniform sampler2D samplerLighting1; -layout(set = 0, binding = 2) uniform sampler2D samplerAlbedo2; -layout(set = 0, binding = 3) uniform sampler2D samplerGbufferDepth; -layout(set = 0, binding = 4) uniform sampler2D samplerGbuffer1Depth; -layout(set = 0, binding = 5) uniform sampler2D samplerGbuffer2Depth; -layout(set = 0, binding = 6) uniform usampler2D samplerSegmentation0; -layout(set = 0, binding = 7) uniform usampler2D samplerSegmentation1; -layout(set = 0, binding = 8) uniform sampler2D samplerLineDepth; -layout(set = 0, binding = 9) uniform sampler2D samplerLine; -layout(set = 0, binding = 10) uniform sampler2D samplerSmoothedDepthLinear; -layout(set = 0, binding = 11) uniform sampler2D samplerPoint; - -layout(set = 0, binding = 12) uniform sampler2D samplerPosition0; -layout(set = 0, binding = 13) uniform sampler2D samplerPosition1; - - -layout(location = 0) in vec2 inUV; -layout(location = 0) out vec4 outColor; -layout(location = 1) out vec4 outDepthLinear; -layout(location = 2) out uvec4 outSegmentation; -layout(location = 3) out vec4 outSegmentationView0; -layout(location = 4) out vec4 outSegmentationView1; -layout(location = 5) out vec4 outPosition; - - -layout(set = 1, binding = 0) uniform CameraBuffer { - mat4 viewMatrix; - mat4 projectionMatrix; - mat4 viewMatrixInverse; - mat4 projectionMatrixInverse; - mat4 prevViewMatrix; - mat4 prevViewMatrixInverse; - float width; - float height; -} cameraBuffer; - -vec4 colors[60] = { - vec4(0.8, 0.4, 0.4 , 1 ), - vec4(0.8, 0.41, 0.24, 1 ), - vec4(0.8, 0.75, 0.32, 1 ), - vec4(0.6, 0.8, 0.4 , 1 ), - vec4(0.35, 0.8, 0.24, 1 ), - vec4(0.32, 0.8, 0.51, 1 ), - vec4(0.4, 0.8, 0.8 , 1 ), - vec4(0.24, 0.63, 0.8 , 1 ), - vec4(0.32, 0.37, 0.8 , 1 ), - vec4(0.6, 0.4, 0.8 , 1 ), - vec4(0.69, 0.24, 0.8 , 1 ), - vec4(0.8, 0.32, 0.61, 1 ), - vec4(0.8, 0.32, 0.32, 1 ), - vec4(0.8, 0.64, 0.4 , 1 ), - vec4(0.8, 0.74, 0.24, 1 ), - vec4(0.56, 0.8, 0.32, 1 ), - vec4(0.4, 0.8, 0.44, 1 ), - vec4(0.24, 0.8, 0.46, 1 ), - vec4(0.32, 0.8, 0.8 , 1 ), - vec4(0.4, 0.56, 0.8 , 1 ), - vec4(0.24, 0.3, 0.8 , 1 ), - vec4(0.56, 0.32, 0.8 , 1 ), - vec4(0.8, 0.4, 0.76, 1 ), - vec4(0.8, 0.24, 0.58, 1 ), - vec4(0.8, 0.24, 0.24, 1 ), - vec4(0.8, 0.61, 0.32, 1 ), - vec4(0.72, 0.8, 0.4 , 1 ), - vec4(0.52, 0.8, 0.24, 1 ), - vec4(0.32, 0.8, 0.37, 1 ), - vec4(0.4, 0.8, 0.68, 1 ), - vec4(0.24, 0.8, 0.8 , 1 ), - vec4(0.32, 0.51, 0.8 , 1 ), - vec4(0.48, 0.4, 0.8 , 1 ), - vec4(0.52, 0.24, 0.8 , 1 ), - vec4(0.8, 0.32, 0.75, 1 ), - vec4(0.8, 0.4, 0.52, 1 ), - vec4(0.8, 0.52, 0.4 , 1 ), - vec4(0.8, 0.58, 0.24, 1 ), - vec4(0.7, 0.8, 0.32, 1 ), - vec4(0.48, 0.8, 0.4 , 1 ), - vec4(0.24, 0.8, 0.3 , 1 ), - vec4(0.32, 0.8, 0.66, 1 ), - vec4(0.4, 0.68, 0.8 , 1 ), - vec4(0.24, 0.46, 0.8 , 1 ), - vec4(0.42, 0.32, 0.8 , 1 ), - vec4(0.72, 0.4, 0.8 , 1 ), - vec4(0.8, 0.24, 0.74, 1 ), - vec4(0.8, 0.32, 0.46, 1 ), - vec4(0.8, 0.46, 0.32, 1 ), - vec4(0.8, 0.76, 0.4 , 1 ), - vec4(0.69, 0.8, 0.24, 1 ), - vec4(0.42, 0.8, 0.32, 1 ), - vec4(0.4, 0.8, 0.56, 1 ), - vec4(0.24, 0.8, 0.63, 1 ), - vec4(0.32, 0.66, 0.8 , 1 ), - vec4(0.4, 0.44, 0.8 , 1 ), - vec4(0.35, 0.24, 0.8 , 1 ), - vec4(0.7, 0.32, 0.8 , 1 ), - vec4(0.8, 0.4, 0.64, 1 ), - vec4(0.8, 0.24, 0.41, 1 ) -}; - - -void main() { - float d0 = texture(samplerGbufferDepth, inUV).x; - float d1 = texture(samplerGbuffer1Depth, inUV).x; - float d2 = texture(samplerGbuffer2Depth, inUV).x; - - float pointDepth = texture(samplerSmoothedDepthLinear, inUV).x; - float dp = (cameraBuffer.projectionMatrix[2][2] * pointDepth + cameraBuffer.projectionMatrix[3][2]) / (-pointDepth); - - vec4 pointColor = texture(samplerPoint, inUV); - - vec4 outColor0 = texture(samplerLighting, inUV); - vec4 outColor1 = texture(samplerLighting1, inUV); - vec4 outColor2 = texture(samplerAlbedo2, inUV); - - vec4 outPos0 = vec4(texture(samplerPosition0, inUV).xyz, d0); - vec4 outPos1 = vec4(texture(samplerPosition1, inUV).xyz, d1); - - // depth composite for 0 and 2 - float factor = step(d0, d2); - float d = min(d0, d2); - outColor = outColor0 * factor + outColor2 * (1 - factor); - outPosition = outPos0; - - // depth composite for 02 and p - if (pointDepth < 0) { - factor = step(d, dp); - d = min(d, dp); - outColor = outColor * factor + pointColor * (1 - factor); - - // convert dp to pointPos - vec2 res = vec2(cameraBuffer.width, cameraBuffer.height); - vec2 pixelNdc = gl_FragCoord.xy / res * 2.0 - 1.0; - vec4 pointPos = cameraBuffer.projectionMatrixInverse * vec4(pixelNdc, dp, 1.0); - pointPos /= pointPos.w; - pointPos.w = dp; - - outPosition = outPosition * factor + pointPos * (1 - factor); - } - - // blend for 02p and 1 - vec3 blend = outColor1.a * outColor1.rgb + (1 - outColor1.a) * outColor.rgb; - factor = step(d, d1); - outColor = vec4((1 - factor) * blend + factor * outColor.rgb, 1.f); - - // outPosition = outPos0 * factor + outPos1 * (1 - factor); - outPosition = outPosition * factor + outPos1 * (1 - factor); - - // TODO: position for points - - // float pointDepth = texture(samplerSmoothedDepthLinear, inUV).x; - // if (pointDepth < 0 && (pointDepth > outPosition.z || outPosition.z == 0)) { - // outColor = vec4(pointColor.xyz, 1); - // } - - outColor = pow(outColor, vec4(1/2.2, 1/2.2, 1/2.2, 1)); - - // vec4 csPosition = cameraBuffer.projectionMatrixInverse * (vec4(inUV * 2 - 1, min(d0, d1), 1)); - outDepthLinear = vec4(vec3(-outPosition.z), 1.); - - uvec4 seg0 = texture(samplerSegmentation0, inUV); - uvec4 seg1 = texture(samplerSegmentation1, inUV); - outSegmentation = d0 < d1 ? seg0 : seg1; - - outSegmentationView0 = mix(vec4(0,0,0,1), colors[outSegmentation.x % 60], sign(outSegmentation.x)); - outSegmentationView1 = mix(vec4(0,0,0,1), colors[outSegmentation.y % 60], sign(outSegmentation.y)); - - vec4 lineColor = texture(samplerLine, inUV); - if (texture(samplerLineDepth, inUV).x < 1) { - outColor = vec4(lineColor.xyz, 1); - } - - outColor = clamp(outColor, vec4(0), vec4(1)); -} diff --git a/mani_skill2/envs/ms2/mpm/shader/point/deferred.frag b/mani_skill2/envs/ms2/mpm/shader/point/deferred.frag deleted file mode 100644 index a62bf8e7e..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/point/deferred.frag +++ /dev/null @@ -1,250 +0,0 @@ -#version 450 - -layout (constant_id = 0) const int NUM_DIRECTIONAL_LIGHTS = 3; -layout (constant_id = 1) const int NUM_POINT_LIGHTS = 10; -layout (constant_id = 2) const int NUM_DIRECTIONAL_LIGHT_SHADOWS = 1; -layout (constant_id = 3) const int NUM_POINT_LIGHT_SHADOWS = 3; -layout (constant_id = 4) const int NUM_TEXTURED_LIGHT_SHADOWS = 1; -layout (constant_id = 5) const int NUM_SPOT_LIGHT_SHADOWS = 10; -layout (constant_id = 6) const int NUM_SPOT_LIGHTS = 10; - -#include "../common/view.glsl" -#include "../common/lights.glsl" -#include "../common/shadow.glsl" - -layout(set = 0, binding = 0) uniform SceneBuffer { - vec4 ambientLight; - DirectionalLight directionalLights[3]; - SpotLight spotLights[10]; - PointLight pointLights[10]; - SpotLight texturedLights[1]; -} sceneBuffer; - -struct LightBuffer { - mat4 viewMatrix; - mat4 viewMatrixInverse; - mat4 projectionMatrix; - mat4 projectionMatrixInverse; - int width; - int height; -}; - -layout(set = 0, binding = 1) uniform ShadowBuffer { - LightBuffer directionalLightBuffers[3]; - LightBuffer spotLightBuffers[10]; - LightBuffer pointLightBuffers[60]; - LightBuffer texturedLightBuffers[1]; -} shadowBuffer; - -layout(set = 0, binding = 2) uniform samplerCube samplerPointLightDepths[3]; -layout(set = 0, binding = 3) uniform sampler2D samplerDirectionalLightDepths[1]; -layout(set = 0, binding = 4) uniform sampler2D samplerTexturedLightDepths[1]; -layout(set = 0, binding = 5) uniform sampler2D samplerSpotLightDepths[10]; -layout(set = 0, binding = 6) uniform sampler2D samplerTexturedLightTextures[1]; - -layout(set = 1, binding = 0) uniform CameraBuffer { - mat4 viewMatrix; - mat4 projectionMatrix; - mat4 viewMatrixInverse; - mat4 projectionMatrixInverse; - mat4 prevViewMatrix; - mat4 prevViewMatrixInverse; - float width; - float height; -} cameraBuffer; - -layout(set = 2, binding = 0) uniform sampler2D samplerAlbedo; -layout(set = 2, binding = 1) uniform sampler2D samplerPosition0; -layout(set = 2, binding = 2) uniform sampler2D samplerSpecular; -layout(set = 2, binding = 3) uniform sampler2D samplerNormal; -layout(set = 2, binding = 4) uniform sampler2D samplerEmission; -layout(set = 2, binding = 5) uniform sampler2D samplerGbufferDepth; -layout(set = 2, binding = 6) uniform sampler2D samplerCustom; - -// IBL -layout(set = 2, binding = 7) uniform samplerCube samplerEnvironment; -layout(set = 2, binding = 8) uniform sampler2D samplerBRDFLUT; - -layout(location = 0) in vec2 inUV; -layout(location = 0) out vec4 outLighting; - -vec4 world2camera(vec4 pos) { - return cameraBuffer.viewMatrix * pos; -} - -vec3 getBackgroundColor(vec3 texcoord) { - texcoord = vec3(-texcoord.y, texcoord.z, -texcoord.x); - return textureLod(samplerEnvironment, texcoord, 0).rgb; -} - -vec3 diffuseIBL(vec3 albedo, vec3 N) { - N = vec3(-N.y, N.z, -N.x); - vec3 color = textureLod(samplerEnvironment, N, 5).rgb; - return color * albedo; -} - -vec3 specularIBL(vec3 fresnel, float roughness, vec3 N, vec3 V) { - float dotNV = max(dot(N, V), 0); - vec3 R = 2 * dot(N, V) * N - V; - R = vec3(-R.y, R.z, -R.x); - vec3 color = textureLod(samplerEnvironment, R, roughness * 5).rgb; - vec2 envBRDF = texture(samplerBRDFLUT, vec2(roughness, dotNV)).xy; - return color * (fresnel * envBRDF.x + envBRDF.y); -} - -void main() { - vec3 albedo = texture(samplerAlbedo, inUV).xyz; - vec3 frm = texture(samplerSpecular, inUV).xyz; - float specular = frm.x; - float roughness = frm.y; - float metallic = frm.z; - - vec3 normal = normalize(texture(samplerNormal, inUV).xyz); - float depth = texture(samplerGbufferDepth, inUV).x; - - vec4 csPosition = cameraBuffer.projectionMatrixInverse * (vec4(inUV * 2 - 1, depth, 1)); - csPosition /= csPosition.w; - - vec3 camDir = -normalize(csPosition.xyz); - - vec3 diffuseAlbedo = albedo * (1 - metallic); - vec3 fresnel = specular * (1 - metallic) + albedo * metallic; - - vec3 color = texture(samplerEmission, inUV).rgb; - - // point light - for (int i = 0; i < NUM_POINT_LIGHT_SHADOWS; ++i) { - vec3 pos = world2camera(vec4(sceneBuffer.pointLights[i].position.xyz, 1.f)).xyz; - mat4 shadowProj = shadowBuffer.pointLightBuffers[6 * i].projectionMatrix; - - vec3 l = pos - csPosition.xyz; - vec3 wsl = vec3(cameraBuffer.viewMatrixInverse * vec4(l, 0)); - - vec3 v = abs(wsl); - vec4 p = shadowProj * vec4(0, 0, -max(max(v.x, v.y), v.z), 1); - float pixelDepth = p.z / p.w; - float shadowDepth = texture(samplerPointLightDepths[i], wsl).x; - - float visibility = step(pixelDepth - shadowDepth, 0); - color += visibility * computePointLight( - sceneBuffer.pointLights[i].emission.rgb, - l, normal, camDir, diffuseAlbedo, roughness, fresnel); - } - - for (int i = NUM_POINT_LIGHT_SHADOWS; i < NUM_POINT_LIGHTS; i++) { - vec3 pos = world2camera(vec4(sceneBuffer.pointLights[i].position.xyz, 1.f)).xyz; - vec3 l = pos - csPosition.xyz; - color += computePointLight( - sceneBuffer.pointLights[i].emission.rgb, - l, normal, camDir, diffuseAlbedo, roughness, fresnel); - } - - // directional light - for (int i = 0; i < NUM_DIRECTIONAL_LIGHT_SHADOWS; ++i) { - mat4 shadowView = shadowBuffer.directionalLightBuffers[i].viewMatrix; - mat4 shadowProj = shadowBuffer.directionalLightBuffers[i].projectionMatrix; - - vec3 lightDir = mat3(cameraBuffer.viewMatrix) * sceneBuffer.directionalLights[i].direction.xyz; - - vec4 ssPosition = shadowView * cameraBuffer.viewMatrixInverse * vec4((csPosition.xyz), 1); - vec4 shadowMapCoord = shadowProj * ssPosition; - shadowMapCoord /= shadowMapCoord.w; - shadowMapCoord.xy = shadowMapCoord.xy * 0.5 + 0.5; - - float resolution = textureSize(samplerDirectionalLightDepths[i], 0).x; - float visibility = ShadowMapPCF( - samplerDirectionalLightDepths[i], shadowMapCoord.xyz, resolution, 1 / resolution, 1); - - color += visibility * computeDirectionalLight( - lightDir, - sceneBuffer.directionalLights[i].emission.rgb, - normal, camDir, diffuseAlbedo, roughness, fresnel); - } - - for (int i = NUM_DIRECTIONAL_LIGHT_SHADOWS; i < NUM_DIRECTIONAL_LIGHTS; ++i) { - color += computeDirectionalLight( - mat3(cameraBuffer.viewMatrix) * sceneBuffer.directionalLights[i].direction.xyz, - sceneBuffer.directionalLights[i].emission.rgb, - normal, camDir, diffuseAlbedo, roughness, fresnel); - } - - // spot light - for (int i = 0; i < NUM_SPOT_LIGHT_SHADOWS; ++i) { - mat4 shadowView = shadowBuffer.spotLightBuffers[i].viewMatrix; - mat4 shadowProj = shadowBuffer.spotLightBuffers[i].projectionMatrix; - - vec3 pos = world2camera(vec4(sceneBuffer.spotLights[i].position.xyz, 1.f)).xyz; - vec3 centerDir = mat3(cameraBuffer.viewMatrix) * sceneBuffer.spotLights[i].direction.xyz; - vec3 l = pos - csPosition.xyz; - - vec4 ssPosition = shadowView * cameraBuffer.viewMatrixInverse * vec4((csPosition.xyz), 1); - vec4 shadowMapCoord = shadowProj * ssPosition; - shadowMapCoord /= shadowMapCoord.w; - shadowMapCoord.xy = shadowMapCoord.xy * 0.5 + 0.5; - - float resolution = textureSize(samplerSpotLightDepths[i], 0).x; - float visibility = ShadowMapPCF( - samplerSpotLightDepths[i], shadowMapCoord.xyz, resolution, 1 / resolution, 1); - - color += visibility * computeSpotLight( - sceneBuffer.spotLights[i].emission.a, - sceneBuffer.spotLights[i].direction.a, - centerDir, - sceneBuffer.spotLights[i].emission.rgb, - l, normal, camDir, diffuseAlbedo, roughness, fresnel); - } - - for (int i = NUM_SPOT_LIGHT_SHADOWS; i < NUM_SPOT_LIGHTS; ++i) { - vec3 pos = world2camera(vec4(sceneBuffer.spotLights[i].position.xyz, 1.f)).xyz; - vec3 l = pos - csPosition.xyz; - vec3 centerDir = mat3(cameraBuffer.viewMatrix) * sceneBuffer.spotLights[i].direction.xyz; - color += computeSpotLight( - sceneBuffer.spotLights[i].emission.a, - sceneBuffer.spotLights[i].direction.a, - centerDir, - sceneBuffer.spotLights[i].emission.rgb, - l, normal, camDir, diffuseAlbedo, roughness, fresnel); - } - - // textured light - for (int i = 0; i < NUM_TEXTURED_LIGHT_SHADOWS; ++i) { - mat4 shadowView = shadowBuffer.texturedLightBuffers[i].viewMatrix; - mat4 shadowProj = shadowBuffer.texturedLightBuffers[i].projectionMatrix; - - vec3 pos = world2camera(vec4(sceneBuffer.texturedLights[i].position.xyz, 1.f)).xyz; - vec3 centerDir = mat3(cameraBuffer.viewMatrix) * sceneBuffer.texturedLights[i].direction.xyz; - vec3 l = pos - csPosition.xyz; - - vec4 ssPosition = shadowView * cameraBuffer.viewMatrixInverse * vec4((csPosition.xyz), 1); - vec4 shadowMapCoord = shadowProj * ssPosition; - shadowMapCoord /= shadowMapCoord.w; - shadowMapCoord.xy = shadowMapCoord.xy * 0.5 + 0.5; - - float resolution = textureSize(samplerTexturedLightDepths[i], 0).x; - float visibility = ShadowMapPCF( - samplerTexturedLightDepths[i], shadowMapCoord.xyz, resolution, 1 / resolution, 1); - visibility *= texture(samplerTexturedLightTextures[i], shadowMapCoord.xy).x; - - color += visibility * computeSpotLight( - sceneBuffer.texturedLights[i].emission.a, - sceneBuffer.texturedLights[i].direction.a, - centerDir, - sceneBuffer.texturedLights[i].emission.rgb, - l, normal, camDir, diffuseAlbedo, roughness, fresnel); - } - - // environmental light - vec3 wnormal = mat3(cameraBuffer.viewMatrixInverse) * normal; - color += diffuseIBL(diffuseAlbedo, wnormal); - color += specularIBL(fresnel, roughness, - wnormal, - mat3(cameraBuffer.viewMatrixInverse) * camDir); - - color += sceneBuffer.ambientLight.rgb * diffuseAlbedo; - - if (depth == 1) { - outLighting = vec4(getBackgroundColor((cameraBuffer.viewMatrixInverse * csPosition).xyz), 1.f); - } else { - outLighting = vec4(color, 1); - } -} diff --git a/mani_skill2/envs/ms2/mpm/shader/point/deferred.vert b/mani_skill2/envs/ms2/mpm/shader/point/deferred.vert deleted file mode 100644 index 520c95688..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/point/deferred.vert +++ /dev/null @@ -1,9 +0,0 @@ -#version 450 - -layout (location = 0) out vec2 outUV; - -void main() -{ - outUV = vec2((gl_VertexIndex << 1) & 2, gl_VertexIndex & 2); - gl_Position = vec4(outUV * 2.0f + -1.0f, 0.0f, 1.0f); -} diff --git a/mani_skill2/envs/ms2/mpm/shader/point/gbuffer.frag b/mani_skill2/envs/ms2/mpm/shader/point/gbuffer.frag deleted file mode 100644 index f834bde49..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/point/gbuffer.frag +++ /dev/null @@ -1,115 +0,0 @@ -#version 450 - -layout(set = 0, binding = 0) uniform CameraBuffer { - mat4 viewMatrix; - mat4 projectionMatrix; - mat4 viewMatrixInverse; - mat4 projectionMatrixInverse; - mat4 prevViewMatrix; - mat4 prevViewMatrixInverse; - float width; - float height; -} cameraBuffer; - -layout(set = 1, binding = 0) uniform ObjectBuffer { - mat4 modelMatrix; - mat4 prevModelMatrix; - uvec4 segmentation; - float transparency; - int shadeFlat; -} objectBuffer; - -layout(set = 2, binding = 0) uniform MaterialBuffer { - vec4 emission; - vec4 baseColor; - float fresnel; - float roughness; - float metallic; - float transmission; - float ior; - float transmissionRoughness; - int textureMask; - int padding1; -} materialBuffer; - -layout(set = 2, binding = 1) uniform sampler2D colorTexture; -layout(set = 2, binding = 2) uniform sampler2D roughnessTexture; -layout(set = 2, binding = 3) uniform sampler2D normalTexture; -layout(set = 2, binding = 4) uniform sampler2D metallicTexture; -layout(set = 2, binding = 5) uniform sampler2D emissionTexture; - -layout(location = 0) in vec4 inPosition; -layout(location = 1) in vec4 inPrevPosition; -layout(location = 2) in vec2 inUV; -layout(location = 3) in flat uvec4 inSegmentation; -layout(location = 4) in vec3 objectCoord; -layout(location = 5) in mat3 inTbn; - -layout(location = 0) out vec4 outAlbedo; -layout(location = 1) out vec4 outPosition0; -layout(location = 2) out vec4 outSpecular; -layout(location = 3) out vec4 outNormal; -layout(location = 4) out uvec4 outSegmentation0; -layout(location = 5) out vec4 outCustom; -layout(location = 6) out vec4 outMotionDirection; -layout(location = 7) out vec4 outEmission; - -void main() { - outCustom = vec4(objectCoord, 1); - outSegmentation0 = inSegmentation; - - outPosition0 = inPosition; - - vec4 p1 = cameraBuffer.projectionMatrix * inPosition; - p1 /= p1.w; - vec2 p1s = p1.xy / p1.z; - - vec4 p2 = cameraBuffer.projectionMatrix * inPrevPosition; - p2 /= p2.w; - vec2 p2s = p2.xy / p2.z; - - outMotionDirection = vec4((p1s - p2s)*0.5, 0, 1); - - if ((materialBuffer.textureMask & 16) != 0) { - outEmission = texture(emissionTexture, inUV); - } else { - outEmission = materialBuffer.emission; - } - - if ((materialBuffer.textureMask & 1) != 0) { - outAlbedo = texture(colorTexture, inUV); - } else { - outAlbedo = materialBuffer.baseColor; - } - - if (outAlbedo.a == 0) { - discard; - } - - outSpecular.r = materialBuffer.fresnel * 0.08; - - if ((materialBuffer.textureMask & 2) != 0) { - outSpecular.g = texture(roughnessTexture, inUV).r; - } else { - outSpecular.g = materialBuffer.roughness; - } - - if ((materialBuffer.textureMask & 8) != 0) { - outSpecular.b = texture(metallicTexture, inUV).r; - } else { - outSpecular.b = materialBuffer.metallic; - } - - if (objectBuffer.shadeFlat == 0) { - if ((materialBuffer.textureMask & 4) != 0) { - outNormal = vec4(normalize(inTbn * (texture(normalTexture, inUV).xyz * 2 - 1)), 0); - } else { - outNormal = vec4(normalize(inTbn * vec3(0, 0, 1)), 0); - } - } else { - vec4 fdx = dFdx(inPosition); - vec4 fdy = dFdy(inPosition); - vec3 normal = -normalize(cross(fdx.xyz, fdy.xyz)); - outNormal = vec4(normal, 0); - } -} diff --git a/mani_skill2/envs/ms2/mpm/shader/point/gbuffer.vert b/mani_skill2/envs/ms2/mpm/shader/point/gbuffer.vert deleted file mode 100644 index 359ca8041..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/point/gbuffer.vert +++ /dev/null @@ -1,50 +0,0 @@ -#version 450 -#extension GL_ARB_separate_shader_objects : enable -#extension GL_ARB_shading_language_420pack : enable - -layout(set = 0, binding = 0) uniform CameraBuffer { - mat4 viewMatrix; - mat4 projectionMatrix; - mat4 viewMatrixInverse; - mat4 projectionMatrixInverse; - mat4 prevViewMatrix; - mat4 prevViewMatrixInverse; - float width; - float height; -} cameraBuffer; - -layout(set = 1, binding = 0) uniform ObjectBuffer { - mat4 modelMatrix; - mat4 prevModelMatrix; - uvec4 segmentation; - float transparency; - int shadeFlat; -} objectBuffer; - -layout(location = 0) in vec3 position; -layout(location = 1) in vec3 normal; -layout(location = 2) in vec2 uv; -layout(location = 3) in vec3 tangent; -layout(location = 4) in vec3 bitangent; - -layout(location = 0) out vec4 outPosition; -layout(location = 1) out vec4 outPrevPosition; -layout(location = 2) out vec2 outUV; -layout(location = 3) out flat uvec4 outSegmentation; -layout(location = 4) out vec3 objectCoord; -layout(location = 5) out mat3 outTbn; - -void main() { - outSegmentation = objectBuffer.segmentation; - mat4 modelView = cameraBuffer.viewMatrix * objectBuffer.modelMatrix; - mat3 normalMatrix = mat3(transpose(inverse(modelView))); - objectCoord = position; - outPosition = modelView * vec4(position, 1); - outPrevPosition = cameraBuffer.prevViewMatrix * objectBuffer.prevModelMatrix * vec4(position, 1); - outUV = uv; - gl_Position = cameraBuffer.projectionMatrix * outPosition; - vec3 outTangent = normalize(normalMatrix * tangent); - vec3 outBitangent = normalize(normalMatrix * bitangent); - vec3 outNormal = normalize(normalMatrix * normal); - outTbn = mat3(outTangent, outBitangent, outNormal); -} diff --git a/mani_skill2/envs/ms2/mpm/shader/point/gbuffer1.frag b/mani_skill2/envs/ms2/mpm/shader/point/gbuffer1.frag deleted file mode 100644 index 172bfbf22..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/point/gbuffer1.frag +++ /dev/null @@ -1,256 +0,0 @@ -#version 450 - -layout (constant_id = 0) const int NUM_DIRECTIONAL_LIGHTS = 3; -layout (constant_id = 1) const int NUM_POINT_LIGHTS = 10; -layout (constant_id = 2) const int NUM_DIRECTIONAL_LIGHT_SHADOWS = 1; -layout (constant_id = 3) const int NUM_POINT_LIGHT_SHADOWS = 3; -layout (constant_id = 4) const int NUM_TEXTURED_LIGHT_SHADOWS = 1; -layout (constant_id = 5) const int NUM_SPOT_LIGHT_SHADOWS = 10; -layout (constant_id = 6) const int NUM_SPOT_LIGHTS = 10; - -layout(set = 0, binding = 0) uniform CameraBuffer { - mat4 viewMatrix; - mat4 projectionMatrix; - mat4 viewMatrixInverse; - mat4 projectionMatrixInverse; - mat4 prevViewMatrix; - mat4 prevViewMatrixInverse; - float width; - float height; -} cameraBuffer; - -layout(set = 1, binding = 0) uniform ObjectBuffer { - mat4 modelMatrix; - mat4 prevModelMatrix; - uvec4 segmentation; - float transparency; - int shadeFlat; -} objectBuffer; - -layout(set = 2, binding = 0) uniform MaterialBuffer { - vec4 emission; - vec4 baseColor; - float fresnel; - float roughness; - float metallic; - float transmission; - float ior; - float transmissionRoughness; - int textureMask; - int padding1; -} materialBuffer; - -layout(set = 2, binding = 1) uniform sampler2D colorTexture; -layout(set = 2, binding = 2) uniform sampler2D roughnessTexture; -layout(set = 2, binding = 3) uniform sampler2D normalTexture; -layout(set = 2, binding = 4) uniform sampler2D metallicTexture; -layout(set = 2, binding = 5) uniform sampler2D emissionTexture; - -#include "../common/lights.glsl" - -layout(set = 3, binding = 0) uniform SceneBuffer { - vec4 ambientLight; - DirectionalLight directionalLights[3]; - SpotLight spotLights[10]; - PointLight pointLights[10]; - SpotLight texturedLights[1]; -} sceneBuffer; - -struct LightBuffer { - mat4 viewMatrix; - mat4 viewMatrixInverse; - mat4 projectionMatrix; - mat4 projectionMatrixInverse; - int width; - int height; -}; - -layout(set = 3, binding = 1) uniform ShadowBuffer { - LightBuffer directionalLightBuffers[3]; - LightBuffer spotLightBuffers[10]; - LightBuffer pointLightBuffers[60]; - LightBuffer texturedLightBuffers[1]; -} shadowBuffer; - -layout(set = 3, binding = 2) uniform samplerCube samplerPointLightDepths[3]; -layout(set = 3, binding = 3) uniform sampler2D samplerDirectionalLightDepths[1]; -layout(set = 3, binding = 4) uniform sampler2D samplerTexturedLightDepths[1]; -layout(set = 3, binding = 5) uniform sampler2D samplerSpotLightDepths[10]; -layout(set = 3, binding = 6) uniform sampler2D samplerTexturedLightTextures[1]; - -layout(set = 3, binding = 7) uniform samplerCube samplerEnvironment; -layout(set = 3, binding = 8) uniform sampler2D samplerBRDFLUT; - -vec4 world2camera(vec4 pos) { - return cameraBuffer.viewMatrix * pos; -} - - -vec3 getBackgroundColor(vec3 texcoord) { - texcoord = vec3(-texcoord.y, texcoord.z, -texcoord.x); - return textureLod(samplerEnvironment, texcoord, 0).rgb; -} - -vec3 diffuseIBL(vec3 albedo, vec3 N) { - N = vec3(-N.y, N.z, -N.x); - vec3 color = textureLod(samplerEnvironment, N, 5).rgb; - return color * albedo; -} - -vec3 specularIBL(vec3 fresnel, float roughness, vec3 N, vec3 V) { - float dotNV = max(dot(N, V), 0); - vec3 R = 2 * dot(N, V) * N - V; - R = vec3(-R.y, R.z, -R.x); - vec3 color = textureLod(samplerEnvironment, R, roughness * 5).rgb; - vec2 envBRDF = texture(samplerBRDFLUT, vec2(roughness, dotNV)).xy; - return color * (fresnel * envBRDF.x + envBRDF.y); -} - - - - - -layout(location = 0) in vec4 inPosition; -layout(location = 1) in vec2 inUV; -layout(location = 2) in flat uvec4 inSegmentation; -layout(location = 3) in vec3 objectCoord; -layout(location = 4) in mat3 inTbn; - -layout(location = 0) out vec4 outLighting1; -layout(location = 1) out vec4 outNormal1; -layout(location = 2) out uvec4 outSegmentation1; -layout(location = 3) out vec4 outPosition1; - -void main() { - outSegmentation1 = inSegmentation; - - vec4 emission; - vec4 albedo; - vec4 frm; - - if ((materialBuffer.textureMask & 16) != 0) { - emission = texture(emissionTexture, inUV); - } else { - emission = materialBuffer.emission; - } - - if ((materialBuffer.textureMask & 1) != 0) { - albedo = texture(colorTexture, inUV); - } else { - albedo = materialBuffer.baseColor; - } - - albedo.a *= (1.f - objectBuffer.transparency); - - if (albedo.a == 0) { - discard; - } - - frm.r = materialBuffer.fresnel * 0.08; - - if ((materialBuffer.textureMask & 2) != 0) { - frm.g = texture(roughnessTexture, inUV).r; - } else { - frm.g = materialBuffer.roughness; - } - - if ((materialBuffer.textureMask & 8) != 0) { - frm.b = texture(metallicTexture, inUV).r; - } else { - frm.b = materialBuffer.metallic; - } - - if (objectBuffer.shadeFlat == 0) { - if ((materialBuffer.textureMask & 4) != 0) { - outNormal1 = vec4(normalize(inTbn * (texture(normalTexture, inUV).xyz * 2 - 1)), 0); - } else { - outNormal1 = vec4(normalize(inTbn * vec3(0, 0, 1)), 0); - } - } else { - vec4 fdx = dFdx(inPosition); - vec4 fdy = dFdy(inPosition); - vec3 normal = -normalize(cross(fdx.xyz, fdy.xyz)); - outNormal1 = vec4(normal, 0); - } - - outPosition1 = inPosition; - - float specular = frm.x; - float roughness = frm.y; - float metallic = frm.z; - - vec3 normal = outNormal1.xyz; - vec4 csPosition = inPosition; - csPosition /= csPosition.w; - - vec3 camDir = -normalize(csPosition.xyz); - - vec3 diffuseAlbedo = albedo.rgb * (1 - metallic); - vec3 fresnel = specular * (1 - metallic) + albedo.rgb * metallic; - - vec3 color = emission.rgb; - - for (int i = 0; i < NUM_POINT_LIGHTS; i++) { - vec3 pos = world2camera(vec4(sceneBuffer.pointLights[i].position.xyz, 1.f)).xyz; - vec3 l = pos - csPosition.xyz; - color += computePointLight( - sceneBuffer.pointLights[i].emission.rgb, - l, normal, camDir, diffuseAlbedo, roughness, fresnel); - } - - for (int i = 0; i < NUM_DIRECTIONAL_LIGHTS; ++i) { - color += computeDirectionalLight( - mat3(cameraBuffer.viewMatrix) * sceneBuffer.directionalLights[i].direction.xyz, - sceneBuffer.directionalLights[i].emission.rgb, - normal, camDir, diffuseAlbedo, roughness, fresnel); - } - - for (int i = 0; i < NUM_SPOT_LIGHTS; ++i) { - vec3 pos = world2camera(vec4(sceneBuffer.spotLights[i].position.xyz, 1.f)).xyz; - vec3 l = pos - csPosition.xyz; - vec3 centerDir = mat3(cameraBuffer.viewMatrix) * sceneBuffer.spotLights[i].direction.xyz; - color += computeSpotLight( - sceneBuffer.spotLights[i].emission.a, - sceneBuffer.spotLights[i].direction.a, - centerDir, - sceneBuffer.spotLights[i].emission.rgb, - l, normal, camDir, diffuseAlbedo, roughness, fresnel); - } - - for (int i = 0; i < NUM_TEXTURED_LIGHT_SHADOWS; ++i) { - mat4 shadowView = shadowBuffer.texturedLightBuffers[i].viewMatrix; - mat4 shadowProj = shadowBuffer.texturedLightBuffers[i].projectionMatrix; - - vec3 pos = world2camera(vec4(sceneBuffer.texturedLights[i].position.xyz, 1.f)).xyz; - vec3 centerDir = mat3(cameraBuffer.viewMatrix) * sceneBuffer.texturedLights[i].direction.xyz; - vec3 l = pos - csPosition.xyz; - - float bias = 0; - - vec4 ssPosition = shadowView * cameraBuffer.viewMatrixInverse * vec4((csPosition.xyz), 1); - ssPosition.z += bias; - vec4 shadowMapCoord = shadowProj * ssPosition; - shadowMapCoord /= shadowMapCoord.w; - shadowMapCoord.xy = shadowMapCoord.xy * 0.5 + 0.5; - - float resolution = textureSize(samplerTexturedLightDepths[i], 0).x; - float visibility = texture(samplerTexturedLightTextures[i], shadowMapCoord.xy).x; - - color += visibility * computeSpotLight( - sceneBuffer.texturedLights[i].emission.a, - sceneBuffer.texturedLights[i].direction.a, - centerDir, - sceneBuffer.texturedLights[i].emission.rgb, - l, normal, camDir, diffuseAlbedo, roughness, fresnel); - } - - vec3 wnormal = mat3(cameraBuffer.viewMatrixInverse) * normal; - color += diffuseIBL(diffuseAlbedo, wnormal); - color += specularIBL(fresnel, roughness, - wnormal, - mat3(cameraBuffer.viewMatrixInverse) * camDir); - - color += sceneBuffer.ambientLight.rgb * diffuseAlbedo; - - outLighting1 = vec4(color, albedo.a); -} diff --git a/mani_skill2/envs/ms2/mpm/shader/point/gbuffer1.vert b/mani_skill2/envs/ms2/mpm/shader/point/gbuffer1.vert deleted file mode 100644 index c225a9504..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/point/gbuffer1.vert +++ /dev/null @@ -1,48 +0,0 @@ -#version 450 -#extension GL_ARB_separate_shader_objects : enable -#extension GL_ARB_shading_language_420pack : enable - -layout(set = 0, binding = 0) uniform CameraBuffer { - mat4 viewMatrix; - mat4 projectionMatrix; - mat4 viewMatrixInverse; - mat4 projectionMatrixInverse; - mat4 prevViewMatrix; - mat4 prevViewMatrixInverse; - float width; - float height; -} cameraBuffer; - -layout(set = 1, binding = 0) uniform ObjectBuffer { - mat4 modelMatrix; - mat4 prevModelMatrix; - uvec4 segmentation; - float transparency; - int shadeFlat; -} objectBuffer; - -layout(location = 0) in vec3 position; -layout(location = 1) in vec3 normal; -layout(location = 2) in vec2 uv; -layout(location = 3) in vec3 tangent; -layout(location = 4) in vec3 bitangent; - -layout(location = 0) out vec4 outPosition; -layout(location = 1) out vec2 outUV; -layout(location = 2) out flat uvec4 outSegmentation; -layout(location = 3) out vec3 objectCoord; -layout(location = 4) out mat3 outTbn; - -void main() { - outSegmentation = objectBuffer.segmentation; - mat4 modelView = cameraBuffer.viewMatrix * objectBuffer.modelMatrix; - mat3 normalMatrix = mat3(transpose(inverse(modelView))); - objectCoord = position; - outPosition = modelView * vec4(position, 1); - outUV = uv; - gl_Position = cameraBuffer.projectionMatrix * outPosition; - vec3 outTangent = normalize(normalMatrix * tangent); - vec3 outBitangent = normalize(normalMatrix * bitangent); - vec3 outNormal = normalize(normalMatrix * normal); - outTbn = mat3(outTangent, outBitangent, outNormal); -} diff --git a/mani_skill2/envs/ms2/mpm/shader/point/gbuffer2.frag b/mani_skill2/envs/ms2/mpm/shader/point/gbuffer2.frag deleted file mode 100644 index 263e52568..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/point/gbuffer2.frag +++ /dev/null @@ -1,38 +0,0 @@ -#version 450 -#extension GL_ARB_separate_shader_objects : enable -#extension GL_ARB_shading_language_420pack : enable - -layout(set = 2, binding = 0) uniform MaterialBuffer { - vec4 emission; - vec4 baseColor; - float fresnel; - float roughness; - float metallic; - float transmission; - float ior; - float transmissionRoughness; - int textureMask; - int padding1; -} materialBuffer; - -layout(set = 2, binding = 1) uniform sampler2D colorTexture; -layout(set = 2, binding = 2) uniform sampler2D roughnessTexture; -layout(set = 2, binding = 3) uniform sampler2D normalTexture; -layout(set = 2, binding = 4) uniform sampler2D metallicTexture; - -layout(location = 0) in vec4 inPosition; -layout(location = 1) in vec2 inUV; - -layout(location = 0) out vec4 outAlbedo2; - -void main() { - if ((materialBuffer.textureMask & 1) != 0) { - outAlbedo2 = texture(colorTexture, inUV); - outAlbedo2.rgb = outAlbedo2.rgb; // sRGB to linear - } else { - outAlbedo2 = materialBuffer.baseColor; - } - if (outAlbedo2.a == 0) { - discard; - } -} diff --git a/mani_skill2/envs/ms2/mpm/shader/point/gbuffer2.vert b/mani_skill2/envs/ms2/mpm/shader/point/gbuffer2.vert deleted file mode 100644 index 14d362118..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/point/gbuffer2.vert +++ /dev/null @@ -1,38 +0,0 @@ -#version 450 -#extension GL_ARB_separate_shader_objects : enable -#extension GL_ARB_shading_language_420pack : enable - -layout(set = 0, binding = 0) uniform CameraBuffer { - mat4 viewMatrix; - mat4 projectionMatrix; - mat4 viewMatrixInverse; - mat4 projectionMatrixInverse; - mat4 prevViewMatrix; - mat4 prevViewMatrixInverse; - float width; - float height; -} cameraBuffer; - -layout(set = 1, binding = 0) uniform ObjectBuffer { - mat4 modelMatrix; - mat4 prevModelMatrix; - uvec4 segmentation; - float transparency; - int shadeFlat; -} objectBuffer; - -layout(location = 0) in vec3 position; -layout(location = 1) in vec3 normal; -layout(location = 2) in vec2 uv; -layout(location = 3) in vec3 tangent; -layout(location = 4) in vec3 bitangent; - -layout(location = 0) out vec4 outPosition; -layout(location = 1) out vec2 outUV; - -void main() { - mat4 modelView = cameraBuffer.viewMatrix * objectBuffer.modelMatrix; - outPosition = modelView * vec4(position, 1); - outUV = uv; - gl_Position = cameraBuffer.projectionMatrix * outPosition; -} diff --git a/mani_skill2/envs/ms2/mpm/shader/point/line.frag b/mani_skill2/envs/ms2/mpm/shader/point/line.frag deleted file mode 100644 index d1516fcc1..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/point/line.frag +++ /dev/null @@ -1,13 +0,0 @@ -#version 450 -#extension GL_ARB_separate_shader_objects : enable -#extension GL_ARB_shading_language_420pack : enable - -layout(location = 0) in vec4 inPosition; -layout(location = 1) in vec4 inPrevPosition; -layout(location = 2) in vec4 inColor; - -layout(location = 0) out vec4 outLine; - -void main() { - outLine = inColor; -} diff --git a/mani_skill2/envs/ms2/mpm/shader/point/line.vert b/mani_skill2/envs/ms2/mpm/shader/point/line.vert deleted file mode 100644 index f78568b07..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/point/line.vert +++ /dev/null @@ -1,38 +0,0 @@ -#version 450 -#extension GL_ARB_separate_shader_objects : enable -#extension GL_ARB_shading_language_420pack : enable - -layout(set = 0, binding = 0) uniform CameraBuffer { - mat4 viewMatrix; - mat4 projectionMatrix; - mat4 viewMatrixInverse; - mat4 projectionMatrixInverse; - mat4 prevViewMatrix; - mat4 prevViewMatrixInverse; - float width; - float height; -} cameraBuffer; - -layout(set = 1, binding = 0) uniform ObjectBuffer { - mat4 modelMatrix; - mat4 prevModelMatrix; - uvec4 segmentation; - float transparency; - int shadeFlat; -} objectBuffer; - -layout(location = 0) in vec3 position; -layout(location = 1) in float scale; -layout(location = 2) in vec4 color; - -layout(location = 0) out vec4 outPosition; -layout(location = 1) out vec4 outPrevPosition; -layout(location = 2) out vec4 outColor; - -void main() { - mat4 modelView = cameraBuffer.viewMatrix * objectBuffer.modelMatrix; - outPosition = modelView * vec4(position, 1); - outPrevPosition = cameraBuffer.prevViewMatrix * objectBuffer.prevModelMatrix * vec4(position, 1); - gl_Position = cameraBuffer.projectionMatrix * outPosition; - outColor = color; -} diff --git a/mani_skill2/envs/ms2/mpm/shader/point/point.frag b/mani_skill2/envs/ms2/mpm/shader/point/point.frag deleted file mode 100644 index e8c1ccea2..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/point/point.frag +++ /dev/null @@ -1,57 +0,0 @@ -#version 450 -#extension GL_ARB_separate_shader_objects : enable -#extension GL_ARB_shading_language_420pack : enable - -layout (constant_id = 0) const int NUM_DIRECTIONAL_LIGHTS = 3; -layout (constant_id = 1) const int NUM_POINT_LIGHTS = 10; -layout (constant_id = 2) const int NUM_SPOT_LIGHTS = 10; -layout (constant_id = 3) const float RESOLUTION_SCALE = 1.0; - -layout(location = 0) in vec4 inColor; -layout(location = 1) in flat vec4 inCenterRadius; - -layout(location = 0) out vec4 outPointColor; -layout(location = 1) out vec4 outPointDepthLinear; - -layout(set = 0, binding = 0) uniform CameraBuffer { - mat4 viewMatrix; - mat4 projectionMatrix; - mat4 viewMatrixInverse; - mat4 projectionMatrixInverse; - mat4 prevViewMatrix; - mat4 prevViewMatrixInverse; - float width; - float height; -} cameraBuffer; - -void main() { - vec3 center = inCenterRadius.xyz; - float radius = inCenterRadius.w; - - vec2 res = vec2(cameraBuffer.width, cameraBuffer.height) * RESOLUTION_SCALE; - vec2 pixelNdc = gl_FragCoord.xy / res * 2.0 - 1.0; - - vec4 dir4 = cameraBuffer.projectionMatrixInverse * vec4(pixelNdc, 1.0, 1.0); - vec3 dir = normalize(dir4.xyz / dir4.w); - - float a = dot(dir, center); - float N = a * a - (dot(center, center) - radius * radius); - - if (N < 0) { - discard; - } - - float d = dot(dir, center) - sqrt(N); - - if (d < 0) { - discard; - } - - vec4 csPosition = vec4(dir * d, 1.0); - vec4 ndc = cameraBuffer.projectionMatrix * csPosition; - ndc /= ndc.w; - gl_FragDepth = ndc.z; - - outPointColor = inColor; - outPointDepthLinear.x = csPosition.z; -} diff --git a/mani_skill2/envs/ms2/mpm/shader/point/point.vert b/mani_skill2/envs/ms2/mpm/shader/point/point.vert deleted file mode 100644 index 4e7ed3731..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/point/point.vert +++ /dev/null @@ -1,76 +0,0 @@ -#version 450 -#extension GL_ARB_separate_shader_objects : enable -#extension GL_ARB_shading_language_420pack : enable - -layout(set = 0, binding = 0) uniform CameraBuffer { - mat4 viewMatrix; - mat4 projectionMatrix; - mat4 viewMatrixInverse; - mat4 projectionMatrixInverse; - mat4 prevViewMatrix; - mat4 prevViewMatrixInverse; - float width; - float height; -} cameraBuffer; - -layout(set = 1, binding = 0) uniform ObjectBuffer { - mat4 modelMatrix; - mat4 prevModelMatrix; - uvec4 segmentation; - float transparency; - int shadeFlat; -} objectBuffer; - -layout(location = 0) in vec3 position; -layout(location = 1) in float scale; -layout(location = 2) in vec4 color; - -layout(location = 0) out vec4 outColor; -layout(location = 1) out vec4 outCenterRadius; - -void main() { - mat4 modelView = cameraBuffer.viewMatrix * objectBuffer.modelMatrix; - vec4 outPosition = modelView * vec4(position, 1); - float radius = scale; - - vec3 v0 = outPosition.xyz + vec3(-radius, -radius, -radius); - vec3 v1 = outPosition.xyz + vec3(-radius, -radius, radius); - vec3 v2 = outPosition.xyz + vec3(-radius, radius, -radius); - vec3 v3 = outPosition.xyz + vec3(-radius, radius, radius); - vec3 v4 = outPosition.xyz + vec3(radius, -radius, -radius); - vec3 v5 = outPosition.xyz + vec3(radius, -radius, radius); - vec3 v6 = outPosition.xyz + vec3(radius, radius, -radius); - vec3 v7 = outPosition.xyz + vec3(radius, radius, radius); - - vec4 p0 = cameraBuffer.projectionMatrix * vec4(v0, 1.0); - vec4 p1 = cameraBuffer.projectionMatrix * vec4(v1, 1.0); - vec4 p2 = cameraBuffer.projectionMatrix * vec4(v2, 1.0); - vec4 p3 = cameraBuffer.projectionMatrix * vec4(v3, 1.0); - vec4 p4 = cameraBuffer.projectionMatrix * vec4(v4, 1.0); - vec4 p5 = cameraBuffer.projectionMatrix * vec4(v5, 1.0); - vec4 p6 = cameraBuffer.projectionMatrix * vec4(v6, 1.0); - vec4 p7 = cameraBuffer.projectionMatrix * vec4(v7, 1.0); - - vec2 s0 = p0.xy / p0.w; - vec2 s1 = p1.xy / p1.w; - vec2 s2 = p2.xy / p2.w; - vec2 s3 = p3.xy / p3.w; - vec2 s4 = p4.xy / p4.w; - vec2 s5 = p5.xy / p5.w; - vec2 s6 = p6.xy / p6.w; - vec2 s7 = p7.xy / p7.w; - - vec2 lower = min(min(min(s0, s1), min(s2, s3)), min(min(s4, s5), min(s6, s7))); - vec2 upper = max(max(max(s0, s1), max(s2, s3)), max(max(s4, s5), max(s6, s7))); - vec2 extent_ndc = upper - lower; - vec2 center_ndc = (lower + upper) / 2.0; - - vec2 extent = extent_ndc * vec2(cameraBuffer.width, cameraBuffer.height) / 2.0; - float size = max(extent.x, extent.y); - - gl_PointSize = size; - gl_Position = vec4(center_ndc, 1.0, 1.0); - outCenterRadius = vec4(outPosition.xyz, radius); - - outColor = color; -} diff --git a/mani_skill2/envs/ms2/mpm/shader/point/shadow.vert b/mani_skill2/envs/ms2/mpm/shader/point/shadow.vert deleted file mode 100644 index c5d1dcc24..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/point/shadow.vert +++ /dev/null @@ -1,30 +0,0 @@ -#version 450 -#extension GL_ARB_separate_shader_objects : enable -#extension GL_ARB_shading_language_420pack : enable - -layout(set = 0, binding = 0) uniform LightBuffer { - mat4 viewMatrix; - mat4 viewMatrixInverse; - mat4 projectionMatrix; - mat4 projectionMatrixInverse; - int width; - int height; -} lightBuffer; - -layout(set = 1, binding = 0) uniform ObjectBuffer { - mat4 modelMatrix; - mat4 prevModelMatrix; - uvec4 segmentation; - float transparency; - int shadeFlat; -} objectBuffer; - -layout(location = 0) in vec3 position; -layout(location = 1) in vec3 normal; -layout(location = 2) in vec2 uv; -layout(location = 3) in vec3 tangent; -layout(location = 4) in vec3 bitangent; - -void main() { - gl_Position = lightBuffer.projectionMatrix * lightBuffer.viewMatrix * objectBuffer.modelMatrix * vec4(position, 1.f); -} diff --git a/mani_skill2/envs/ms2/mpm/shader/point/shadow_point.frag b/mani_skill2/envs/ms2/mpm/shader/point/shadow_point.frag deleted file mode 100644 index b84fa3224..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/point/shadow_point.frag +++ /dev/null @@ -1,39 +0,0 @@ -#version 450 -#extension GL_ARB_separate_shader_objects : enable -#extension GL_ARB_shading_language_420pack : enable - -layout(location = 0) in vec4 inPosition; -layout(location = 1) in flat vec4 inNdcRadius; - -layout(set = 0, binding = 0) uniform LightBuffer { - mat4 viewMatrix; - mat4 viewMatrixInverse; - mat4 projectionMatrix; - mat4 projectionMatrixInverse; - int width; - int height; -} lightBuffer; - -layout(set = 1, binding = 0) uniform ObjectBuffer { - mat4 modelMatrix; - mat4 prevModelMatrix; - uvec4 segmentation; - float transparency; - int shadeFlat; -} objectBuffer; - -void main() { - if (gl_PointCoord.s * gl_PointCoord.s + gl_PointCoord.t * gl_PointCoord.t > 1) { - discard; - } - // vec2 centerNdc = inNdcRadius.xy; - // vec2 res = vec2(lightBuffer.width, lightBuffer.height) * RESOLUTION_SCALE; - // vec2 pixelNdc = gl_FragCoord.xy / res * 2.0 - 1.0; - // vec2 offsetNdc = pixelNdc - centerNdc; - // vec2 offset = offsetNdc * (-inPosition.z) / vec2(lightBuffer.projectionMatrix[0][0], lightBuffer.projectionMatrix[1][1]); - // float radius = inNdcRadius.w; - // offset /= radius; - // if (offset.x * offset.x + offset.y * offset.y > 1) { - // discard; - // } -} diff --git a/mani_skill2/envs/ms2/mpm/shader/point/shadow_point.vert b/mani_skill2/envs/ms2/mpm/shader/point/shadow_point.vert deleted file mode 100644 index b86fd2b3a..000000000 --- a/mani_skill2/envs/ms2/mpm/shader/point/shadow_point.vert +++ /dev/null @@ -1,39 +0,0 @@ -#version 450 -#extension GL_ARB_separate_shader_objects : enable -#extension GL_ARB_shading_language_420pack : enable - -layout(set = 0, binding = 0) uniform LightBuffer { - mat4 viewMatrix; - mat4 viewMatrixInverse; - mat4 projectionMatrix; - mat4 projectionMatrixInverse; - int width; - int height; -} lightBuffer; - -layout(set = 1, binding = 0) uniform ObjectBuffer { - mat4 modelMatrix; - mat4 prevModelMatrix; - uvec4 segmentation; - float transparency; - int shadeFlat; -} objectBuffer; - -layout(location = 0) in vec3 position; -layout(location = 1) in float scale; -layout(location = 2) in vec4 color; - -layout(location = 0) out vec4 outPosition; -layout(location = 1) out flat vec4 outNdcRadius; - -void main() { - mat4 modelView = lightBuffer.viewMatrix * objectBuffer.modelMatrix; - outPosition = modelView * vec4(position, 1); - - float radius = scale; - - gl_PointSize = lightBuffer.projectionMatrix[0][0] * lightBuffer.width * radius; - - gl_Position = lightBuffer.projectionMatrix * outPosition; - outNdcRadius = vec4(gl_Position.xyz / gl_Position.w, radius); -} diff --git a/mani_skill2/envs/ms2/mpm/utils.py b/mani_skill2/envs/ms2/mpm/utils.py deleted file mode 100644 index fc6240161..000000000 --- a/mani_skill2/envs/ms2/mpm/utils.py +++ /dev/null @@ -1,146 +0,0 @@ -from collections import OrderedDict -from typing import Union - -import h5py -import numpy as np -import sapien -import trimesh - - -def actor2meshes(actor: sapien.Entity, visual=False, return_primitives=False): - primitives = [] - meshes = [] - if visual: - bodies = actor.get_visual_bodies() - for body in bodies: - if body.type == "mesh": - for shape in body.get_render_shapes(): - mesh = trimesh.Trimesh( - shape.mesh.vertices * body.scale, - shape.mesh.indices.reshape((-1, 3)), - ) - mesh.apply_transform(body.local_pose.to_transformation_matrix()) - meshes.append(mesh) - elif body.type == "box": - half_lengths = body.half_lengths - if return_primitives: - primitives.append(("box", half_lengths, body.local_pose)) - else: - mesh = trimesh.creation.box(half_lengths * 2) - meshes.append(mesh) - elif body.type == "capsule": - radius = body.radius - half_length = body.half_length - if return_primitives: - primitives.append( - ("capsule", [radius, half_length], body.local_pose) - ) - else: - raise NotImplementedError() - else: - raise NotImplementedError() - else: - collisions = actor.get_collision_shapes() - meshes = [] - for s in collisions: - s: sapien.CollisionShape - g: sapien.CollisionGeometry = s.geometry - - if isinstance(g, sapien.ConvexMeshGeometry): - mesh = trimesh.Trimesh(g.vertices * g.scale, g.indices.reshape(-1, 3)) - assert np.allclose(g.rotation, np.array([1, 0, 0, 0])) - mesh.apply_transform(s.get_local_pose().to_transformation_matrix()) - meshes.append(mesh) - elif isinstance(g, sapien.BoxGeometry): - half_lengths = g.half_lengths - if return_primitives: - primitives.append(("box", half_lengths, s.get_local_pose())) - else: - mesh = trimesh.creation.box(half_lengths * 2) - meshes.append(mesh) - elif isinstance(g, sapien.CapsuleGeometry): - radius = g.radius - half_length = g.half_length - if return_primitives: - primitives.append( - ("capsule", [radius, half_length], s.get_local_pose()) - ) - else: - raise NotImplementedError() - else: - raise NotImplementedError() - - if return_primitives: - return meshes, primitives - return meshes - - -def trimesh2sdf(meshes, margin, dx, bbox=None): - if meshes is None: - return None - mesh = trimesh.util.concatenate(meshes) - - if bbox is None: - bbox = mesh.bounds.copy() - - sdfs = [] - normals = [] - for mesh in meshes: - center = (bbox[0] + bbox[1]) / 2 - res = np.ceil((bbox[1] - bbox[0] + margin * 2) / dx).astype(int) - lower = center - res * dx / 2.0 - - points = np.zeros((res[0], res[1], res[2], 3)) - x = np.arange(0.5, res[0]) * dx + lower[0] - y = np.arange(0.5, res[1]) * dx + lower[1] - z = np.arange(0.5, res[2]) * dx + lower[2] - - points[..., 0] += x[:, None, None] - points[..., 1] += y[None, :, None] - points[..., 2] += z[None, None, :] - - points = points.reshape((-1, 3)) - - query = trimesh.proximity.ProximityQuery(mesh) - sdf = query.signed_distance(points) * -1.0 - - surface_points, _, tri_id = query.on_surface(points) - face_normal = mesh.face_normals[tri_id] - normal = (points - surface_points) * np.sign(sdf)[..., None] - length = np.linalg.norm(normal, axis=-1) - mask = length < 1e6 - normal[mask] = face_normal[mask] - normal = normal / (np.linalg.norm(normal, axis=-1, keepdims=True) + 1e-8) - sdf = sdf.reshape(res) - normal = normal.reshape((res[0], res[1], res[2], 3)) - - sdfs.append(sdf) - normals.append(normal) - - if len(sdfs) == 1: - sdf = sdfs[0] - normal = normals[0] - else: - sdfs = np.stack(sdfs) - normals = np.stack(normals) - index = np.expand_dims(sdfs.argmin(0), 0) - sdf = np.take_along_axis(sdfs, index, 0)[0] - normal = np.take_along_axis(normals, np.expand_dims(index, -1), 0)[0] - - return { - "sdf": sdf, - "normal": normal, - "position": lower, - "scale": np.ones(3) * dx, - "dim": res, - } - - -def load_h5_as_dict(h5file: Union[h5py.File, h5py.Group]): - out = OrderedDict() - for key in h5file.keys(): - if isinstance(h5file[key], h5py.Group): - out[key] = load_h5_as_dict(h5file[key]) - else: - out[key] = h5file[key][:] - return out diff --git a/mani_skill2/envs/ms2/mpm/write_env.py b/mani_skill2/envs/ms2/mpm/write_env.py deleted file mode 100644 index 17be03af4..000000000 --- a/mani_skill2/envs/ms2/mpm/write_env.py +++ /dev/null @@ -1,298 +0,0 @@ -from collections import OrderedDict - -import h5py -import numpy as np -import sapien -import warp as wp -from mpm.height_rasterizer import rasterize_clear_kernel, rasterize_kernel -from transforms3d.euler import euler2quat - -from mani_skill2 import ASSET_DIR -from mani_skill2.agents.robots.panda.variants import PandaStick -from mani_skill2.envs.ms2.mpm.base_env import MPMBaseEnv, MPMModelBuilder, MPMSimulator -from mani_skill2.envs.ms2.mpm.utils import load_h5_as_dict -from mani_skill2.sensors.camera import CameraConfig -from mani_skill2.utils.registration import register_env -from mani_skill2.utils.structs.pose import vectorize_pose - - -@wp.kernel -def success_iou_kernel( - goal_image: wp.array(dtype=int, ndim=2), - current_image: wp.array(dtype=int, ndim=2), - out: wp.array(dtype=int), -): - h, w = wp.tid() - a = goal_image[h, w] < 40 - an = goal_image[h, w] < 50 - b = current_image[h, w] < 40 - bn = current_image[h, w] < 50 - - # intersection - if (an and b) or (a and bn): - wp.atomic_add(out, 0, 1) - - # union - if a or b: - wp.atomic_add(out, 1, 1) - - -@register_env("Write-v0", max_episode_steps=200) -class WriteEnv(MPMBaseEnv): - def __init__( - self, - *args, - level_dir="write/levels", - **kwargs, - ): - self.level_dir = ASSET_DIR / level_dir - self.all_filepaths = sorted(self.level_dir.glob("*.h5")) - if len(self.all_filepaths) == 0: - raise RuntimeError( - "Please download required assets for Write:" - "`python -m mani_skill2.utils.download_asset write`" - ) - super().__init__(*args, **kwargs) - - def reset(self, seed=None, options=None): - if options is None: - options = dict() - self.level_file = options.pop("level_file", None) - return super().reset(seed=seed, options=options) - - def _get_obs_extra(self): - return OrderedDict( - tcp_pose=vectorize_pose(self.end_effector.pose), - goal=self.goal_image_display_numpy, - ) - - def _initialize_mpm(self): - if self.level_file is not None: - filepath = self.level_dir / self.level_file - assert filepath is not None - else: - filepath = self._episode_rng.choice(self.all_filepaths) - - self.info = load_h5_as_dict(h5py.File(str(filepath), "r")) - - n = len(self.info["goal"]) - self.image_size = 64 - self.image_world_size = 0.21 - self.goal_array = wp.array(self.info["goal"], dtype=wp.vec3, device="cuda") - self.goal_image = wp.empty( - (self.image_size, self.image_size), dtype=int, device="cuda" - ) - self.current_image = wp.empty( - (self.image_size, self.image_size), dtype=int, device="cuda" - ) - self.iou_buffer = wp.empty(2, dtype=int, device="cuda") - - wp.launch( - rasterize_clear_kernel, - dim=(self.image_size, self.image_size), - inputs=[self.goal_image, 0], - device="cuda", - ) - xy_scale = self.image_size / self.image_world_size - z_scale = 1000 - radius = int(0.007 * xy_scale) - wp.launch( - rasterize_kernel, - dim=n, - inputs=[ - self.goal_array, - wp.vec3(self.image_world_size / 2, self.image_world_size / 2, 0.0), - xy_scale, - z_scale, - radius, - self.image_size, - self.image_size, - self.goal_image, - ], - device="cuda", - ) - - self.goal_image_display_numpy = self.goal_image.numpy()[:, ::-1] - self.goal_image_display_numpy = np.clip( - self.goal_image_display_numpy, 0, 255 - ).astype(np.uint8) - - self._episode_rng = np.random.RandomState(self._episode_seed) - - self.model_builder.clear_particles() - E = 3e5 - nu = 0.1 - mu, lam = E / (2 * (1 + nu)), E * nu / ((1 + nu) * (1 - 2 * nu)) - type = 0 - ys = 2e3 - - height_map = np.ones((42, 42), dtype=np.float32) * 0.05 - self.model_builder.add_mpm_from_height_map( - pos=(0.0, 0.0, 0.0), - vel=(0.0, 0.0, 0.0), - dx=0.005, - height_map=height_map, - density=3.0e3, - mu_lambda_ys=(mu, lam, ys), - friction_cohesion=(0.0, 0.0, 0.0), - type=type, - jitter=True, - color=(0.65237011, 0.14198029, 0.02201299), - random_state=self._episode_rng, - ) - - self.model_builder.init_model_state(self.mpm_model, self.mpm_states) - self.mpm_model.struct.static_ke = 100.0 - self.mpm_model.struct.static_kd = 0.0 - self.mpm_model.struct.static_mu = 1.0 - self.mpm_model.struct.static_ka = 0.0 - - self.mpm_model.struct.body_ke = 100.0 - self.mpm_model.struct.body_kd = 0.0 - self.mpm_model.struct.body_mu = 1.0 - self.mpm_model.struct.body_ka = 0.0 - - self.mpm_model.adaptive_grid = False - self.mpm_model.grid_contact = False - self.mpm_model.particle_contact = True - self.mpm_model.struct.body_sticky = 0 - self.mpm_model.struct.ground_sticky = 1 - - self.mpm_model.struct.particle_radius = 0.005 - - def _setup_mpm(self): - self.model_builder = MPMModelBuilder() - self.model_builder.set_mpm_domain(domain_size=[0.5, 0.5, 0.5], grid_length=0.01) - self.model_builder.reserve_mpm_particles(count=self.max_particles) - self._setup_mpm_bodies() - self.mpm_simulator = MPMSimulator(device="cuda") - self.mpm_model = self.model_builder.finalize(device="cuda") - self.mpm_model.gravity = np.array((0.0, 0.0, -9.81), dtype=np.float32) - self.mpm_model.struct.ground_normal = wp.vec3(0.0, 0.0, 1.0) - self.mpm_model.struct.particle_radius = 0.005 - self.mpm_states = [ - self.mpm_model.state() for _ in range(self._mpm_step_per_sapien_step + 1) - ] - - def _register_sensors(self): - p, q = [-0.2, 0, 0.3], euler2quat(0, np.pi / 6, 0) - return CameraConfig("base_camera", p, q, 128, 128, np.pi / 2, 0.001, 10) - - def _register_render_cameras(self): - p, q = [-0.3, 0, 0.4], euler2quat(0, np.pi / 5, 0) - return CameraConfig("render_camera", p, q, 512, 512, 1, 0.001, 10) - - def _load_agent(self): - self.agent = PandaStick( - self._scene, - self._control_freq, - control_mode=self._control_mode, - ) - self.end_effector = self.agent.robot.get_links()[-1] - - def _initialize_agent(self): - noise = self._episode_rng.uniform([-0.1] * 7, [0.1] * 7) - # fmt: off - qpos = np.array([-0.029177314, 0.10816099, 0.03054934, -2.1639752, -0.0013982388, 2.2785723, 0.79039097]) + noise - # fmt: on - - self.agent.reset(qpos) - self.agent.robot.set_pose(sapien.Pose([-0.55, 0, 0])) - - def _load_actors(self): - super()._load_actors() - b = self._scene.create_actor_builder() - b.add_box_collision(half_size=[0.15, 0.02, 0.04]) - b.add_box_visual(half_size=[0.15, 0.02, 0.04]) - w0 = b.build_kinematic("wall") - w1 = b.build_kinematic("wall") - w2 = b.build_kinematic("wall") - w3 = b.build_kinematic("wall") - - w0.set_pose(sapien.Pose([0, -0.13, 0.04])) - w1.set_pose(sapien.Pose([0, 0.13, 0.04])) - w2.set_pose(sapien.Pose([-0.13, 0, 0.04], [0.7071068, 0, 0, 0.7071068])) - w3.set_pose(sapien.Pose([0.13, 0, 0.04], [0.7071068, 0, 0, 0.7071068])) - self.walls = [w0, w1, w2, w3] - - def _get_coupling_actors(self): - return [ - l for l in self.agent.robot.get_links() if l.name == "panda_hand" - ] + self.walls - - def step(self, *args, **kwargs): - self._iou = None - return super().step(*args, **kwargs) - - def _compute_iou(self): - if self._iou is not None: - return self._iou - - n = len(self.info["goal"]) - self.iou_buffer.zero_() - - wp.launch( - rasterize_clear_kernel, - dim=(self.image_size, self.image_size), - inputs=[self.current_image, 0], - device="cuda", - ) - xy_scale = self.image_size / self.image_world_size - z_scale = 1000 - radius = int(0.007 * xy_scale) - wp.launch( - rasterize_kernel, - dim=n, - inputs=[ - self.mpm_states[0].struct.particle_q, - wp.vec3(self.image_world_size / 2, self.image_world_size / 2, 0.0), - xy_scale, - z_scale, - radius, - self.image_size, - self.image_size, - self.current_image, - ], - device="cuda", - ) - - wp.launch( - success_iou_kernel, - dim=(self.image_size, self.image_size), - inputs=[self.goal_image, self.current_image, self.iou_buffer], - device="cuda", - ) - inter, union = self.iou_buffer.numpy() - self._iou = inter / union - return self._iou - - def evaluate(self, **kwargs): - return OrderedDict(success=self._compute_iou() > 0.8, iou=self._compute_iou()) - - def compute_dense_reward(self, reward_info=False, **kwargs): - # iou (goal): [0, 1] - iou = self._compute_iou() - - # reaching reward - gripper_mat = self.end_effector.get_pose().to_transformation_matrix() - gripper_bottom_mat = np.array( - [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0.02], [0, 0, 0, 1]] - ) - bottom_mat = gripper_mat @ gripper_bottom_mat - bottom_pos = sapien.Pose.from_transformation_matrix(bottom_mat).p - particles_x = self.get_mpm_state()["x"] - distance = np.min(np.linalg.norm(particles_x - bottom_pos, axis=-1)) - reaching_reward = 1 - np.tanh(10.0 * distance) - - # vertical reward - gripper_mat = self.end_effector.pose.to_transformation_matrix() - z = gripper_mat[:3, 2] - angle = np.arcsin(np.clip(np.linalg.norm(np.cross(z, [0, 0, -1])), -1, 1)) - reward_orientation = 1 - angle - - return iou + 0.1 * reaching_reward + 0.1 * reward_orientation - - def compute_normalized_dense_reward(self, **kwargs): - return self.compute_dense_reward( - **kwargs - ) # no normalization for now since original reward scale is already low diff --git a/mani_skill2/envs/sapien_env.py b/mani_skill2/envs/sapien_env.py index 33bbc97c6..7c9f112f7 100644 --- a/mani_skill2/envs/sapien_env.py +++ b/mani_skill2/envs/sapien_env.py @@ -17,9 +17,9 @@ from sapien.utils import Viewer from mani_skill2 import logger +from mani_skill2.agents import REGISTERED_AGENTS from mani_skill2.agents.base_agent import BaseAgent from mani_skill2.agents.multi_agent import MultiAgent -from mani_skill2.agents.robots import ROBOTS from mani_skill2.envs.scene import ManiSkillScene from mani_skill2.envs.utils.observations.observations import ( sensor_data_to_pointcloud, @@ -288,7 +288,11 @@ def _load_agent(self): agent_cls = robot_uid # robot_uids = self._agent_cls.uid else: - agent_cls = ROBOTS[robot_uid] + if robot_uid not in REGISTERED_AGENTS: + raise RuntimeError( + f"Agent {robot_uid} not found in the dict of registered agents. If the id is not a typo then make sure to apply the @register_agent() decorator." + ) + agent_cls = REGISTERED_AGENTS[robot_uid].agent_cls agent: BaseAgent = agent_cls( self._scene, self._control_freq, diff --git a/mani_skill2/envs/scenes/__init__.py b/mani_skill2/envs/scenes/__init__.py index 1c69ac41f..d7df6be93 100644 --- a/mani_skill2/envs/scenes/__init__.py +++ b/mani_skill2/envs/scenes/__init__.py @@ -13,10 +13,11 @@ scene_builders = { "ReplicaCAD": ReplicaCADSceneBuilder, - "ArchitecTHOR": ArchitecTHORSceneBuilder, - "ProcTHOR": ProcTHORSceneBuilder, - "RoboTHOR": RoboTHORSceneBuilder, - "iTHOR": iTHORSceneBuilder, + # TODO these will be added in final release later once articulations are fixed + # "ArchitecTHOR": ArchitecTHORSceneBuilder, + # "ProcTHOR": ProcTHORSceneBuilder, + # "RoboTHOR": RoboTHORSceneBuilder, + # "iTHOR": iTHORSceneBuilder, } # Register environments just for benchmarking/exploration and to be creatable by just ID, these don't have any specific tasks designed in them. diff --git a/mani_skill2/envs/tasks/fmb/fmb.py b/mani_skill2/envs/tasks/fmb/fmb.py index f1df0fbac..4443360f3 100644 --- a/mani_skill2/envs/tasks/fmb/fmb.py +++ b/mani_skill2/envs/tasks/fmb/fmb.py @@ -7,9 +7,6 @@ import torch from transforms3d.euler import euler2quat -from mani_skill2.agents.robots import ( # a dictionary mapping robot name to robot class that inherits BaseAgent - ROBOTS, -) from mani_skill2.agents.robots.panda.panda import Panda from mani_skill2.envs.sapien_env import BaseEnv from mani_skill2.sensors.camera import CameraConfig diff --git a/mani_skill2/envs/tasks/open_cabinet_drawer.py b/mani_skill2/envs/tasks/open_cabinet_drawer.py index e98211846..057eca355 100644 --- a/mani_skill2/envs/tasks/open_cabinet_drawer.py +++ b/mani_skill2/envs/tasks/open_cabinet_drawer.py @@ -22,7 +22,8 @@ from mani_skill2.utils.structs.pose import Pose -@register_env("OpenCabinetDrawer-v1", max_episode_steps=100) +# TODO (stao): we need to cut the meshes of all the cabinets in this dataset for gpu sim, not registering task for now +# @register_env("OpenCabinetDrawer-v1", max_episode_steps=100) class OpenCabinetDrawerEnv(BaseEnv): """ Task Description diff --git a/mani_skill2/utils/scene_builder/ai2thor/scene_builder.py b/mani_skill2/utils/scene_builder/ai2thor/scene_builder.py index 5bc028203..2853cb797 100644 --- a/mani_skill2/utils/scene_builder/ai2thor/scene_builder.py +++ b/mani_skill2/utils/scene_builder/ai2thor/scene_builder.py @@ -115,7 +115,7 @@ def build( if self.scene_dataset == "ProcTHOR": # for some reason the scene needs to rotate around y-axis by 90 degrees for ProcTHOR scenes from hssd dataset bg_q = transforms3d.quaternions.qmult( - q, + bg_q, transforms3d.quaternions.axangle2quat( np.array([0, -1, 0]), theta=np.deg2rad(90) ), diff --git a/tests/test_gpu_envs.py b/tests/test_gpu_envs.py index f34c14066..0564e578d 100644 --- a/tests/test_gpu_envs.py +++ b/tests/test_gpu_envs.py @@ -23,6 +23,23 @@ @pytest.mark.gpu_sim @pytest.mark.parametrize("env_id", ENV_IDS) +def test_all_envs(env_id): + env = gym.make_vec( + env_id, + num_envs=16, + vectorization_mode="custom", + vector_kwargs=dict(obs_mode="state", sim_cfg=LOW_MEM_SIM_CFG), + ) + obs, _ = env.reset() + action_space = env.action_space + for _ in range(5): + obs, rew, terminated, truncated, info = env.step(action_space.sample()) + env.close() + del env + + +@pytest.mark.gpu_sim +@pytest.mark.parametrize("env_id", STATIONARY_ENV_IDS) @pytest.mark.parametrize("obs_mode", OBS_MODES) def test_envs_obs_modes(env_id, obs_mode): def assert_device(x): @@ -169,7 +186,7 @@ def test_raw_sim_states(): @pytest.mark.gpu_sim @pytest.mark.parametrize("env_id", ENV_IDS) -@pytest.mark.parametrize("robot_uids", ROBOTS) +@pytest.mark.parametrize("robot_uids", STATIONARY_ENV_IDS) def test_robots(env_id, robot_uids): if env_id in [ "PandaAvoidObstacles-v0", diff --git a/tests/test_wrappers.py b/tests/test_wrappers.py index 50a465305..6b7a1af93 100644 --- a/tests/test_wrappers.py +++ b/tests/test_wrappers.py @@ -9,11 +9,17 @@ ) from mani_skill2.utils.wrappers.visual_encoders import VisualEncoderWrapper from mani_skill2.vector.wrappers.gymnasium import ManiSkillVectorEnv -from tests.utils import ENV_IDS, LOW_MEM_SIM_CFG, MULTI_AGENT_ENV_IDS, OBS_MODES +from tests.utils import ( + ENV_IDS, + LOW_MEM_SIM_CFG, + MULTI_AGENT_ENV_IDS, + OBS_MODES, + STATIONARY_ENV_IDS, +) @pytest.mark.gpu_sim -@pytest.mark.parametrize("env_id", ENV_IDS) +@pytest.mark.parametrize("env_id", STATIONARY_ENV_IDS) @pytest.mark.parametrize("obs_mode", OBS_MODES) def test_recordepisode_wrapper_gpu(env_id, obs_mode): env = gym.make( @@ -68,7 +74,7 @@ def test_recordepisode_wrapper(env_id, obs_mode): @pytest.mark.gpu_sim -@pytest.mark.parametrize("env_id", ENV_IDS[:1]) +@pytest.mark.parametrize("env_id", STATIONARY_ENV_IDS[:1]) @pytest.mark.parametrize("obs_mode", OBS_MODES[:1]) def test_recordepisode_wrapper_gpu_render_sensor(env_id, obs_mode): env = gym.make( @@ -100,7 +106,7 @@ def test_recordepisode_wrapper_gpu_render_sensor(env_id, obs_mode): del env -@pytest.mark.parametrize("env_id", ENV_IDS) +@pytest.mark.parametrize("env_id", STATIONARY_ENV_IDS) @pytest.mark.parametrize("obs_mode", OBS_MODES) def test_recordepisode_wrapper_render_sensor(env_id, obs_mode): env = gym.make( @@ -123,7 +129,7 @@ def test_recordepisode_wrapper_render_sensor(env_id, obs_mode): @pytest.mark.gpu_sim -@pytest.mark.parametrize("env_id", ENV_IDS[:1]) +@pytest.mark.parametrize("env_id", STATIONARY_ENV_IDS[:1]) @pytest.mark.parametrize("obs_mode", OBS_MODES[:1]) def test_recordepisode_wrapper_partial_reset_gpu(env_id, obs_mode): env = gym.make( @@ -156,7 +162,7 @@ def test_recordepisode_wrapper_partial_reset_gpu(env_id, obs_mode): del env -@pytest.mark.parametrize("env_id", ENV_IDS[:1]) +@pytest.mark.parametrize("env_id", STATIONARY_ENV_IDS[:1]) @pytest.mark.parametrize("obs_mode", OBS_MODES[:1]) def test_recordepisode_wrapper_partial_reset(env_id, obs_mode): env = gym.make( @@ -189,7 +195,7 @@ def test_recordepisode_wrapper_partial_reset(env_id, obs_mode): @pytest.mark.gpu_sim -@pytest.mark.parametrize("env_id", [ENV_IDS[0]]) +@pytest.mark.parametrize("env_id", STATIONARY_ENV_IDS[:1]) def test_visualencoders_gpu(env_id): env = gym.make( env_id, @@ -225,7 +231,7 @@ def test_visualencoders_gpu(env_id): @pytest.mark.gpu_sim -@pytest.mark.parametrize("env_id", [ENV_IDS[0]]) +@pytest.mark.parametrize("env_id", STATIONARY_ENV_IDS[:1]) def test_visualencoder_flatten_gpu(env_id): env = gym.make( env_id, diff --git a/tests/utils.py b/tests/utils.py index 6b3600815..0ce04a611 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -4,15 +4,11 @@ import torch from mani_skill2.utils.common import flatten_dict_keys +from mani_skill2.utils.registration import REGISTERED_ENVS from mani_skill2.utils.sapien_utils import to_numpy # TODO (stao): reactivate old tasks once fixed -ENV_IDS = [ - # "LiftCube-v0", - "PickCube-v1", - "StackCube-v1", - "PickSingleYCB-v1", -] +ENV_IDS = list(REGISTERED_ENVS.keys()) MULTI_AGENT_ENV_IDS = ["TwoRobotStackCube-v1", "TwoRobotPickCube-v1"] STATIONARY_ENV_IDS = [