[update] update position retargeting example to be compataible with s…

…apien 3

example/position_retargeting/README.md

@@ -1,6 +1,6 @@
-## Retarget Robot Motion from Hand Object Pose Dataset
+# Retarget Robot Motion from Hand Object Pose Dataset
 
-### Setting up DexYCB Dataset
+## Setting up DexYCB Dataset
 
 This example illustrates how you can utilize the impressive DexYCB dataset to create a robot motion trajectory.
 The [DexYCB](https://dex-ycb.github.io/) is a hand-object dataset developed by NVIDIA.
@@ -13,6 +13,7 @@ In this case, we will be using the `20200709-subject-01.tar.gz` subset from DexY
 2. Download `models` and `calibration`, and keep them alongside the `20200709-subject-01.tar.gz`.
 
 ```Log
+.
 ├── 20200709-subject-01
 ├── calibration
 └── models
@@ -26,16 +27,24 @@ cd example/position_retargeting
 python dataset.py --dexycb-dir=PATH_TO_YOUR_DEXYCB_DIR_ROOT
 ```
 
-### Setting up manopth
+You will get something similar like this:
+
+```shell
+50
+Counter({'002_master_chef_can': 12, '005_tomato_soup_can': 9, '004_sugar_box': 6, '003_cracker_box': 6, '008_pudding_box': 4, '006_mustard_bottle': 4, '009_gelatin_box': 3, '007_tuna_fish_can': 2, '019_pitcher_base': 1, '024_bowl': 1, '021_bleach_cleanser': 1, '010_potted_meat_can': 1})
+dict_keys(['hand_pose', 'object_pose', 'extrinsics', 'ycb_ids', 'hand_shape', 'object_mesh_file', 'capture_name'])
+```
+
+## Setting up manopth
 
 Now, we will set up manopth similar to how it is done in [dex-ycb-toolkit](https://github.com/NVlabs/dex-ycb-toolkit).
 
 1. Download manopth in this directory, the manopth should be located
-   at `dex_retargeting/example/position_retargeting/dex_ycb`
+   at `dex_retargeting/example/position_retargeting`
 
     ```shell
     git clone https://github.com/hassony2/manopth
-    pip3 install chumpy opencv-python # install manopth dependencies
+    pip install chumpy opencv-python # install manopth dependencies
     ```
 
 2. Download MANO models and locally install manopth
@@ -44,33 +53,33 @@ Now, we will set up manopth similar to how it is done in [dex-ycb-toolkit](https
 
     ```shell
     cd manopth
-    pip3 install -e .
+    pip install -e .
     unzip mano_v1_2.zip
     cd mano
     ln -s ../mano_v1_2/models models
     ```
 
-### Installing Additional Python Dependencies
+## Installing Additional Python Dependencies
 
 ```shell
-pip install tyro pyyaml
+pip install tyro pyyaml sapien==3.0.0b0
 ```
 
-### Visualizing Human Hand-Object Interaction
+## Visualizing Human Hand-Object Interaction
 
-Before proceeding to retargeting, we can first visualize the original dataset in SAPIEN viewer. The hand mesh is
-computed via manopth/
+Before proceeding to retargeting, we can first visualize the original dataset in SAPIEN renderer. The hand mesh is
+computed via manopth.
 
 ```shell
-python3 visualize_hand_object.py --dexycb-dir=PATH_TO_YOUR_DEXYCB_DIR_ROOT
-# Press q to exit viewer
+python visualize_hand_object.py --dexycb-dir=PATH_TO_YOUR_DEXYCB_DIR_ROOT
+# Close the viewer window to quit
 ```
 
-### Visualizing Robot Hand-Object Interaction
+## Visualizing Robot Hand-Object Interaction
 
 Visualize the retargeting results for multiple robot hands along with the human hand.
 
 ```shell
-python3 visualize_hand_object.py --dexycb-dir=PATH_TO_YOUR_DEXYCB_DIR_ROOT --robots allegro shadow svh
-# Press q to exit viewer
+python visualize_hand_object.py --dexycb-dir=PATH_TO_YOUR_DEXYCB_DIR_ROOT --robots allegro shadow svh
+# Close the viewer window to quit
 ```

example/position_retargeting/hand_robot_viewer.py

@@ -1,10 +1,12 @@
+import tempfile
 from pathlib import Path
 from typing import Dict, List
 
 import numpy as np
-import sapien.core as sapien
+import sapien
 import transforms3d.quaternions
 
+from dex_retargeting import yourdfpy as urdf
 from dex_retargeting.constants import RobotName, HandType, get_default_config_path, RetargetingType
 from dex_retargeting.retargeting_config import RetargetingConfig
 from dex_retargeting.seq_retarget import SeqRetargeting
@@ -32,46 +34,47 @@ def prepare_vector_retargeting(joint_pos: np.array, link_hand_indices_pairs: np.
     return task_link_pos - origin_link_pos
 
 
-class LPFilter:
-    def __init__(self, control_freq, cutoff_freq):
-        dt = 1 / control_freq
-        wc = cutoff_freq * 2 * np.pi
-        y_cos = 1 - np.cos(wc * dt)
-        self.alpha = -y_cos + np.sqrt(y_cos**2 + 2 * y_cos)
-        self.y = 0
-        self.is_init = False
-
-    def next(self, x):
-        self.y = self.y + self.alpha * (x - self.y)
-        return self.y.copy()
-
-    def init(self, y):
-        self.y = y.copy()
-        self.is_init = True
-
-
 class RobotHandDatasetSAPIENViewer(HandDatasetSAPIENViewer):
     def __init__(self, robot_names: List[RobotName], hand_type: HandType, headless=False, use_ray_tracing=False):
         super().__init__(headless=headless, use_ray_tracing=use_ray_tracing)
 
         self.robots: List[sapien.Articulation] = []
         self.robot_file_names: List[str] = []
         self.retargetings: List[SeqRetargeting] = []
+        self.retarget2sapien: List[np.ndarray] = []
 
         # Load optimizer and filter
+        loader = self.scene.create_urdf_loader()
+        loader.fix_root_link = True
+        loader.load_multiple_collisions_from_file = True
         for robot_name in robot_names:
             config_path = get_default_config_path(robot_name, RetargetingType.position, hand_type)
 
             # Add 6-DoF dummy joint at the root of each robot to make them move freely in the space
             override = dict(add_dummy_free_joint=True)
             config = RetargetingConfig.load_from_file(config_path, override=override)
-            retargeting = config.build(self.scene)
-            robot = retargeting.optimizer.robot
+            retargeting = config.build()
             robot_file_name = Path(config.urdf_path).stem
-            self.robots.append(robot)
             self.robot_file_names.append(robot_file_name)
             self.retargetings.append(retargeting)
 
+            # Build robot
+            urdf_path = Path(config.urdf_path)
+            if "glb" not in urdf_path.stem:
+                urdf_path = str(urdf_path).replace(".urdf", "_glb.urdf")
+
+            robot_urdf = urdf.URDF.load(str(urdf_path), add_dummy_free_joints=True, build_scene_graph=False)
+            urdf_name = urdf_path.split("/")[-1]
+            temp_dir = tempfile.mkdtemp(prefix="dex_retargeting-")
+            temp_path = f"{temp_dir}/{urdf_name}"
+            robot_urdf.write_xml_file(temp_path)
+
+            robot = loader.load(temp_path)
+            self.robots.append(robot)
+            sapien_joint_names = [joint.name for joint in robot.get_active_joints()]
+            retarget2sapien = np.array([retargeting.joint_names.index(n) for n in sapien_joint_names]).astype(int)
+            self.retarget2sapien.append(retarget2sapien)
+
     def load_object_hand(self, data: Dict):
         super().load_object_hand(data)
         ycb_ids = data["ycb_ids"]
@@ -86,7 +89,7 @@ def load_object_hand(self, data: Dict):
     def render_dexycb_data(self, data: Dict, fps=5, y_offset=0.8):
         # Set table and viewer pose for better visual effect only
         global_y_offset = -y_offset * len(self.robots) / 2
-        self.table.set_pose(sapien.Pose([0.5, global_y_offset, 0]))
+        self.table.set_pose(sapien.Pose([0.5, global_y_offset + 0.2, 0]))
         if self.viewer is not None:
             self.viewer.set_camera_xyz(1.5, global_y_offset, 1)
 
@@ -132,14 +135,14 @@ def render_dexycb_data(self, data: Dict, fps=5, y_offset=0.8):
             self._update_hand(vertex)
 
             # Update poses for robot hands
-            for robot, retargeting in zip(self.robots, self.retargetings):
+            for robot, retargeting, retarget2sapien in zip(self.robots, self.retargetings, self.retarget2sapien):
                 indices = retargeting.optimizer.target_link_human_indices
                 ref_value = joint[indices, :]
-                qpos = retargeting.retarget(ref_value)
+                qpos = retargeting.retarget(ref_value)[retarget2sapien]
                 robot.set_qpos(qpos)
 
             for k in range(60 // fps):
                 self.viewer.render()
 
-        self.viewer.toggle_pause(True)
+        self.viewer.paused = True
         self.viewer.render()

example/position_retargeting/hand_viewer.py

@@ -1,10 +1,12 @@
 from typing import Dict, List, Optional
 
 import numpy as np
+import sapien
 import sapien.core as sapien
 import torch
+from pytransform3d import transformations as pt
+from sapien import internal_renderer as R
 from sapien.asset import create_dome_envmap
-from sapien.core.pysapien import renderer as R
 from sapien.utils import Viewer
 
 from dataset import YCB_CLASSES
@@ -39,25 +41,18 @@ def compute_smooth_shading_normal_np(vertices, indices):
 class HandDatasetSAPIENViewer:
     def __init__(self, headless=False, use_ray_tracing=False):
         # Setup
-        engine = sapien.Engine()
-        engine.set_log_level("warning")
-
-        if use_ray_tracing:
-            sapien.render_config.camera_shader_dir = "rt"
-            sapien.render_config.viewer_shader_dir = "rt"
-            sapien.render_config.rt_samples_per_pixel = 32
-            sapien.render_config.rt_use_denoiser = True
+        if not use_ray_tracing:
+            sapien.render.set_viewer_shader_dir("default")
+            sapien.render.set_camera_shader_dir("default")
         else:
-            sapien.render_config.camera_shader_dir = "ibl"
-            sapien.render_config.viewer_shader_dir = "ibl"
-
-        renderer = sapien.SapienRenderer(offscreen_only=headless)
-        engine.set_renderer(renderer)
-        renderer.set_log_level("error")
+            sapien.render.set_viewer_shader_dir("rt")
+            sapien.render.set_camera_shader_dir("rt")
+            sapien.render.set_ray_tracing_samples_per_pixel(16)
+            sapien.render.set_ray_tracing_path_depth(8)
+            sapien.render.set_ray_tracing_denoiser("oidn")
 
         # Scene
-        scene_config = sapien.SceneConfig()
-        scene = engine.create_scene(scene_config)
+        scene = sapien.Scene()
         scene.set_timestep(1 / 240)
 
         # Lighting
@@ -75,7 +70,7 @@ def __init__(self, headless=False, use_ray_tracing=False):
         )
 
         # Add ground
-        visual_material = renderer.create_material()
+        visual_material = sapien.render.RenderMaterial()
         visual_material.set_base_color(np.array([0.5, 0.5, 0.5, 1]))
         visual_material.set_roughness(0.7)
         visual_material.set_metallic(1)
@@ -84,16 +79,16 @@ def __init__(self, headless=False, use_ray_tracing=False):
 
         # Viewer
         if not headless:
-            viewer = Viewer(renderer)
+            viewer = Viewer()
             viewer.set_scene(scene)
             viewer.set_camera_xyz(1.5, 0, 1)
             viewer.set_camera_rpy(0, -0.6, 3.14)
-            viewer.toggle_axes(0)
+            viewer.control_window.toggle_origin_frame(False)
             self.viewer = viewer
         self.gui = not headless
 
-        # Table
-        white_diffuse = renderer.create_material()
+        # Create table
+        white_diffuse = sapien.render.RenderMaterial()
         white_diffuse.set_base_color(np.array([0.8, 0.8, 0.8, 1]))
         white_diffuse.set_roughness(0.9)
         builder = scene.create_actor_builder()
@@ -115,17 +110,16 @@ def __init__(self, headless=False, use_ray_tracing=False):
         self.table.set_pose(sapien.Pose([0.5, 0, 0]))
 
         # Caches
-        self.engine = engine
-        self.renderer = renderer
+        sapien.render.set_log_level("error")
         self.scene = scene
-        self.internal_scene: R.Scene = scene.get_renderer_scene()._internal_scene
-        self.context: R.Context = renderer._internal_context
+        self.internal_scene: R.Scene = scene.render_system._internal_scene
+        self.context: R.Context = sapien.render.SapienRenderer()._internal_context
         self.mat_hand = self.context.create_material(np.zeros(4), np.array([0.96, 0.75, 0.69, 1]), 0.0, 0.8, 0)
 
         self.mano_layer: Optional[MANOLayer] = None
         self.mano_face: Optional[np.ndarray] = None
         self.camera_pose: Optional[sapien.Pose] = None
-        self.objects: List[sapien.ActorStatic] = []
+        self.objects: List[sapien.Entity] = []
         self.nodes: List[R.Node] = []
 
     def clear_all(self):
@@ -152,7 +146,8 @@ def load_object_hand(self, data: Dict):
 
         self.mano_layer = MANOLayer("right", hand_shape.astype(np.float32))
         self.mano_face = self.mano_layer.f.cpu().numpy()
-        self.camera_pose = sapien.Pose.from_transformation_matrix(extrinsic_mat).inv()
+        pose_vec = pt.pq_from_transform(extrinsic_mat)
+        self.camera_pose = sapien.Pose(pose_vec[0:3], pose_vec[3:7]).inv()
 
     def _load_ycb_object(self, ycb_id, ycb_mesh_file):
         builder = self.scene.create_actor_builder()
@@ -213,5 +208,5 @@ def render_dexycb_data(self, data: Dict, fps=10):
             for _ in range(step_per_frame):
                 self.viewer.render()
 
+        self.viewer.paused = True
         self.viewer.render()
-        self.viewer.toggle_pause(True)

example/position_retargeting/visualize_hand_object.py

@@ -1,33 +1,54 @@
-import argparse
 from pathlib import Path
 from typing import Optional, Tuple, List
 
+import numpy as np
 import tyro
 
 from dataset import DexYCBVideoDataset
 from dex_retargeting.constants import RobotName, HandType
+from dex_retargeting.retargeting_config import RetargetingConfig
 from hand_robot_viewer import RobotHandDatasetSAPIENViewer
 from hand_viewer import HandDatasetSAPIENViewer
-from dex_retargeting.retargeting_config import RetargetingConfig
+
+# For numpy version compatibility
+np.bool = bool
+np.int = int
+np.float = float
+np.str = str
+np.complex = complex
+np.object = object
+np.unicode = np.unicode_
 
 
 def viz_hand_object(robots: Optional[Tuple[RobotName]], data_root: Path, fps: int):
-    dataset = DexYCBVideoDataset(data_root)
+    dataset = DexYCBVideoDataset(data_root, hand_type="right")
     if robots is None:
         viewer = HandDatasetSAPIENViewer(headless=False)
     else:
         viewer = RobotHandDatasetSAPIENViewer(list(robots), HandType.right, headless=False)
 
-    sampled_data = dataset[4]
+    # Data ID, feel free to change it to visualize different trajectory
+    data_id = 4
+
+    sampled_data = dataset[data_id]
     for key, value in sampled_data.items():
         if "pose" not in key:
             print(f"{key}: {value}")
     viewer.load_object_hand(sampled_data)
     viewer.render_dexycb_data(sampled_data, fps)
 
 
-def main(dexycb_dir: str, robots: Optional[List[RobotName]] = None, fps: int = 5):
-    print(robots)
+def main(dexycb_dir: str, robots: Optional[List[RobotName]] = None, fps: int = 10):
+    """
+    Render the human and robot trajectories for grasping object inside DexYCB dataset.
+    The human trajectory is visualized as provided, while the robot trajectory is generated from position retargeting
+
+    Args:
+        dexycb_dir: Data root path to the dexycb dataset
+        robots: The names of robots to render, if None, render human hand trajectory only
+        fps: frequency to render hand-object trajectory
+
+    """
     data_root = Path(dexycb_dir).absolute()
     robot_dir = Path(__file__).absolute().parent.parent.parent / "assets" / "robots" / "hands"
     RetargetingConfig.set_default_urdf_dir(robot_dir)