new task: in hand rotating in_hand_task

mani_skill2/agents/robots/__init__.py

@@ -5,6 +5,7 @@
 from .panda import Panda
 from .xarm import XArm7Ability
 from .xmate3 import Xmate3Robotiq
+from .allegro_hand import AllegroHandRightTouch, AllegroHandRight, AllegroHandLeft
 
 ROBOTS = {
     "panda": Panda,
@@ -15,6 +16,9 @@
     # Dexterous Hand
     "dclaw": DClaw,
     "xarm7_ability": XArm7Ability,
+    "allegro_hand_right": AllegroHandRight,
+    "allegro_hand_left": AllegroHandLeft,
+    "allegro_hand_right_touch": AllegroHandRightTouch,
     # Locomotion
     "anymal-c": ANYmalC,
 }

mani_skill2/agents/robots/allegro_hand/__init__.py

@@ -0,0 +1,2 @@
+from .allegro import AllegroHandLeft, AllegroHandRight
+from .allegro_touch import AllegroHandRightTouch

mani_skill2/agents/robots/allegro_hand/allegro.py

@@ -0,0 +1,151 @@
+from copy import deepcopy
+from typing import List
+
+import sapien
+import torch
+
+from mani_skill2 import PACKAGE_ASSET_DIR
+from mani_skill2.agents.base_agent import BaseAgent
+from mani_skill2.agents.controllers import *
+from mani_skill2.utils.sapien_utils import (
+    get_obj_by_name,
+)
+from mani_skill2.utils.sapien_utils import get_objs_by_names
+from mani_skill2.utils.structs.pose import vectorize_pose
+
+
+class AllegroHandRight(BaseAgent):
+    uid = "allegro_hand_right"
+    urdf_path = f"{PACKAGE_ASSET_DIR}/robots/allegro/allegro_hand_right_glb.urdf"
+    urdf_config = dict(
+        _materials=dict(
+            tip=dict(static_friction=2.0, dynamic_friction=1.0, restitution=0.0)
+        ),
+        link={
+            "link_3.0_tip": dict(
+                material="tip", patch_radius=0.1, min_patch_radius=0.1
+            ),
+            "link_7.0_tip": dict(
+                material="tip", patch_radius=0.1, min_patch_radius=0.1
+            ),
+            "link_11.0_tip": dict(
+                material="tip", patch_radius=0.1, min_patch_radius=0.1
+            ),
+            "link_15.0_tip": dict(
+                material="tip", patch_radius=0.1, min_patch_radius=0.1
+            ),
+        },
+    )
+    sensor_configs = {}
+
+    def __init__(self, *args, **kwargs):
+        self.joint_names = [
+            "joint_0.0",
+            "joint_1.0",
+            "joint_2.0",
+            "joint_3.0",
+            "joint_4.0",
+            "joint_5.0",
+            "joint_6.0",
+            "joint_7.0",
+            "joint_8.0",
+            "joint_9.0",
+            "joint_10.0",
+            "joint_11.0",
+            "joint_12.0",
+            "joint_13.0",
+            "joint_14.0",
+            "joint_15.0",
+        ]
+
+        self.joint_stiffness = 4e2
+        self.joint_damping = 1e1
+        self.joint_force_limit = 5e1
+
+        # Order: thumb finger, index finger, middle finger, ring finger
+        self.tip_link_names = [
+            "link_15.0_tip",
+            "link_3.0_tip",
+            "link_7.0_tip",
+            "link_11.0_tip",
+        ]
+
+        self.palm_link_name = "palm"
+        super().__init__(*args, **kwargs)
+
+    def _after_init(self):
+        self.tip_links: List[sapien.Entity] = get_objs_by_names(
+            self.robot.get_links(), self.tip_link_names
+        )
+        self.palm_link: sapien.Entity = get_obj_by_name(
+            self.robot.get_links(), self.palm_link_name
+        )
+
+    @property
+    def controller_configs(self):
+        # -------------------------------------------------------------------------- #
+        # Arm
+        # -------------------------------------------------------------------------- #
+        joint_pos = PDJointPosControllerConfig(
+            self.joint_names,
+            None,
+            None,
+            self.joint_stiffness,
+            self.joint_damping,
+            self.joint_force_limit,
+            normalize_action=False,
+        )
+        joint_delta_pos = PDJointPosControllerConfig(
+            self.joint_names,
+            -0.1,
+            0.1,
+            self.joint_stiffness,
+            self.joint_damping,
+            self.joint_force_limit,
+            use_delta=True,
+        )
+        joint_target_delta_pos = deepcopy(joint_delta_pos)
+        joint_target_delta_pos.use_target = True
+
+        controller_configs = dict(
+            pd_joint_delta_pos=joint_delta_pos,
+            pd_joint_pos=joint_pos,
+            pd_joint_target_delta_pos=joint_target_delta_pos,
+        )
+
+        # Make a deepcopy in case users modify any config
+        return deepcopy_dict(controller_configs)
+
+    def get_proprioception(self):
+        """
+        Get the proprioceptive state of the agent.
+        """
+        obs = super().get_proprioception()
+        obs.update(
+            {
+                "palm_pose": self.palm_pose,
+                "tip_poses": self.tip_poses.reshape(-1, len(self.tip_links) * 7),
+            }
+        )
+
+        return obs
+
+    @property
+    def tip_poses(self):
+        """
+        Get the tip pose for each of the finger, four fingers in total
+        """
+        tip_poses = [vectorize_pose(link.pose) for link in self.tip_links]
+        return torch.stack(tip_poses, dim=-2)
+
+    @property
+    def palm_pose(self):
+        """
+        Get the palm pose for allegro hand
+        """
+        return vectorize_pose(self.palm_link.pose)
+
+
+class AllegroHandLeft(AllegroHandRight):
+    uid = "allegro_hand_left"
+    urdf_path = f"{PACKAGE_ASSET_DIR}/robots/allegro/allegro_hand_left.urdf"

