Rotate Cube Task using Trifingerpro robot (#249)

* update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1 * update rotate_cube v1
haosulab · Mar 28, 2024 · 9737f5d · 9737f5d
1 parent 2100487
commit 9737f5d
Show file tree

Hide file tree

Showing 76 changed files with 149,804 additions and 4 deletions.
diff --git a/docs/source/tasks/index.md b/docs/source/tasks/index.md
@@ -256,4 +256,28 @@ Using the D'Claw robot, rotate a [ROBEL valve](https://sites.google.com/view/rob
 
 :::{figure} https://github.com/haosulab/ManiSkill2/raw/dev/figures/environment_demos/RotateValve-v1_rt.png
 :alt: rotate valve task
-:::
+:::
+
+
+### RotateCubeLevel0-v1
+
+(Note there is Level0, Level1, ... to Level4)
+
+:::{dropdown} Task Card
+:icon: note
+:color: primary
+
+**Task Description:**
+Using the TriFingerPro robot, rotate a cube
+
+**Supported Robots: TriFingerPro**
+
+**Randomizations:**
+- Level 0: Random goal position on the table, no orientation.
+- Level 1:  Random goal position on the table, including yaw orientation.
+- Level 2: Fixed goal position in the air with x,y = 0.  No orientation.
+- Level 3: Random goal position in the air, no orientation.
+- Level 4: Random goal pose in the air, including orientation.
+
+**Success Conditions:**
+- The rotated cube should be within 0.02 m of the goal position and 0.1 rad of the goal orientation.
diff --git a/examples/baselines/ppo/README.md b/examples/baselines/ppo/README.md
@@ -20,6 +20,21 @@ python ppo.py --env_id="PickSingleYCB-v1" \
 python ppo.py --env_id="TwoRobotStackCube-v1" \
    --num_envs=1024 --update_epochs=8 --num_minibatches=32 \
    --total_timesteps=40_000_000 --num-steps=100 --num-eval-steps=100
+python ppo.py --env_id="RotateCubeLevel0-v1" \
+   --num_envs=128 --update_epochs=8 --num_minibatches=32 \
+   --total_timesteps=50_000_000 --num-steps=250 --num-eval-steps=250
+python ppo.py --env_id="RotateCubeLevel1-v1" \
+   --num_envs=128 --update_epochs=8 --num_minibatches=32 \
+   --total_timesteps=50_000_000 --num-steps=250 --num-eval-steps=250
+python ppo.py --env_id="RotateCubeLevel2-v1" \
+   --num_envs=128 --update_epochs=8 --num_minibatches=32 \
+   --total_timesteps=50_000_000 --num-steps=250 --num-eval-steps=250
+python ppo.py --env_id="RotateCubeLevel3-v1" \
+   --num_envs=128 --update_epochs=8 --num_minibatches=32 \
+   --total_timesteps=50_000_000 --num-steps=250 --num-eval-steps=250
+python ppo.py --env_id="RotateCubeLevel4-v1" \
+   --num_envs=1024 --update_epochs=8 --num_minibatches=32 \
+   --total_timesteps=500_000_000 --num-steps=250 --num-eval-steps=250
 ```
 
 <!-- TODO (stao, arnav) clean up the baseline code to be slightly nicer (track FPS, and update time separately), and put results onto wandb. Also merge code with CleanRL Repo (stao can ask costa to help do that) -->
diff --git a/mani_skill/agents/robots/__init__.py b/mani_skill/agents/robots/__init__.py
@@ -6,6 +6,7 @@
 from .panda import Panda, PandaRealSensed435
 from .xarm import XArm7Ability
 from .xmate3 import Xmate3Robotiq
+from .trifingerpro import TriFingerPro
 
 # = {
 #     "panda": Panda,

diff --git a/mani_skill/agents/robots/trifingerpro/__init__.py b/mani_skill/agents/robots/trifingerpro/__init__.py
@@ -0,0 +1 @@
+from .trifingerpro import TriFingerPro
diff --git a/mani_skill/agents/robots/trifingerpro/trifingerpro.py b/mani_skill/agents/robots/trifingerpro/trifingerpro.py
@@ -0,0 +1,176 @@
+from copy import deepcopy
+from typing import List
+
+import sapien
+import torch
+
+from mani_skill import PACKAGE_ASSET_DIR
+from mani_skill.agents.base_agent import BaseAgent
+from mani_skill.agents.controllers import *
+from mani_skill.agents.registration import register_agent
+from mani_skill.agents.utils import (
+    get_active_joint_indices,
+)
+from mani_skill.utils.sapien_utils import (
+    get_objs_by_names,
+)
+from mani_skill.utils.structs.pose import vectorize_pose
+
+
+@register_agent()
+class TriFingerPro(BaseAgent):
+    """
+    Modified from https://github.com/NVIDIA-Omniverse/IsaacGymEnvs/blob/main/isaacgymenvs/tasks/trifinger.py
+
+    """
+    uid = "trifingerpro"
+    urdf_path = f"{PACKAGE_ASSET_DIR}/robots/trifinger/trifingerpro.urdf"
+    urdf_config = dict(
+        _materials=dict(
+            tip=dict(static_friction=2.0, dynamic_friction=1.0, restitution=0.0)
+        ),
+        link=dict(
+            finger_tip_link_0=dict(material="tip"),
+            finger_tip_link_120=dict(material="tip"),
+            finger_tip_link_240=dict(material="tip"),
+        ),
+    )
+    sensor_configs = {}
+
+    def __init__(self, *args, **kwargs):
+        self.joint_names = [
+            # "base_to_upper_holder_joint",
+            # "finger_upper_visuals_joint_0",
+            # "finger_middle_visuals_joint_0",
+            # "finger_lower_to_tip_joint_0",
+            "finger_base_to_upper_joint_0",
+            "finger_upper_to_middle_joint_0",
+            "finger_middle_to_lower_joint_0",
+            # "finger_upper_visuals_joint_120",
+            # "finger_middle_visuals_joint_120",
+            # "finger_lower_to_tip_joint_120",
+            "finger_base_to_upper_joint_120",
+            "finger_upper_to_middle_joint_120",
+            "finger_middle_to_lower_joint_120",
+            # "finger_upper_visuals_joint_240",
+            # "finger_middle_visuals_joint_240",
+            # "finger_lower_to_tip_joint_240",
+            "finger_base_to_upper_joint_240",
+            "finger_upper_to_middle_joint_240",
+            "finger_middle_to_lower_joint_240",
+            # "holder_to_finger_0",
+            # "holder_to_finger_120",
+            # "holder_to_finger_240",
+        ]
+
+        self.joint_stiffness = 1e2
+        self.joint_damping = 1e1
+        self.joint_force_limit = 2e1
+        self.tip_link_names = ["finger_tip_link_0", "finger_tip_link_120", "finger_tip_link_240"]
+        self.root_joint_names = ["finger_base_to_upper_joint_0", "finger_base_to_upper_joint_120", "finger_base_to_upper_joint_240"]
+
+        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.root_joints = [
+            self.robot.find_joint_by_name(n) for n in self.root_joint_names
+        ]
+        self.root_joint_indices = get_active_joint_indices(
+            self.robot, self.root_joint_names
+        )
+
+    @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
+
+        # PD joint velocity
+        pd_joint_vel = PDJointVelControllerConfig(
+            self.joint_names,
+            -1.0,
+            1.0,
+            self.joint_damping,  # this might need to be tuned separately
+            self.joint_force_limit,
+        )
+
+        # PD joint position and velocity
+        joint_pos_vel = PDJointPosVelControllerConfig(
+            self.joint_names,
+            None,
+            None,
+            self.joint_stiffness,
+            self.joint_damping,
+            self.joint_force_limit,
+            normalize_action=False,
+        )
+        joint_delta_pos_vel = PDJointPosVelControllerConfig(
+            self.joint_names,
+            -0.1,
+            0.1,
+            self.joint_stiffness,
+            self.joint_damping,
+            self.joint_force_limit,
+            use_delta=True,
+        )
+
+        controller_configs = dict(
+            pd_joint_delta_pos=dict(joint=joint_delta_pos),
+            pd_joint_pos=dict(joint=joint_pos),
+            pd_joint_target_delta_pos=dict(joint=joint_target_delta_pos),
+            # Caution to use the following controllers
+            pd_joint_vel=dict(joint=pd_joint_vel),
+            pd_joint_pos_vel=dict(joint=joint_pos_vel),
+            pd_joint_delta_pos_vel=dict(joint=joint_delta_pos_vel),
+        )
+
+        # 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({"tip_poses": self.tip_poses.view(-1, 21)})
+        obs.update({"tip_velocities": self.tip_velocities().view(-1, 9)})
+        return obs
+
+    @property
+    def tip_poses(self):
+        """
+        Get the tip pose for each of the finger, three fingers in total
+        """
+        tip_poses = [vectorize_pose(link.pose) for link in self.tip_links]
+        return torch.stack(tip_poses, dim=-1)
+
+    # @property
+    def tip_velocities(self):
+        """
+        Get the tip velocity for each of the finger, three fingers in total
+        """
+        tip_velocities = [link.linear_velocity for link in self.tip_links]
+        return torch.stack(tip_velocities, dim=-1)
diff --git a/mani_skill/assets/robots/trifinger/cube_multicolor_rrc.urdf b/mani_skill/assets/robots/trifinger/cube_multicolor_rrc.urdf
@@ -0,0 +1,20 @@
+<?xml version="1.0"?>
+<robot name="object">
+  <link name="object">
+    <visual>
+      <origin xyz="0 0 0"/>
+      <geometry>
+        <mesh filename="meshes/cube_multicolor.obj" scale=".065 .065 .065"/>
+      </geometry>
+    </visual>
+    <collision>
+      <origin xyz="0 0 0"/>
+      <geometry>
+        <box size="0.065 0.065 0.065"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <density value="291.3"/>
+    </inertial>
+  </link>
+</robot>
diff --git a/mani_skill/assets/robots/trifinger/high_table_boundary.urdf b/mani_skill/assets/robots/trifinger/high_table_boundary.urdf
@@ -0,0 +1,33 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- =================================================================================== -->
+<!-- |    This document was autogenerated by xacro from /home/felixwidmaier/ws/fingers/workspace/src/catkin/robots/robot_properties/robot_properties_fingers/xacro/stage.urdf.xacro | -->
+<!-- |    EDITING THIS FILE BY HAND IS NOT RECOMMENDED                                 | -->
+<!-- =================================================================================== -->
+<robot name="stage">
+  <material name="stage_material">
+    <color rgba="0.8 0.8 0.8 1.0"/>
+  </material>
+  <!-- add the "stage" -->
+  <link name="stage_link">
+    <visual>
+      <origin rpy="0 0 0" xyz="0 0 0"/>
+      <geometry>
+        <mesh filename="package://robot_properties_fingers/meshes/high_table_boundary.stl" scale="1 1 1"/>
+      </geometry>
+      <material name="stage_material"/>
+    </visual>
+
+    <collision>
+      <origin rpy="0 0 0" xyz="0 0 0"/>
+      <geometry>
+        <mesh filename="package://robot_properties_fingers/meshes/high_table_boundary.stl" scale="1 1 1"/>
+      </geometry>
+    </collision>
+
+    <inertial>
+      <origin rpy="0 0 0" xyz="0 0 0"/>
+      <mass value="1"/>
+      <inertia ixx="0.00770833333333" ixy="0" ixz="0" iyy="0.00770833333333" iyz="0" izz="0.015"/>
+    </inertial>
+  </link>
+</robot>