Spaces:
Running on Zero
Running on Zero
| # Project EmbodiedGen | |
| # | |
| # Copyright (c) 2025 Horizon Robotics. All Rights Reserved. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or | |
| # implied. See the License for the specific language governing | |
| # permissions and limitations under the License. | |
| import concurrent.futures | |
| import dataclasses | |
| import os | |
| import threading | |
| import time | |
| import warnings | |
| from collections import defaultdict | |
| from dataclasses import dataclass, field | |
| warnings.filterwarnings("ignore", category=UserWarning) | |
| from typing import Literal | |
| import imageio | |
| import numpy as np | |
| import sapien.core as sapien | |
| import sapien.physx as physx | |
| import torch | |
| import tyro | |
| from PIL import Image, ImageColor | |
| from scipy.spatial.transform import Rotation as R | |
| from tqdm import tqdm | |
| from embodied_gen.data.utils import DiffrastRender | |
| from embodied_gen.utils.geometry import MeshInfo | |
| from embodied_gen.utils.io_utils import ( | |
| URDFFile, | |
| load_json, | |
| load_mesh_info, | |
| write_json, | |
| ) | |
| from embodied_gen.utils.log import logger | |
| COLORMAP = list(set(ImageColor.colormap.values())) | |
| COLOR_PALETTE = np.array( | |
| [ImageColor.getrgb(c) for c in COLORMAP], dtype=np.uint8 | |
| ) | |
| MAX_NUM_WORKERS = 10 | |
| MAX_GRASPS_PER_WORKER_PROCESS = 8 | |
| SAPIEN_DEFAULT_STATIC_FRICTION = 0.3 | |
| SAPIEN_DEFAULT_DYNAMIC_FRICTION = 0.3 | |
| SAPIEN_DEFAULT_RESTITUTION = 0.1 | |
| _EVALUATOR_KEEPALIVE: list["EvalGrasps"] = [] | |
| _EVALUATOR_KEEPALIVE_LOCK = threading.Lock() | |
| __all__ = [ | |
| "EvalGraspsConfig", | |
| "GripperControlInfo", | |
| "SapienSceneManager", | |
| "EvalGrasps", | |
| "run_eval", | |
| "entrypoint", | |
| ] | |
| class EvalGraspsConfig: | |
| urdf_paths: list[str] = field(default_factory=list) | |
| gripper_name: str = "franka_panda" | |
| gripper_urdf_path: str = ( | |
| f"embodied_gen/scripts/affordance_annot/{gripper_name}/gripper.urdf" | |
| ) | |
| output_dirs: list[str] = field(default_factory=list) | |
| sim_freq: int = 240 | |
| video_fps: int = 30 | |
| render_interval: int = 24 | |
| image_hw: tuple[int, int] = (512, 512) | |
| num_cameras: int = 1 | |
| fovy_deg: float = 45.0 | |
| sim_backend: str = "cpu" | |
| ray_tracing: bool = False | |
| gripper_move_speed: float = 0.1 | |
| gripper_close_speed: float = 0.01 | |
| action_wait_seconds: float = 0.5 | |
| lift_height: float = 0.5 | |
| sweep_distance: float = 0.5 | |
| num_workers: int | None = None | |
| max_grasps_per_worker_process: int = MAX_GRASPS_PER_WORKER_PROCESS | |
| gripper_drive_stiffness: float = 2500.0 | |
| gripper_static_friction: float = SAPIEN_DEFAULT_STATIC_FRICTION | |
| gripper_dynamic_friction: float = SAPIEN_DEFAULT_DYNAMIC_FRICTION | |
| gripper_restitution: float = SAPIEN_DEFAULT_RESTITUTION | |
| root_drive_stiffness: float = 1e6 | |
| root_drive_damping: float = 1e4 | |
| object_bottom_clearance: float = 0.1 | |
| hold_object_during_close: bool = True | |
| slip_translation_threshold: float = 0.05 | |
| slip_rotation_threshold_deg: float = 30.0 | |
| check_close_slip: bool = True | |
| debug_mode: bool = False | |
| save_failed_grasps: bool = False | |
| class GripperControlInfo: | |
| open_qpos: float | |
| close_qpos: float | |
| drive_damping: float | |
| force_limit: float | |
| class GripperActionAgent: | |
| def __init__( | |
| self, | |
| scene_manager: "SapienSceneManager", | |
| gripper: physx.PhysxArticulation, | |
| root_target: sapien.Entity, | |
| root_drive: physx.PhysxDriveComponent, | |
| ) -> None: | |
| self.scene_manager = scene_manager | |
| self.gripper = gripper | |
| self.root_target = root_target | |
| self.root_drive = root_drive | |
| def set_action(self, action: np.ndarray | torch.Tensor) -> None: | |
| if isinstance(action, torch.Tensor): | |
| values = action.detach().cpu().numpy().reshape(-1) | |
| else: | |
| values = np.asarray(action).reshape(-1) | |
| if values.size != 8: | |
| raise ValueError( | |
| f"Gripper action must have 8 values, got {values.size}." | |
| ) | |
| self.root_target.set_pose( | |
| sapien.Pose( | |
| p=values[:3].astype(np.float64), | |
| q=values[3:7].astype(np.float64), | |
| ) | |
| ) | |
| self.scene_manager.set_gripper_target( | |
| self.gripper, | |
| float(values[7]), | |
| ) | |
| class SapienSceneManager: | |
| def __init__( | |
| self, sim_freq: int, ray_tracing: bool, device: str = "cpu" | |
| ) -> None: | |
| self.sim_freq = sim_freq | |
| self.ray_tracing = ray_tracing | |
| self.device = device | |
| self.renderer = sapien.SapienRenderer() | |
| self.scene = self.setup_scene() | |
| self.cameras: list[sapien.render.RenderCameraComponent] = [] | |
| self.actors: dict[str, sapien.pysapien.Entity] = {} | |
| def setup_scene(self) -> sapien.Scene: | |
| # Ray tracing settings | |
| if self.ray_tracing: | |
| sapien.render.set_camera_shader_dir("rt") | |
| sapien.render.set_ray_tracing_samples_per_pixel(64) | |
| sapien.render.set_ray_tracing_path_depth(10) | |
| sapien.render.set_ray_tracing_denoiser("oidn") | |
| scene = sapien.Scene() | |
| scene.set_timestep(1 / self.sim_freq) | |
| # Add lighting | |
| scene.set_ambient_light([0.2, 0.2, 0.2]) | |
| scene.add_directional_light( | |
| direction=[0, 1, -1], | |
| color=[1.5, 1.45, 1.4], | |
| shadow=True, | |
| shadow_map_size=2048, | |
| ) | |
| scene.add_directional_light( | |
| direction=[0, -0.5, 1], color=[0.8, 0.8, 0.85], shadow=False | |
| ) | |
| scene.add_directional_light( | |
| direction=[0, -1, 1], color=[1.0, 1.0, 1.0], shadow=False | |
| ) | |
| ground_material = self.renderer.create_material() | |
| ground_material.base_color = [0.5, 0.5, 0.5, 1] # rgba, gray | |
| ground_material.roughness = 0.7 | |
| ground_material.metallic = 0.0 | |
| scene.add_ground(0, render_material=ground_material) | |
| return scene | |
| def set_ground_color(self, rgba: list[float]) -> None: | |
| """Update the default ground plane material color for this scene.""" | |
| for actor in self.scene.get_all_actors(): | |
| if actor.name != "ground": | |
| continue | |
| for component in actor.components: | |
| render_shapes = getattr(component, "render_shapes", None) | |
| if render_shapes is None: | |
| continue | |
| for render_shape in render_shapes: | |
| render_shape.material.set_base_color(rgba) | |
| return | |
| raise ValueError("Ground actor not found in the scene.") | |
| def create_camera( | |
| self, | |
| cam_name: str, | |
| pose: sapien.Pose, | |
| image_hw: tuple[int, int], | |
| fovy_deg: float, | |
| ) -> sapien.render.RenderCameraComponent: | |
| cam_actor = self.scene.create_actor_builder().build_kinematic() | |
| cam_actor.set_pose(pose) | |
| camera = self.scene.add_mounted_camera( | |
| name=cam_name, | |
| mount=cam_actor, | |
| pose=sapien.Pose(p=[0, 0, 0], q=[1, 0, 0, 0]), | |
| width=image_hw[1], | |
| height=image_hw[0], | |
| fovy=np.deg2rad(fovy_deg), | |
| near=0.01, | |
| far=100, | |
| ) | |
| self.cameras.append(camera) | |
| return camera | |
| def add_ring_cameras( | |
| self, | |
| num_cameras: int, | |
| radius: float, | |
| height: float, | |
| target_pt: list[float], | |
| image_hw: tuple[int, int], | |
| fovy_deg: float, | |
| ) -> list[sapien.render.RenderCameraComponent]: | |
| angle_step = 2 * np.pi / num_cameras | |
| world_up_vec = np.array([0.0, 0.0, 1.0]) | |
| target_pt = np.array(target_pt) | |
| for i in range(num_cameras): | |
| angle = i * angle_step | |
| cam_x = target_pt[0] + radius * np.cos(angle) | |
| cam_y = target_pt[1] + radius * np.sin(angle) | |
| cam_z = height | |
| eye_pos = [cam_x, cam_y, cam_z] | |
| forward_vec = target_pt - eye_pos | |
| forward_vec = forward_vec / np.linalg.norm(forward_vec) | |
| temp_right_vec = np.cross(forward_vec, world_up_vec) | |
| if np.linalg.norm(temp_right_vec) < 1e-6: | |
| temp_right_vec = np.array([1.0, 0.0, 0.0]) | |
| if np.abs(np.dot(temp_right_vec, forward_vec)) > 0.99: | |
| temp_right_vec = np.array([0.0, 1.0, 0.0]) | |
| right_vec = temp_right_vec / np.linalg.norm(temp_right_vec) | |
| up_vec = np.cross(right_vec, forward_vec) | |
| rotation_matrix = np.array([forward_vec, -right_vec, up_vec]).T | |
| rot = R.from_matrix(rotation_matrix) | |
| scipy_quat = rot.as_quat() # (x, y, z, w) | |
| quat = [ | |
| scipy_quat[3], | |
| scipy_quat[0], | |
| scipy_quat[1], | |
| scipy_quat[2], | |
| ] # (w, x, y, z) | |
| self.create_camera( | |
| f"camera_{i}", | |
| sapien.Pose(p=eye_pos, q=quat), | |
| image_hw, | |
| fovy_deg, | |
| ) | |
| return self.cameras | |
| def add_actor( | |
| self, | |
| mesh_info: MeshInfo, | |
| position: list[float], | |
| ) -> sapien.Entity: | |
| builder = self.scene.create_actor_builder() | |
| material = physx.PhysxMaterial( | |
| static_friction=float( | |
| np.clip(mesh_info.static_friction, 0.1, 0.7) | |
| ), | |
| dynamic_friction=float( | |
| np.clip(mesh_info.dynamic_friction, 0.1, 0.6) | |
| ), | |
| restitution=0.05, | |
| ) | |
| builder.set_physx_body_type("dynamic") | |
| builder.add_multiple_convex_collisions_from_file( | |
| mesh_info.collision_mesh_path, | |
| scale=mesh_info.collision_mesh_scale, | |
| pose=mesh_info.collision_local_pose, | |
| material=material, | |
| ) | |
| builder.add_visual_from_file( | |
| mesh_info.visual_mesh_path, | |
| scale=mesh_info.visual_mesh_scale, | |
| pose=mesh_info.visual_local_pose, | |
| ) | |
| builder.set_initial_pose(sapien.Pose(p=position)) | |
| actor = builder.build(name=mesh_info.actor_name) | |
| if mesh_info.mass is not None and hasattr( | |
| actor.components[1], "set_mass" | |
| ): | |
| actor.components[1].set_mass(mesh_info.mass) | |
| return actor | |
| def add_gripper( | |
| self, | |
| gripper_urdf_path: str, | |
| root_pose: sapien.Pose, | |
| open_qpos: float, | |
| drive_stiffness: float, | |
| drive_damping: float, | |
| force_limit: float, | |
| static_friction: float, | |
| dynamic_friction: float, | |
| restitution: float, | |
| root_drive_stiffness: float, | |
| root_drive_damping: float, | |
| ) -> tuple[ | |
| physx.PhysxArticulation, | |
| sapien.Entity, | |
| physx.PhysxDriveComponent, | |
| ]: | |
| loader = self.scene.create_urdf_loader() | |
| loader.fix_root_link = False | |
| gripper = loader.load(gripper_urdf_path) | |
| gripper.set_root_pose(root_pose) | |
| self.set_articulation_contact_material( | |
| gripper, | |
| physx.PhysxMaterial( | |
| static_friction, | |
| dynamic_friction, | |
| restitution, | |
| ), | |
| ) | |
| for link in gripper.get_links(): | |
| link.set_disable_gravity(True) | |
| root_target = self.scene.create_actor_builder().build_kinematic( | |
| name="gripper_root_target" | |
| ) | |
| root_target.set_pose(root_pose) | |
| root_drive = self.scene.create_drive( | |
| root_target, | |
| sapien.Pose(), | |
| gripper.get_root(), | |
| sapien.Pose(), | |
| ) | |
| root_drive.set_limit_x( | |
| 0.0, 0.0, root_drive_stiffness, root_drive_damping | |
| ) | |
| root_drive.set_limit_y( | |
| 0.0, 0.0, root_drive_stiffness, root_drive_damping | |
| ) | |
| root_drive.set_limit_z( | |
| 0.0, 0.0, root_drive_stiffness, root_drive_damping | |
| ) | |
| root_drive.set_limit_twist( | |
| 0.0, 0.0, root_drive_stiffness, root_drive_damping | |
| ) | |
| root_drive.set_limit_cone( | |
| 0.0, 0.0, root_drive_stiffness, root_drive_damping | |
| ) | |
| root_drive.set_drive_property_x( | |
| root_drive_stiffness, root_drive_damping | |
| ) | |
| root_drive.set_drive_property_y( | |
| root_drive_stiffness, root_drive_damping | |
| ) | |
| root_drive.set_drive_property_z( | |
| root_drive_stiffness, root_drive_damping | |
| ) | |
| root_drive.set_drive_property_slerp( | |
| root_drive_stiffness, | |
| root_drive_damping, | |
| ) | |
| root_drive.set_inv_mass_scales(0.0, 1.0) | |
| qpos = [] | |
| for joint in gripper.get_active_joints(): | |
| joint.set_drive_properties( | |
| drive_stiffness, | |
| drive_damping, | |
| force_limit, | |
| ) | |
| joint.set_friction(0.1) | |
| limit = np.asarray(joint.get_limits()).reshape(-1, 2)[0] | |
| target = float(np.clip(open_qpos, limit[0], limit[1])) | |
| joint.set_drive_target(target) | |
| qpos.append(target) | |
| if qpos: | |
| gripper.set_qpos(qpos) | |
| gripper.set_qvel(np.zeros(len(qpos), dtype=np.float32)) | |
| return gripper, root_target, root_drive | |
| def set_articulation_contact_material( | |
| articulation: physx.PhysxArticulation, | |
| material: physx.PhysxMaterial, | |
| ) -> None: | |
| for link in articulation.get_links(): | |
| for shape in link.get_collision_shapes(): | |
| shape.set_physical_material(material) | |
| def set_gripper_target( | |
| self, | |
| gripper: physx.PhysxArticulation, | |
| finger_qpos: float, | |
| ) -> None: | |
| for joint in gripper.get_active_joints(): | |
| limit = np.asarray(joint.get_limits()).reshape(-1, 2)[0] | |
| target = float(np.clip(finger_qpos, limit[0], limit[1])) | |
| joint.set_drive_target(target) | |
| def render_images( | |
| self, | |
| camera, | |
| render_keys: list[ | |
| Literal[ | |
| "Color", | |
| "Segmentation", | |
| "Normal", | |
| "Mask", | |
| "Depth", | |
| "Foreground", | |
| ] | |
| ] = None, | |
| ) -> dict[str, Image.Image]: | |
| if render_keys is None: | |
| render_keys = [ | |
| "Color", | |
| "Segmentation", | |
| "Normal", | |
| "Mask", | |
| "Depth", | |
| "Foreground", | |
| ] | |
| results: dict[str, Image.Image] = {} | |
| if "Color" in render_keys: | |
| color = camera.get_picture("Color") | |
| color_rgb = (np.clip(color[..., :3], 0, 1) * 255).astype(np.uint8) | |
| results["Color"] = Image.fromarray(color_rgb) | |
| if "Mask" in render_keys: | |
| alpha = (np.clip(color[..., 3], 0, 1) * 255).astype(np.uint8) | |
| results["Mask"] = Image.fromarray(alpha) | |
| if "Segmentation" in render_keys: | |
| seg_labels = camera.get_picture("Segmentation") | |
| label0 = seg_labels[..., 0].astype(np.uint8) | |
| seg_color = COLOR_PALETTE[label0] | |
| results["Segmentation"] = Image.fromarray(seg_color) | |
| if "Foreground" in render_keys: | |
| seg_labels = camera.get_picture("Segmentation") | |
| label0 = seg_labels[..., 0] | |
| mask = np.where((label0 > 1), 255, 0).astype(np.uint8) | |
| color = camera.get_picture("Color") | |
| color_rgb = (np.clip(color[..., :3], 0, 1) * 255).astype(np.uint8) | |
| foreground = np.concatenate([color_rgb, mask[..., None]], axis=-1) | |
| results["Foreground"] = Image.fromarray(foreground) | |
| if "Normal" in render_keys: | |
| normal = camera.get_picture("Normal")[..., :3] | |
| normal_img = (((normal + 1) / 2) * 255).astype(np.uint8) | |
| results["Normal"] = Image.fromarray(normal_img) | |
| if "Depth" in render_keys: | |
| position_map = camera.get_picture("Position") | |
| depth = -position_map[..., 2] | |
| alpha = torch.tensor(color[..., 3], dtype=torch.float32) | |
| norm_depth = DiffrastRender.normalize_map_by_mask( | |
| torch.tensor(depth), alpha | |
| ) | |
| depth_img = (norm_depth * 255).to(torch.uint8).numpy() | |
| results["Depth"] = Image.fromarray(depth_img) | |
| return results | |
| def step_action( | |
| self, | |
| agent: GripperActionAgent | None, | |
| action: np.ndarray | torch.Tensor | None, | |
| render_keys: list[str], | |
| sim_steps_per_control: int = 1, | |
| render: bool = True, | |
| ) -> dict: | |
| if action is not None: | |
| agent.set_action(action) | |
| frames = defaultdict(defaultdict) | |
| for _ in range(sim_steps_per_control): | |
| self.scene.step() | |
| if not render: | |
| return frames | |
| self.scene.update_render() | |
| for camera in self.cameras: | |
| camera.take_picture() | |
| images = self.render_images(camera, render_keys=render_keys) | |
| frames[camera.name].update(images) | |
| return frames | |
| class EvalGrasps: | |
| def __init__( | |
| self, | |
| cfg: EvalGraspsConfig, | |
| mesh_info: MeshInfo, | |
| gripper_control: GripperControlInfo, | |
| grasp: dict, | |
| ) -> None: | |
| self.cfg = cfg | |
| self.mesh_info = mesh_info | |
| self.gripper_control = gripper_control | |
| self.grasp = grasp | |
| self.init_scene() | |
| if self.cfg.debug_mode: | |
| self.setup_cameras(mesh_info) | |
| self.object_actor = self.add_object(mesh_info) | |
| self.frames = {} | |
| self.step_index = 0 | |
| self.world_x_offset = np.array( | |
| [self.cfg.sweep_distance, 0.0, 0.0], | |
| dtype=np.float64, | |
| ) | |
| self.world_y_offset = np.array( | |
| [0.0, self.cfg.sweep_distance, 0.0], | |
| dtype=np.float64, | |
| ) | |
| def should_render(self) -> bool: | |
| return ( | |
| self.cfg.debug_mode | |
| and self.step_index % max(1, self.cfg.render_interval) == 0 | |
| ) | |
| def init_scene(self) -> None: | |
| self.scene_manager = SapienSceneManager( | |
| sim_freq=self.cfg.sim_freq, | |
| ray_tracing=self.cfg.ray_tracing, | |
| device=self.cfg.sim_backend, | |
| ) | |
| if self.cfg.debug_mode: | |
| self.scene_manager.set_ground_color([0.78, 0.90, 0.72, 1.0]) | |
| def setup_cameras(self, mesh_info: MeshInfo) -> None: | |
| bounds = mesh_info.transformed_mesh.bounds.copy() | |
| mesh_size = bounds[1] - bounds[0] | |
| mesh_height = max(float(mesh_size[2]), 0.5) | |
| target_pt = [ | |
| float((bounds[0, 0] + bounds[1, 0]) * 0.5), | |
| float((bounds[0, 1] + bounds[1, 1]) * 0.5), | |
| float(bounds[1, 2]), | |
| ] | |
| aspect = self.cfg.image_hw[1] / self.cfg.image_hw[0] | |
| half_vfov = np.deg2rad(self.cfg.fovy_deg) * 0.5 | |
| half_hfov = np.arctan(np.tan(half_vfov) * aspect) | |
| half_min_fov = min(half_vfov, half_hfov) | |
| framing_half_extents = np.array( | |
| [ | |
| max(float(mesh_size[0]) * 0.5, 1e-6), | |
| max(float(mesh_size[1]) * 0.5, 1e-6), | |
| mesh_height, | |
| ], | |
| dtype=np.float64, | |
| ) | |
| camera_radius = max( | |
| float( | |
| np.linalg.norm(framing_half_extents) | |
| / np.sin(half_min_fov) | |
| * 1.5 | |
| ), | |
| 2.5, | |
| ) | |
| camera_height = target_pt[2] | |
| self.scene_manager.add_ring_cameras( | |
| num_cameras=self.cfg.num_cameras, | |
| radius=camera_radius, | |
| height=camera_height, | |
| target_pt=target_pt, | |
| image_hw=self.cfg.image_hw, | |
| fovy_deg=self.cfg.fovy_deg, | |
| ) | |
| def add_object( | |
| self, | |
| mesh_info: MeshInfo, | |
| ) -> sapien.Entity: | |
| spawn_z = self.cfg.object_bottom_clearance - float( | |
| mesh_info.transformed_mesh.bounds[0, 2] | |
| ) | |
| actor = self.scene_manager.add_actor(mesh_info, [0.0, 0.0, spawn_z]) | |
| self.set_object_gravity(not self.cfg.hold_object_during_close, actor) | |
| return actor | |
| def set_object_gravity( | |
| self, | |
| enabled: bool, | |
| actor: sapien.Entity | None = None, | |
| ) -> None: | |
| actor = self.object_actor if actor is None else actor | |
| for component in actor.components: | |
| set_disable_gravity = getattr( | |
| component, | |
| "set_disable_gravity", | |
| None, | |
| ) | |
| if set_disable_gravity is not None: | |
| set_disable_gravity(not enabled) | |
| if enabled: | |
| wake_up = getattr(component, "wake_up", None) | |
| if wake_up is not None: | |
| wake_up() | |
| def run_actions( | |
| self, | |
| agent: GripperActionAgent, | |
| actions: np.ndarray, | |
| ) -> None: | |
| for action in actions: | |
| this_frames = self.scene_manager.step_action( | |
| agent=agent, | |
| action=action, | |
| render_keys=["Color"], | |
| sim_steps_per_control=1, | |
| render=self.should_render(), | |
| ) | |
| self.update_frames(this_frames) | |
| self.step_index += 1 | |
| def make_pose( | |
| position: np.ndarray, | |
| quat_wxyz: np.ndarray, | |
| ) -> sapien.Pose: | |
| return sapien.Pose( | |
| p=[float(v) for v in position], | |
| q=[float(v) for v in quat_wxyz], | |
| ) | |
| def pose_matrix(pose: sapien.Pose) -> np.ndarray: | |
| quat_wxyz = np.asarray(pose.q, dtype=np.float64) | |
| quat_norm = np.linalg.norm(quat_wxyz) | |
| if quat_norm <= 0.0: | |
| raise ValueError("Pose quaternion has zero norm.") | |
| quat_wxyz = quat_wxyz / quat_norm | |
| matrix = np.eye(4, dtype=np.float64) | |
| matrix[:3, :3] = R.from_quat( | |
| [ | |
| quat_wxyz[1], | |
| quat_wxyz[2], | |
| quat_wxyz[3], | |
| quat_wxyz[0], | |
| ] | |
| ).as_matrix() | |
| matrix[:3, 3] = np.asarray(pose.p, dtype=np.float64) | |
| return matrix | |
| def pose_delta( | |
| reference_relative_pose: np.ndarray, | |
| current_relative_pose: np.ndarray, | |
| ) -> dict[str, float]: | |
| delta_pose = ( | |
| np.linalg.inv(reference_relative_pose) @ current_relative_pose | |
| ) | |
| translation_m = float(np.linalg.norm(delta_pose[:3, 3])) | |
| rotation_rad = float(R.from_matrix(delta_pose[:3, :3]).magnitude()) | |
| return { | |
| "translation_m": translation_m, | |
| "rotation_rad": rotation_rad, | |
| "rotation_deg": float(np.rad2deg(rotation_rad)), | |
| } | |
| def object_rel_pose( | |
| self, | |
| agent: GripperActionAgent, | |
| ) -> np.ndarray: | |
| gripper_pose = self.pose_matrix(agent.gripper.get_root_pose()) | |
| object_pose = self.pose_matrix(self.object_actor.get_pose()) | |
| return np.linalg.inv(gripper_pose) @ object_pose | |
| def object_bottom_z(self) -> float: | |
| object_pose = self.pose_matrix(self.object_actor.get_pose()) | |
| vertices = np.asarray(self.mesh_info.transformed_mesh.vertices) | |
| world_vertices = (object_pose[:3, :3] @ vertices.T).T + object_pose[ | |
| :3, 3 | |
| ] | |
| return float(np.min(world_vertices[:, 2])) | |
| def check_slip( | |
| self, | |
| agent: GripperActionAgent, | |
| reference_relative_pose: np.ndarray, | |
| ) -> dict[str, float | bool]: | |
| delta = self.pose_delta( | |
| reference_relative_pose, | |
| self.object_rel_pose(agent), | |
| ) | |
| slipped = ( | |
| delta["translation_m"] > self.cfg.slip_translation_threshold | |
| or delta["rotation_deg"] > self.cfg.slip_rotation_threshold_deg | |
| ) | |
| return { | |
| **delta, | |
| "slipped": slipped, | |
| } | |
| def checked_action( | |
| self, | |
| step_name: str, | |
| agent: GripperActionAgent, | |
| action_fn, | |
| *, | |
| check_slip: bool = True, | |
| ) -> dict: | |
| reference_relative_pose = self.object_rel_pose(agent) | |
| action_fn() | |
| self.wait_scene(self.cfg.action_wait_seconds) | |
| delta = self.check_slip( | |
| agent, | |
| reference_relative_pose, | |
| ) | |
| slipped = bool(delta["slipped"]) | |
| success = not slipped if check_slip else True | |
| return { | |
| "step": step_name, | |
| "success": success, | |
| "checked": check_slip, | |
| "slipped": slipped, | |
| "reason": None if success else "relative_pose_slip", | |
| "relative_pose_delta": { | |
| "translation_m": delta["translation_m"], | |
| "rotation_rad": delta["rotation_rad"], | |
| "rotation_deg": delta["rotation_deg"], | |
| }, | |
| } | |
| def get_steps( | |
| self, | |
| distance: float, | |
| speed: float, | |
| ) -> int: | |
| if distance <= 0.0: | |
| return 2 | |
| duration = distance / speed | |
| return max(2, int(np.ceil(duration * self.cfg.sim_freq))) | |
| def move_to_pose( | |
| self, | |
| agent: GripperActionAgent, | |
| start_pose: sapien.Pose, | |
| end_pose: sapien.Pose, | |
| finger_qpos: float | None = None, | |
| ) -> None: | |
| if finger_qpos is None: | |
| finger_qpos = self.gripper_control.close_qpos | |
| start_pos = np.asarray(start_pose.p, dtype=np.float64) | |
| end_pos = np.asarray(end_pose.p, dtype=np.float64) | |
| travel_distance = float(np.linalg.norm(end_pos - start_pos)) | |
| num_steps = self.get_steps( | |
| travel_distance, | |
| self.cfg.gripper_move_speed, | |
| ) | |
| alphas = np.linspace(0.0, 1.0, num_steps) | |
| quat = np.asarray(start_pose.q, dtype=np.float64) | |
| actions = [] | |
| for alpha in alphas: | |
| pose = self.make_pose( | |
| (1.0 - alpha) * start_pos + alpha * end_pos, | |
| quat, | |
| ) | |
| actions.append(self.make_action(pose, finger_qpos)) | |
| self.run_actions(agent, np.stack(actions, axis=0)) | |
| def move_gripper( | |
| self, | |
| agent: GripperActionAgent, | |
| offset: np.ndarray, | |
| ) -> None: | |
| start_pose = agent.gripper.get_root_pose() | |
| end_pose = self.make_pose( | |
| np.asarray(start_pose.p, dtype=np.float64) | |
| + np.asarray(offset, dtype=np.float64), | |
| np.asarray(start_pose.q, dtype=np.float64), | |
| ) | |
| self.move_to_pose(agent, start_pose, end_pose) | |
| def init_grasp_pose(self) -> sapien.Pose: | |
| orientation = self.grasp["orientation"] | |
| self.grasp = sapien.Pose( | |
| p=[float(v) for v in self.grasp["position"]], | |
| q=[ | |
| float(orientation["w"]), | |
| float(orientation["xyz"][0]), | |
| float(orientation["xyz"][1]), | |
| float(orientation["xyz"][2]), | |
| ], | |
| ) | |
| gripper_pose = self.object_actor.get_pose() * self.grasp | |
| return gripper_pose | |
| def make_action( | |
| self, | |
| pose: sapien.Pose, | |
| finger_qpos: float, | |
| ) -> np.ndarray: | |
| return np.array( | |
| [ | |
| float(pose.p[0]), | |
| float(pose.p[1]), | |
| float(pose.p[2]), | |
| float(pose.q[0]), | |
| float(pose.q[1]), | |
| float(pose.q[2]), | |
| float(pose.q[3]), | |
| float(finger_qpos), | |
| ], | |
| dtype=np.float32, | |
| ) | |
| def add_gripper_agent(self) -> GripperActionAgent: | |
| gripper_pose = self.init_grasp_pose() | |
| gripper, root_target, root_drive = self.scene_manager.add_gripper( | |
| self.cfg.gripper_urdf_path, | |
| gripper_pose, | |
| self.gripper_control.open_qpos, | |
| self.cfg.gripper_drive_stiffness, | |
| self.gripper_control.drive_damping, | |
| self.gripper_control.force_limit, | |
| self.cfg.gripper_static_friction, | |
| self.cfg.gripper_dynamic_friction, | |
| self.cfg.gripper_restitution, | |
| self.cfg.root_drive_stiffness, | |
| self.cfg.root_drive_damping, | |
| ) | |
| return GripperActionAgent( | |
| self.scene_manager, | |
| gripper, | |
| root_target, | |
| root_drive, | |
| ) | |
| def update_frames(self, this_frames: dict) -> None: | |
| for cam_name, images_list in this_frames.items(): | |
| if cam_name not in self.frames: | |
| self.frames[cam_name] = [] | |
| self.frames[cam_name].append(np.array(images_list["Color"])) | |
| def wait_scene(self, seconds: float) -> None: | |
| num_steps = int(self.cfg.sim_freq * seconds) | |
| for _ in range(num_steps): | |
| this_frames = self.scene_manager.step_action( | |
| agent=None, | |
| action=None, | |
| render_keys=["Color"], | |
| sim_steps_per_control=1, | |
| render=self.should_render(), | |
| ) | |
| self.update_frames(this_frames) | |
| self.step_index += 1 | |
| def close_gripper(self, agent: GripperActionAgent) -> None: | |
| num_close_steps = self.get_steps( | |
| abs( | |
| self.gripper_control.open_qpos | |
| - self.gripper_control.close_qpos | |
| ), | |
| self.cfg.gripper_close_speed, | |
| ) | |
| root_pose = agent.gripper.get_root_pose() | |
| actions = [ | |
| self.make_action( | |
| root_pose, | |
| (1.0 - alpha) * self.gripper_control.open_qpos | |
| + alpha * self.gripper_control.close_qpos, | |
| ) | |
| for alpha in np.linspace(0.0, 1.0, max(1, num_close_steps)) | |
| ] | |
| self.run_actions(agent, np.stack(actions, axis=0)) | |
| def summarize_steps(step_results: list[dict]) -> dict: | |
| failure = next( | |
| ( | |
| step_result | |
| for step_result in step_results | |
| if not step_result["success"] | |
| ), | |
| None, | |
| ) | |
| result = { | |
| "success": failure is None, | |
| "relative_pose_deltas": { | |
| step_result["step"]: step_result["relative_pose_delta"] | |
| for step_result in step_results | |
| }, | |
| } | |
| if failure is not None: | |
| result.update( | |
| { | |
| "failure_step": failure["step"], | |
| "reason": failure["reason"], | |
| } | |
| ) | |
| return result | |
| def run(self) -> dict: | |
| step_results = [] | |
| self.wait_scene(2.0) | |
| this_agent = self.add_gripper_agent() | |
| self.wait_scene(self.cfg.action_wait_seconds) | |
| # Close the gripper around the object. | |
| step_result = self.checked_action( | |
| "close_gripper", | |
| this_agent, | |
| lambda: self.close_gripper(this_agent), | |
| check_slip=self.cfg.check_close_slip, | |
| ) | |
| step_results.append(step_result) | |
| if not step_result["success"]: | |
| return self.summarize_steps(step_results) | |
| if self.cfg.hold_object_during_close: | |
| self.set_object_gravity(True) | |
| # Lift the grasped object upward. | |
| step_result = self.checked_action( | |
| "lift_gripper", | |
| this_agent, | |
| lambda: self.move_gripper( | |
| this_agent, | |
| np.array( | |
| [0.0, 0.0, self.cfg.lift_height], | |
| dtype=np.float64, | |
| ), | |
| ), | |
| ) | |
| step_results.append(step_result) | |
| if not step_result["success"]: | |
| return self.summarize_steps(step_results) | |
| shake_start_pose = this_agent.gripper.get_root_pose() | |
| shake_end_pose = self.make_pose( | |
| np.asarray(shake_start_pose.p, dtype=np.float64) | |
| + self.world_x_offset, | |
| np.asarray(shake_start_pose.q, dtype=np.float64), | |
| ) | |
| # Move the gripper outward along the X axis. | |
| step_result = self.checked_action( | |
| "move_gripper_out_x", | |
| this_agent, | |
| lambda: self.move_to_pose( | |
| this_agent, | |
| shake_start_pose, | |
| shake_end_pose, | |
| ), | |
| ) | |
| step_results.append(step_result) | |
| if not step_result["success"]: | |
| return self.summarize_steps(step_results) | |
| # Move the gripper back from the X-axis offset. | |
| step_result = self.checked_action( | |
| "move_gripper_back_x", | |
| this_agent, | |
| lambda: self.move_to_pose( | |
| this_agent, | |
| this_agent.gripper.get_root_pose(), | |
| shake_start_pose, | |
| ), | |
| ) | |
| step_results.append(step_result) | |
| if not step_result["success"]: | |
| return self.summarize_steps(step_results) | |
| shake_start_pose = this_agent.gripper.get_root_pose() | |
| shake_end_pose = self.make_pose( | |
| np.asarray(shake_start_pose.p, dtype=np.float64) | |
| + self.world_y_offset, | |
| np.asarray(shake_start_pose.q, dtype=np.float64), | |
| ) | |
| # Move the gripper outward along the Y axis. | |
| step_result = self.checked_action( | |
| "move_gripper_out_y", | |
| this_agent, | |
| lambda: self.move_to_pose( | |
| this_agent, | |
| shake_start_pose, | |
| shake_end_pose, | |
| ), | |
| ) | |
| step_results.append(step_result) | |
| if not step_result["success"]: | |
| return self.summarize_steps(step_results) | |
| # Move the gripper back from the Y-axis offset. | |
| step_result = self.checked_action( | |
| "move_gripper_back_y", | |
| this_agent, | |
| lambda: self.move_to_pose( | |
| this_agent, | |
| this_agent.gripper.get_root_pose(), | |
| shake_start_pose, | |
| ), | |
| ) | |
| step_results.append(step_result) | |
| if not step_result["success"]: | |
| return self.summarize_steps(step_results) | |
| # Lower the grasped object back down. | |
| descend_offset = np.array( | |
| [0.0, 0.0, 0.01 - self.object_bottom_z()], | |
| dtype=np.float64, | |
| ) | |
| step_result = self.checked_action( | |
| "descend_gripper", | |
| this_agent, | |
| lambda: self.move_gripper(this_agent, descend_offset), | |
| ) | |
| step_results.append(step_result) | |
| if not step_result["success"]: | |
| return self.summarize_steps(step_results) | |
| return self.summarize_steps(step_results) | |
| def export_video(self, output_dir: str) -> None: | |
| if not self.cfg.debug_mode: | |
| return | |
| os.makedirs(output_dir, exist_ok=True) | |
| for cam_name, this_frames in self.frames.items(): | |
| this_output_path = os.path.join(output_dir, f"{cam_name}.mp4") | |
| imageio.mimsave( | |
| this_output_path, this_frames, fps=self.cfg.video_fps | |
| ) | |
| def _validate_config(cfg: EvalGraspsConfig) -> None: | |
| if not cfg.urdf_paths: | |
| raise ValueError("urdf_paths must be provided.") | |
| if len(cfg.output_dirs) == 0: | |
| cfg.output_dirs = [ | |
| os.path.join(os.path.dirname(path), "affordance") | |
| for path in cfg.urdf_paths | |
| ] | |
| if len(cfg.urdf_paths) != len(cfg.output_dirs): | |
| raise ValueError( | |
| "urdf_paths and output_dirs must have the same length, " | |
| f"got {len(cfg.urdf_paths)} and {len(cfg.output_dirs)}." | |
| ) | |
| def _set_single_config( | |
| cfg: EvalGraspsConfig, | |
| urdf_path: str, | |
| output_dir: str, | |
| ) -> EvalGraspsConfig: | |
| return dataclasses.replace( | |
| cfg, | |
| urdf_paths=[urdf_path], | |
| output_dirs=[output_dir], | |
| ) | |
| def _eval_grasp( | |
| cfg: EvalGraspsConfig, | |
| grasp_id: str, | |
| grasp: dict, | |
| ) -> dict: | |
| output_dir = cfg.output_dirs[0] | |
| mesh_info = load_mesh_info(cfg.urdf_paths[0]) | |
| gripper_control = GripperControlInfo( | |
| **URDFFile(cfg.gripper_urdf_path).get_prismatic_joint_control_info() | |
| ) | |
| evaluator = EvalGrasps(cfg, mesh_info, gripper_control, grasp) | |
| result = evaluator.run() | |
| evaluator.export_video( | |
| os.path.join(output_dir, "eval_grasp_renders", grasp_id) | |
| ) | |
| if threading.current_thread() is not threading.main_thread(): | |
| with _EVALUATOR_KEEPALIVE_LOCK: | |
| _EVALUATOR_KEEPALIVE.append(evaluator) | |
| return { | |
| "grasp_id": grasp_id, | |
| "output_dir": output_dir, | |
| "simulation_result": result, | |
| } | |
| def _release_thread_evaluators() -> None: | |
| with _EVALUATOR_KEEPALIVE_LOCK: | |
| evaluators = list(_EVALUATOR_KEEPALIVE) | |
| _EVALUATOR_KEEPALIVE.clear() | |
| for evaluator in evaluators: | |
| del evaluator | |
| def _eval_grasp_task(task: tuple[EvalGraspsConfig, str, dict]) -> dict: | |
| cfg, grasp_id, grasp = task | |
| return _eval_grasp(cfg, grasp_id, grasp) | |
| def _slip_thresholds(cfg: EvalGraspsConfig) -> dict[str, float]: | |
| return { | |
| "translation_m": cfg.slip_translation_threshold, | |
| "rotation_deg": cfg.slip_rotation_threshold_deg, | |
| } | |
| def _is_successful_grasp_result(result: dict) -> bool: | |
| simulation_result = result.get("simulation_result", {}) | |
| if not isinstance(simulation_result, dict): | |
| return False | |
| return bool(simulation_result.get("success", False)) | |
| def _grasp_success_by_id(results: list[dict]) -> dict[str, bool]: | |
| return { | |
| result["grasp_id"]: _is_successful_grasp_result(result) | |
| for result in results | |
| } | |
| def _build_eval_summary( | |
| cfg: EvalGraspsConfig, | |
| results: list[dict], | |
| elapsed_seconds: float, | |
| ) -> dict: | |
| successful_grasp_ids = [ | |
| result["grasp_id"] | |
| for result in results | |
| if _is_successful_grasp_result(result) | |
| ] | |
| failed_grasp_ids = [ | |
| result["grasp_id"] | |
| for result in results | |
| if not _is_successful_grasp_result(result) | |
| ] | |
| num_grasps = len(results) | |
| num_success = len(successful_grasp_ids) | |
| return { | |
| "num_grasps": num_grasps, | |
| "num_success": num_success, | |
| "num_failure": len(failed_grasp_ids), | |
| "success_rate": num_success / max(1, num_grasps), | |
| "successful_grasp_ids": successful_grasp_ids, | |
| "failed_grasp_ids": failed_grasp_ids, | |
| "slip_thresholds": _slip_thresholds(cfg), | |
| "elapsed_seconds": elapsed_seconds, | |
| } | |
| def _run_single_eval(cfg: EvalGraspsConfig) -> None: | |
| start_time = time.perf_counter() | |
| output_dir = cfg.output_dirs[0] | |
| os.makedirs(output_dir, exist_ok=True) | |
| urdf_path = cfg.urdf_paths[0] | |
| urdf = URDFFile(urdf_path) | |
| affordance_path = urdf.get_affordance_annot_path() | |
| affordance_payload = load_json(affordance_path) | |
| affordances = affordance_payload.get("affordances", []) | |
| grasps = {} | |
| for affordance in affordances: | |
| grasps.update(affordance.get("grasp_group", {})) | |
| if not grasps: | |
| return | |
| num_workers = cfg.num_workers | |
| if num_workers is None: | |
| num_workers = min(len(grasps), os.cpu_count() or 1, MAX_NUM_WORKERS) | |
| else: | |
| num_workers = min(num_workers, len(grasps), MAX_NUM_WORKERS) | |
| logger.info( | |
| f"Running {len(grasps)} grasp poses with {num_workers} worker(s)" | |
| ) | |
| results = [] | |
| if num_workers == 1: | |
| for grasp_id, grasp in tqdm(grasps.items(), desc="Evaluating grasps"): | |
| results.append(_eval_grasp(cfg, grasp_id, grasp)) | |
| else: | |
| worker_cfg = dataclasses.replace(cfg, num_workers=1) | |
| tasks = [ | |
| (worker_cfg, grasp_id, grasp) for grasp_id, grasp in grasps.items() | |
| ] | |
| try: | |
| with concurrent.futures.ThreadPoolExecutor( | |
| max_workers=num_workers | |
| ) as executor: | |
| futures = [ | |
| executor.submit(_eval_grasp_task, task) for task in tasks | |
| ] | |
| for future in tqdm( | |
| concurrent.futures.as_completed(futures), | |
| total=len(futures), | |
| desc="Evaluating grasps", | |
| ): | |
| results.append(future.result()) | |
| results.sort( | |
| key=lambda result: list(grasps.keys()).index( | |
| result["grasp_id"] | |
| ) | |
| ) | |
| finally: | |
| _release_thread_evaluators() | |
| elapsed_seconds = time.perf_counter() - start_time | |
| summary = _build_eval_summary(cfg, results, elapsed_seconds) | |
| failed_grasp_ids = set(summary["failed_grasp_ids"]) | |
| grasp_success_by_id = _grasp_success_by_id(results) | |
| for affordance in affordance_payload.get("affordances", []): | |
| grasp_group = affordance.get("grasp_group") | |
| if isinstance(grasp_group, dict): | |
| if cfg.debug_mode or cfg.save_failed_grasps: | |
| for grasp_id, grasp in grasp_group.items(): | |
| if ( | |
| isinstance(grasp, dict) | |
| and grasp_id in grasp_success_by_id | |
| ): | |
| grasp["simulation_success"] = grasp_success_by_id[ | |
| grasp_id | |
| ] | |
| continue | |
| for failed_grasp_id in failed_grasp_ids: | |
| grasp_group.pop(failed_grasp_id, None) | |
| write_json(affordance_payload, affordance_path) | |
| logger.info(f"Grasp evaluation simulation took {elapsed_seconds:.2f}s") | |
| logger.info(f"Updated grasp simulation results in {affordance_path}") | |
| logger.info( | |
| f"Grasp success rate: {summary['num_success']}/" | |
| f"{summary['num_grasps']} ({summary['success_rate']:.2%})" | |
| ) | |
| def _run_single_eval_safe(cfg: EvalGraspsConfig) -> bool: | |
| urdf_path = cfg.urdf_paths[0] | |
| try: | |
| _run_single_eval(cfg) | |
| except Exception as exc: | |
| logger.error( | |
| "Grasp evaluation failed for URDF {}: {}".format( | |
| urdf_path, | |
| exc, | |
| ) | |
| ) | |
| return False | |
| return True | |
| def run_eval(cfg: EvalGraspsConfig) -> None: | |
| _validate_config(cfg) | |
| for idx, (urdf_path, output_dir) in enumerate( | |
| zip(cfg.urdf_paths, cfg.output_dirs), | |
| start=1, | |
| ): | |
| single_cfg = _set_single_config(cfg, urdf_path, output_dir) | |
| logger.info( | |
| f"Running grasp evaluation for URDF {idx}/{len(cfg.urdf_paths)}: " | |
| f"{urdf_path}" | |
| ) | |
| _run_single_eval_safe(single_cfg) | |
| def entrypoint(*args, **kwargs) -> None: | |
| cfg = tyro.cli(EvalGraspsConfig) | |
| run_eval(cfg) | |
| if __name__ == "__main__": | |
| entrypoint() | |