mani_skill2/agents/robots/allegro_hand/allegro_touch.py

@@ -0,0 +1,153 @@
+import itertools
+from typing import List, Dict, Tuple, Optional
+
+import numpy as np
+import sapien
+import torch
+from sapien import physx
+
+from mani_skill2 import PACKAGE_ASSET_DIR
+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,
+)
+from mani_skill2.utils.sapien_utils import get_objs_by_names
+from mani_skill2.utils.structs.actor import Actor
+
+
+class AllegroHandRightTouch(AllegroHandRight):
+    uid = "allegro_hand_right_touch"
+    urdf_path = f"{PACKAGE_ASSET_DIR}/robots/allegro/variation/allegro_hand_right_fsr_simple.urdf"
+
+    def __init__(self, *args, **kwargs):
+        # Order: thumb finger, index finger, middle finger, ring finger, from finger root to fingertip
+        self.finger_fsr_link_names = [
+            # allegro thumb has a different hardware design compared with other fingers
+            "link_14.0_fsr",
+            "link_15.0_fsr",
+            "link_15.0_tip_fsr",
+            # the hardware design of index, middle and ring finger are the same
+            "link_1.0_fsr",
+            "link_2.0_fsr",
+            "link_3.0_tip_fsr",
+            "link_5.0_fsr",
+            "link_6.0_fsr",
+            "link_7.0_tip_fsr",
+            "link_9.0_fsr",
+            "link_10.0_fsr",
+            "link_11.0_tip_fsr",
+        ]
+        self.palm_fsr_link_names = [
+            "link_base_fsr",
+            "link_0.0_fsr",
+            "link_4.0_fsr",
+            "link_8.0_fsr",
+        ]
+
+        super().__init__(*args, **kwargs)
+
+        self.pair_query: Dict[
+            str, Tuple[physx.PhysxGpuContactPairImpulseQuery, Tuple[int, int, int]]
+        ] = dict()
+        self.body_query: Optional[
+            Tuple[physx.PhysxGpuContactBodyImpulseQuery, Tuple[int, int, int]]
+        ] = None
+
+    def _after_init(self):
+        super()._after_init()
+        self.fsr_links: List[Actor] = get_objs_by_names(
+            self.robot.get_links(),
+            self.palm_fsr_link_names + self.finger_fsr_link_names,
+        )
+
+    def get_fsr_obj_impulse(self, obj: Actor = None):
+        if physx.is_gpu_enabled():
+            px: sapien.physx.PhysxGpuSystem = self.scene.px
+            # Create contact query if it is not existed
+            if obj.name not in self.pair_query:
+                bodies = list(zip(*[link._bodies for link in self.fsr_links]))
+                bodies = list(itertools.chain(*bodies))
+                obj_bodies = [
+                    elem for item in obj._bodies for elem in itertools.repeat(item, 2)
+                ]
+                body_pairs = list(zip(bodies, obj_bodies))
+                query = px.gpu_create_contact_pair_impulse_query(body_pairs)
+                self.pair_query[obj.name] = (
+                    query,
+                    (len(obj._bodies), len(self.fsr_links), 3),
+                )
+
+            # Query contact buffer
+            query, contacts_shape = self.pair_query[obj.name]
+            px.gpu_query_contact_pair_impulses(query)
+            contacts = (
+                query.cuda_impulses.torch()
+                .clone()
+                .reshape((len(self.fsr_links), *contacts_shape))
+            )  # [n, 16, 3]
+
+            return contacts
+
+        else:
+            internal_fsr_links = [link._bodies[0].entity for link in self.fsr_links]
+            contacts = self.scene.get_contacts()
+            obj_contacts = get_multiple_pairwise_contacts(
+                contacts, obj._bodies[0].entity, internal_fsr_links
+            )
+            sorted_contacts = [obj_contacts[link] for link in internal_fsr_links]
+            contact_forces = [
+                compute_total_impulse(contact) for contact in sorted_contacts
+            ]
+
+            return np.stack(contact_forces)
+
+    def get_fsr_impulse(self):
+        if physx.is_gpu_enabled():
+            px: sapien.physx.PhysxGpuSystem = self.scene.px
+            # Create contact query if it is not existed
+            if self.body_query is None:
+                # Convert the order of links so that the link from the same sub-scene will come together
+                # It makes life easier for reshape
+                bodies = list(zip(*[link._bodies for link in self.fsr_links]))
+                bodies = list(itertools.chain(*bodies))
+
+                query = px.gpu_create_contact_body_impulse_query(bodies)
+                self.body_query = (
+                    query,
+                    (len(self.fsr_links[0]._bodies), len(self.fsr_links), 3),
+                )
+
+            # Query contact buffer
+            query, contacts_shape = self.body_query
+            px.gpu_query_contact_body_impulses(query)
+            contacts = (
+                query.cuda_impulses.torch().clone().reshape(*contacts_shape)
+            )  # [n, 16, 3]
+
+            return contacts
+
+        else:
+            internal_fsr_links = [link._bodies[0].entity for link in self.fsr_links]
+            contacts = self.scene.get_contacts()
+            contact_map = get_actors_contacts(contacts, internal_fsr_links)
+            sorted_contacts = [contact_map[link] for link in internal_fsr_links]
+            contact_forces = [
+                compute_total_impulse(contact) for contact in sorted_contacts
+            ]
+
+            contact_impulse = torch.from_numpy(
+                np.stack(contact_forces)[None, ...]
+            )  # [1, 16, 3]
+            return contact_impulse
+
+    def get_proprioception(self):
+        """
+        Get the proprioceptive state of the agent.
+        """
+        obs = super().get_proprioception()
+        fsr_impulse = self.get_fsr_impulse()
+        obs.update({"fsr_impulse": torch.linalg.norm(fsr_impulse, dim=-1)})
+
+        return obs