diff --git a/scripts/tutorials/compute_viewpoint.py b/scripts/tutorials/compute_viewpoint.py new file mode 100644 index 00000000..4d30713c --- /dev/null +++ b/scripts/tutorials/compute_viewpoint.py @@ -0,0 +1,59 @@ +import argparse + +import numpy as np +from scipy.spatial.transform import Rotation as R + + +def compute_lookat(eye, euler_deg, distance=2.0): + """Compute lookat point from eye position and euler angles. + + Args: + eye: Camera eye position [x, y, z] + euler_deg: Camera rotation in degrees [roll, pitch, yaw] + distance: Distance from eye to lookat point (default: 2.0) + """ + eye = np.array(eye) + + rot = R.from_euler("xyz", euler_deg, degrees=True) + rot_matrix = rot.as_matrix() + + forward = rot_matrix @ np.array([0, 0, -1]) + + lookat = eye + forward * distance + + print(f" self.viewer.eye = ({eye[0]:.5f}, {eye[1]:.5f}, {eye[2]:.5f})") + print(f" self.viewer.lookat = ({lookat[0]:.5f}, {lookat[1]:.5f}, {lookat[2]:.5f})") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Compute camera lookat point from eye position and euler angles", + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + + parser.add_argument( + "--eye", + type=float, + nargs=3, + required=True, + metavar=("X", "Y", "Z"), + help="Camera eye position [x, y, z]", + ) + parser.add_argument( + "--euler", + type=float, + nargs=3, + required=True, + metavar=("ROLL", "PITCH", "YAW"), + help="Camera rotation in degrees [roll, pitch, yaw]", + ) + parser.add_argument( + "--distance", + type=float, + default=2.0, + help="Distance from eye to lookat point (default: 2.0)", + ) + + args = parser.parse_args() + + compute_lookat(args.eye, args.euler, distance=args.distance) diff --git a/scripts/tutorials/marble_compose.py b/scripts/tutorials/marble_compose.py index 5f32ca5c..9d7dd494 100644 --- a/scripts/tutorials/marble_compose.py +++ b/scripts/tutorials/marble_compose.py @@ -8,6 +8,7 @@ """ import argparse +import os from pathlib import Path import numpy as np @@ -28,6 +29,12 @@ "cube": { "single": "LeIsaac-SO101-LiftCube-v0", }, + "burger": { + "dual": "LeIsaac-SO101-AssembleHamburger-BiArm-v0", + }, + "sausage": { + "dual": "LeIsaac-SO101-SausageCut-BiArm-Direct-v0", + }, } TASK_CONFIG = { @@ -58,7 +65,7 @@ ], "source_scene": "scenes/lightwheel_toyroom/scene.usd", "assets_subpath": "scenes/lightwheel_toyroom/Assets", - "table_name": "KidRoom_Table01", + "platform_name": "KidRoom_Table01", }, "orange": { "objects": ["Orange001", "Orange002", "Orange003", "Plate"], @@ -69,13 +76,32 @@ "objects": ["Table038_01", "cloth"], "source_scene": "scenes/lightwheel_bedroom/scene.usd", "assets_subpath": "scenes/lightwheel_bedroom", - "table_name": "Table038_01", + "platform_name": "Table038_01", }, "cube": { "objects": ["cube"], "source_scene": "scenes/table_with_cube/scene.usd", "assets_subpath": "scenes/table_with_cube/cube", }, + "burger": { + "objects": [ + "Burger_ChoppingBlock", + "Burger_Plate", + "Burger_Beef_Patties001", + "Burger_Cheese001", + "Burger_Bread002", + ], + "source_scene": "scenes/kitchen_with_burger/scene.usd", + "assets_subpath": "scenes/kitchen_with_burger/objects/burger/Assets", + "parent_prim": "Burger", + "platform_name": "Burger_ChoppingBlock", + }, + "sausage": { + "objects": ["ChoppingBlock"], + "source_scene": "scenes/kitchen_with_sausage/scene.usd", + "assets_subpath": "scenes/kitchen_with_sausage/objects", + "platform_name": "ChoppingBlock", + }, } @@ -135,20 +161,16 @@ def compute_scene_transform(orig_pos, orig_quat, target_pos, target_quat): def get_object_usd_path(task_type: str, obj_name: str, assets_base: str) -> str: """Resolve USD file path for an object""" - config = TASK_CONFIG[task_type] - table_name = config.get("table_name") - if task_type == "toys": - if obj_name == table_name: + platform_name = TASK_CONFIG[task_type].get("platform_name") + if obj_name == platform_name: return f"{assets_base}/{obj_name}/{obj_name}.usd" name = obj_name[:-3] if obj_name.endswith("_01") else obj_name return f"{assets_base}/Kit1/{name}.usd" - if task_type == "orange": - return f"{assets_base}/{obj_name}/{obj_name}.usd" - if task_type == "cloth": - if obj_name == table_name: + platform_name = TASK_CONFIG[task_type].get("platform_name") + if obj_name == platform_name: folder = obj_name[:-3] if obj_name.endswith("_01") else obj_name return f"{assets_base}/LW_Loft/Loft/{folder}/{folder}.usd" return f"{assets_base}/cloth/cloth.usd" @@ -156,10 +178,16 @@ def get_object_usd_path(task_type: str, obj_name: str, assets_base: str) -> str: if task_type == "cube": return f"{assets_base}/cube.usd" - raise ValueError(f"Unknown object: {obj_name}") + # Common pattern: orange, burger, sausage + if task_type in ("orange", "burger", "sausage"): + return f"{assets_base}/{obj_name}/{obj_name}.usd" + + raise ValueError(f"Unknown task type: {task_type}") -def read_layout_from_usd(usd_path: str, object_names: list[str]) -> dict[str, dict]: +def read_layout_from_usd( + usd_path: str, object_names: list[str], parent_prim_name: str | None = None +) -> dict[str, dict]: """Read object poses from USD (first-level children of root prim)""" from pxr import Usd, UsdGeom @@ -167,7 +195,12 @@ def read_layout_from_usd(usd_path: str, object_names: list[str]) -> dict[str, di if not stage: raise RuntimeError(f"Cannot open: {usd_path}") - root = stage.GetDefaultPrim() or stage.GetPrimAtPath("/") + if parent_prim_name: + root = stage.GetPrimAtPath(f"/World/{parent_prim_name}") + if not root.IsValid(): + raise RuntimeError(f"Parent prim '{parent_prim_name}' not found in {usd_path}") + else: + root = stage.GetDefaultPrim() or stage.GetPrimAtPath("/") object_set = set(object_names) layout = {} @@ -210,13 +243,11 @@ def load_robot_pose(task_type: str, use_dual_arm: bool = False) -> dict[str, lis if use_dual_arm: env_id = task_env.get("dual") if not env_id: - print(f"[WARN] Dual-arm not supported for '{task_type}', using single-arm") - env_id = task_env["single"] - use_dual_arm = False - else: - print(f"[INFO] Using dual-arm config: {env_id}") + raise ValueError(f"Dual-arm not supported for '{task_type}'") else: - env_id = task_env["single"] + env_id = task_env.get("single") + if not env_id: + raise ValueError(f"Single-arm not supported for '{task_type}'") env_cfg = parse_env_cfg(env_id, device="cpu", num_envs=1) @@ -240,7 +271,7 @@ def compose_scene( assets_base: str, target_pos: list[float], target_quat: list[float], - include_table: bool = False, + include_platform: bool = False, use_dual_arm: bool = False, ) -> str: """ @@ -249,32 +280,32 @@ def compose_scene( 2. Place objects at their original positions Args: - include_table: Include table in output (default: False) + include_platform: Include platform (table/board) in output (default: False) use_dual_arm: Use dual-arm configuration (toys/cloth only, uses left_arm as reference) """ from pxr import Gf, Usd, UsdGeom config = TASK_CONFIG[task_type] - table_name = config.get("table_name") + platform_name = config.get("platform_name") - # Warn if table not supported - if include_table and not table_name: - print(f"[WARN] Table not supported for '{task_type}'") - include_table = False + # Warn if platform not supported + if include_platform and not platform_name: + print(f"[WARN] Platform not supported for '{task_type}'") + include_platform = False # Build source USD path and read layout source_usd = f"{assets_base}/{config['source_scene']}" print(f"[INFO] Reading layout from: {source_usd}") - layout = read_layout_from_usd(source_usd, config["objects"]) + layout = read_layout_from_usd(source_usd, config["objects"], parent_prim_name=config.get("parent_prim")) # Compute scene transform based on reference point - if include_table and table_name: - # Use table as reference point - table_pose = layout.get(table_name) - if not table_pose: - raise RuntimeError(f"Table '{table_name}' not found in source USD") - orig_pos, orig_quat = table_pose["pos"], table_pose["rot"] - print(f"[INFO] Using table '{table_name}' as reference: pos={orig_pos}") + if include_platform and platform_name: + # Use platform as reference point + platform_pose = layout.get(platform_name) + if not platform_pose: + raise RuntimeError(f"Platform '{platform_name}' not found in source USD") + orig_pos, orig_quat = platform_pose["pos"], platform_pose["rot"] + print(f"[INFO] Using platform '{platform_name}' as reference: pos={orig_pos}") else: # Use robot as reference point (dual-arm uses left_arm) robot = load_robot_pose(task_type, use_dual_arm=use_dual_arm) @@ -283,18 +314,30 @@ def compose_scene( scene_pos, scene_quat = compute_scene_transform(orig_pos, orig_quat, target_pos, target_quat) - # Filter out table if not included - if not include_table and table_name: - layout = {k: v for k, v in layout.items() if k != table_name} + # Filter out platform if not included + if not include_platform and platform_name: + layout = {k: v for k, v in layout.items() if k != platform_name} # Create output stage stage = Usd.Stage.CreateNew(output_usd) world = UsdGeom.Xform.Define(stage, "/World") stage.SetDefaultPrim(world.GetPrim()) - # Add background with transform + # Helper to compute relative path from output USD to referenced USD + output_dir = Path(output_usd).resolve().parent + + def to_relative(abs_path: str) -> str: + try: + return os.path.relpath(abs_path, output_dir) + except ValueError: + # On Windows, relpath fails across drives - use absolute path + return abs_path + + # Add background with transform (use relative path) bg_prim = stage.DefinePrim("/World/Scene") - bg_prim.GetReferences().AddReference(background_usd) + bg_relative = to_relative(background_usd) + bg_prim.GetReferences().AddReference(bg_relative) + print(f"[INFO] Background reference: {bg_relative}") bg_xform = UsdGeom.Xformable(bg_prim) bg_xform.ClearXformOpOrder() bg_xform.AddTranslateOp().Set(Gf.Vec3d(*scene_pos)) @@ -302,19 +345,31 @@ def compose_scene( # Add objects at original positions assets_path = f"{assets_base}/{config['assets_subpath']}" + + # Determine parent path for objects (burger uses /World/Burger, others use /World) + parent_prim_name = config.get("parent_prim") + if parent_prim_name: + objects_parent = f"/World/{parent_prim_name}" + # Create parent Xform + UsdGeom.Xform.Define(stage, objects_parent) + print(f"[INFO] Created parent group: {objects_parent}") + else: + objects_parent = "/World" + for name, pose in layout.items(): usd_path = get_object_usd_path(task_type, name, assets_path) if not Path(usd_path).exists(): print(f"[WARN] {usd_path} not found, skipping") continue - prim = stage.DefinePrim(f"/World/{name}") - prim.GetReferences().AddReference(usd_path) + prim = stage.DefinePrim(f"{objects_parent}/{name}") + usd_relative = to_relative(usd_path) + prim.GetReferences().AddReference(usd_relative) xform = UsdGeom.Xformable(prim) xform.ClearXformOpOrder() xform.AddTranslateOp().Set(Gf.Vec3d(*pose["pos"])) xform.AddOrientOp().Set(Gf.Quatf(*pose["rot"])) - print(f"[OK] Added {name}") + print(f"[OK] Added {name} -> {usd_relative}") stage.Save() print(f"[OK] Saved: {output_usd}") @@ -342,18 +397,16 @@ def compose_scene( metavar=("W", "X", "Y", "Z"), help="Target robot quaternion", ) - parser.add_argument("--include-table", action="store_true", help="Include table in output") + parser.add_argument("--include-platform", action="store_true", help="Include working platform in output") parser.add_argument( "--dual-arm", action="store_true", - help="Use dual-arm configuration (toys/cloth only, uses left_arm as reference)", + help="Use dual-arm configuration (uses left_arm as reference)", ) args = parser.parse_args() # Set LEISAAC_ASSETS_ROOT to the provided assets base path - import os - os.environ["LEISAAC_ASSETS_ROOT"] = args.assets_base print(f"[INFO] Set LEISAAC_ASSETS_ROOT={args.assets_base}") @@ -374,7 +427,7 @@ def compose_scene( assets_base=args.assets_base, target_pos=args.target_pos, target_quat=args.target_quat, - include_table=args.include_table, + include_platform=args.include_platform, use_dual_arm=args.dual_arm, ) diff --git a/source/leisaac/leisaac/assets/robots/lerobot.py b/source/leisaac/leisaac/assets/robots/lerobot.py index 1d80101f..23d5ca03 100644 --- a/source/leisaac/leisaac/assets/robots/lerobot.py +++ b/source/leisaac/leisaac/assets/robots/lerobot.py @@ -7,6 +7,7 @@ """Configuration for the SO101 Follower Robot.""" SO101_FOLLOWER_ASSET_PATH = Path(ASSETS_ROOT) / "robots" / "so101_follower.usd" +SO101_KINFE_ASSET_PATH = Path(ASSETS_ROOT) / "robots" / "so101_knife.usd" SO101_FOLLOWER_CFG = ArticulationCfg( spawn=sim_utils.UsdFileCfg( @@ -52,6 +53,57 @@ soft_joint_pos_limit_factor=1.0, ) +SO101_KINFE_CFG = ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=str(SO101_KINFE_ASSET_PATH), + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, + ), + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=True, + solver_position_iteration_count=4, + solver_velocity_iteration_count=4, + fix_root_link=True, + ), + ), + init_state=ArticulationCfg.InitialStateCfg( + pos=(1.4, -2.3, 0), + rot=(0.0, 0.0, 0.0, 1.0), + joint_pos={ + "shoulder_pan": 0.0, + "shoulder_lift": 0.0, + "elbow_flex": 0.0, + "wrist_flex": 0.0, + "wrist_roll": 0.0, + "gripper": 0.0, + }, + ), + actuators={ + "sts3215-gripper": ImplicitActuatorCfg( + joint_names_expr=["gripper"], + effort_limit_sim=10, + velocity_limit_sim=10, + stiffness=17.8, + damping=0.60, + ), + "sts3215-arm": ImplicitActuatorCfg( + joint_names_expr=[ + "shoulder_pan", + "shoulder_lift", + "elbow_flex", + "wrist_flex", + "wrist_roll", + ], + effort_limit_sim=10, + velocity_limit_sim=10, + stiffness=17.8, + damping=0.60, + ), + }, + soft_joint_pos_limit_factor=1.0, +) + + # joint limit written in USD (degree) SO101_FOLLOWER_USD_JOINT_LIMLITS = { "shoulder_pan": (-110.0, 110.0), diff --git a/source/leisaac/leisaac/assets/scenes/kitchen.py b/source/leisaac/leisaac/assets/scenes/kitchen.py index 7a061ace..a8c80f46 100644 --- a/source/leisaac/leisaac/assets/scenes/kitchen.py +++ b/source/leisaac/leisaac/assets/scenes/kitchen.py @@ -15,10 +15,22 @@ ) ) -KITCHEN_WITH_HAMBURGER_USD_PATH = str(SCENES_ROOT / "kitchen_with_hamburger" / "scene.usd") +KITCHEN_WITH_HAMBURGER_USD_PATH = str(SCENES_ROOT / "kitchen_with_burger" / "scene.usd") KITCHEN_WITH_HAMBURGER_CFG = AssetBaseCfg( spawn=sim_utils.UsdFileCfg( usd_path=KITCHEN_WITH_HAMBURGER_USD_PATH, ) ) + + +KITCHEN_WITH_SAUSAGE_USD_PATH = str(SCENES_ROOT / "kitchen_with_sausage" / "scene.usd") + +KITCHEN_WITH_SAUSAGE_CFG = AssetBaseCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=KITCHEN_WITH_SAUSAGE_USD_PATH, + ) +) + +# Sausage object USD path +SAUSAGE_USD_PATH = str(SCENES_ROOT / "kitchen_with_sausage" / "objects" / "Sausage001" / "Sausage001.usd") diff --git a/source/leisaac/leisaac/assets/scenes/loftroom.py b/source/leisaac/leisaac/assets/scenes/loftroom.py new file mode 100644 index 00000000..8d2cb432 --- /dev/null +++ b/source/leisaac/leisaac/assets/scenes/loftroom.py @@ -0,0 +1,17 @@ +from pathlib import Path + +import isaaclab.sim as sim_utils +from isaaclab.assets import AssetBaseCfg +from leisaac.utils.constant import ASSETS_ROOT + +"""Configuration for the Loft Room Scene""" +SCENES_ROOT = Path(ASSETS_ROOT) / "scenes" + +# Base loft room scene (for sausage cutting - no sausage included) +LIGHTWHEEL_LOFTROOM_USD_PATH = str(SCENES_ROOT / "loftroom" / "scene.usd") + +LIGHTWHEEL_LOFTROOM_CFG = AssetBaseCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=LIGHTWHEEL_LOFTROOM_USD_PATH, + ) +) diff --git a/source/leisaac/leisaac/enhance/assets/__init__.py b/source/leisaac/leisaac/enhance/assets/__init__.py index 8f605fb1..bd92e895 100644 --- a/source/leisaac/leisaac/enhance/assets/__init__.py +++ b/source/leisaac/leisaac/enhance/assets/__init__.py @@ -1,2 +1,4 @@ from .cloth_object.cloth_object import ClothObject from .cloth_object.cloth_object_cfg import ClothObjectCfg +from .cuttable_object.cuttable_object import CuttableObject, SingleCuttableObject +from .cuttable_object.cuttable_object_cfg import CuttableObjectCfg diff --git a/source/leisaac/leisaac/enhance/assets/cuttable_object/__init__.py b/source/leisaac/leisaac/enhance/assets/cuttable_object/__init__.py new file mode 100644 index 00000000..b643999f --- /dev/null +++ b/source/leisaac/leisaac/enhance/assets/cuttable_object/__init__.py @@ -0,0 +1,4 @@ +"""Cuttable object asset module.""" + +from .cuttable_object import CuttableObject, SingleCuttableObject +from .cuttable_object_cfg import CuttableObjectCfg diff --git a/source/leisaac/leisaac/enhance/assets/cuttable_object/cuttable_object.py b/source/leisaac/leisaac/enhance/assets/cuttable_object/cuttable_object.py new file mode 100644 index 00000000..84283295 --- /dev/null +++ b/source/leisaac/leisaac/enhance/assets/cuttable_object/cuttable_object.py @@ -0,0 +1,163 @@ +"""Cuttable object asset for mesh cutting simulation.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import TYPE_CHECKING + +import isaaclab.sim as sim_utils +import omni +import torch +from isaacsim.core.utils.prims import delete_prim +from leisaac.utils.collision_checker import Collision_Checker +from leisaac.utils.cutMeshNode import cutMeshNode +from pxr import UsdGeom + +if TYPE_CHECKING: + from .cuttable_object_cfg import CuttableObjectCfg + + +class SingleCuttableObject: + """Single cuttable object that wraps cutMeshNode functionality.""" + + def __init__( + self, + prim_path: str, + usd_path: str, + mesh_subfix: str, + trigger_subfix: str, + knife_path: str, + base_pos: tuple[float, float, float], + base_quat: tuple[float, float, float, float], + ): + self._prim_path = prim_path + self._usd_path = usd_path + self._mesh_path = f"{prim_path}/{mesh_subfix}" if mesh_subfix else prim_path + self._trigger_path = f"{prim_path}/{trigger_subfix}" + self._knife_path = knife_path + self._base_pos = base_pos + self._base_quat = base_quat + + self._stage = None + self._collision_checker = None + self._dummy_db = None + self._last_collision = False + + def _setup_collision_checker(self): + self._collision_checker = Collision_Checker( + stage=self._stage, + prim_path0=self._trigger_path, + prim_path1=self._knife_path, + ) + self._last_collision = False + + def _setup_cut_mesh_node(self): + self._dummy_db.inputs.cut_mesh_path = self._mesh_path + self._dummy_db.inputs.knife_mesh_path = self._knife_path + self._dummy_db.internal_state = cutMeshNode.internal_state() + self._dummy_db.inputs.cutEventIn = False + + def initialize(self): + self._stage = omni.usd.get_context().get_stage() + self._dummy_db = SimpleNamespace() + self._dummy_db.inputs = SimpleNamespace() + self._setup_cut_mesh_node() + self._setup_collision_checker() + + def step(self): + if_collision, _, _ = self._collision_checker.meshes_aabb_collide() + if if_collision == self._last_collision and if_collision: + if_collision = False + else: + self._last_collision = if_collision + self._dummy_db.inputs.cutEventIn = if_collision + cutMeshNode.compute(self._dummy_db) + + def _respawn(self, position: tuple[float, float, float], orientation: tuple[float, float, float, float]): + """Delete and respawn the object at the given pose, then reinitialize cutting system.""" + delete_prim(self._prim_path) + cfg = sim_utils.UsdFileCfg(usd_path=self._usd_path) + cfg.func(self._prim_path, cfg, translation=position, orientation=orientation) + self._setup_cut_mesh_node() + self._setup_collision_checker() + + def reset(self): + """Reset to base position.""" + self._respawn(self._base_pos, self._base_quat) + + def reset_to(self, position: tuple[float, float, float], orientation: tuple[float, float, float, float]): + """Reset to specified position.""" + self._respawn(position, orientation) + + @property + def piece_count(self) -> int: + mesh_prim = self._stage.GetPrimAtPath(self._mesh_path) + return len([p for p in mesh_prim.GetChildren() if p.IsActive()]) + + @property + def root_pose_w(self) -> torch.Tensor: + prim = self._stage.GetPrimAtPath(self._prim_path) + if not prim.IsValid(): + return torch.tensor([*self._base_pos, *self._base_quat], dtype=torch.float32) + + xformable = UsdGeom.Xformable(prim) + transform = xformable.ComputeLocalToWorldTransform(0) + pos = transform.ExtractTranslation() + rot = transform.ExtractRotationQuat() + return torch.tensor([pos[0], pos[1], pos[2], rot.GetReal(), *rot.GetImaginary()], dtype=torch.float32) + + +class CuttableObject: + """Manages cuttable object instances across environments.""" + + cfg: CuttableObjectCfg + + def __init__(self, cfg: CuttableObjectCfg): + self.cfg = cfg + self.cuttable_objects: list[SingleCuttableObject] = [] + + matching_prims = sim_utils.find_matching_prim_paths(self.cfg.prim_path) + for prim_path in matching_prims: + self.cuttable_objects.append( + SingleCuttableObject( + prim_path=prim_path, + usd_path=self.cfg.usd_path, + mesh_subfix=self.cfg.mesh_subfix, + trigger_subfix=self.cfg.trigger_subfix, + knife_path=self.cfg.knife_prim_path, + base_pos=self.cfg.base_pos, + base_quat=self.cfg.base_quat, + ) + ) + + def initialize(self): + for obj in self.cuttable_objects: + obj.initialize() + + def step(self): + for obj in self.cuttable_objects: + obj.step() + + def reset(self, env_ids: list[int] | None = None): + if env_ids is None: + env_ids = range(len(self.cuttable_objects)) + for env_id in env_ids: + self.cuttable_objects[env_id].reset() + + def reset_to(self, env_ids: list[int], positions: torch.Tensor, orientations: torch.Tensor): + for i, env_id in enumerate(env_ids): + pos = tuple(positions[i].cpu().numpy()) + ori = tuple(orientations[i].cpu().numpy()) + self.cuttable_objects[env_id].reset_to(pos, ori) + + @property + def piece_counts(self) -> torch.Tensor: + return torch.tensor([obj.piece_count for obj in self.cuttable_objects], dtype=torch.int32) + + def check_success(self, min_count: int = 2) -> torch.Tensor: + return self.piece_counts >= min_count + + @property + def root_pose_w(self) -> torch.Tensor: + poses = [obj.root_pose_w for obj in self.cuttable_objects] + return torch.stack(poses, dim=0) diff --git a/source/leisaac/leisaac/enhance/assets/cuttable_object/cuttable_object_cfg.py b/source/leisaac/leisaac/enhance/assets/cuttable_object/cuttable_object_cfg.py new file mode 100644 index 00000000..2e6a57e1 --- /dev/null +++ b/source/leisaac/leisaac/enhance/assets/cuttable_object/cuttable_object_cfg.py @@ -0,0 +1,26 @@ +"""Configuration for cuttable objects.""" + +from isaaclab.assets import AssetBaseCfg +from isaaclab.utils import configclass + +from .cuttable_object import SingleCuttableObject + + +@configclass +class CuttableObjectCfg(AssetBaseCfg): + """Configuration for the cuttable object.""" + + usd_path: str = "" + """Path to the USD file for the cuttable object.""" + mesh_subfix: str = "" + """Subfix path to the mesh prim.""" + trigger_subfix: str = "Trigger/Cube" + """Subfix path to the trigger collision mesh.""" + knife_prim_path: str = "" + """Full prim path to the knife mesh.""" + base_pos: tuple[float, float, float] = (0.0, 0.0, 0.0) + """Base position of the object.""" + base_quat: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0) + """Base quaternion (wxyz) of the object.""" + class_type: type = SingleCuttableObject + """Class type of the cuttable object.""" diff --git a/source/leisaac/leisaac/enhance/envs/mdp/events.py b/source/leisaac/leisaac/enhance/envs/mdp/events.py index 2853083c..261c3fe8 100644 --- a/source/leisaac/leisaac/enhance/envs/mdp/events.py +++ b/source/leisaac/leisaac/enhance/envs/mdp/events.py @@ -122,3 +122,129 @@ def disable_rigid_body_gravity( prim_path, sim_utils.RigidBodyPropertiesCfg(disable_gravity=True), ) + + +def reset_mixed_objects_uniform( + env: ManagerBasedRLEnv, + env_ids: torch.Tensor, + pose_range: dict[str, tuple[float, float]], + velocity_range: dict[str, tuple[float, float]], + rigid_asset_cfg: list[SceneEntityCfg], + deformable_asset_cfg: list[SceneEntityCfg], +): + """Reset mixed rigid and deformable objects together with the same randomization offset. + + This allows rigid and deformable objects to be randomized as a group, + sharing the same random offset across all objects. + """ + # Sample common noise (pose) + range_list = [pose_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] + ranges = torch.tensor(range_list, device=env.device) + rand_samples_pose = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (len(env_ids), 6), device=env.device) + + # Sample common noise (velocity) + range_list_vel = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] + ranges_vel = torch.tensor(range_list_vel, device=env.device) + rand_samples_vel = math_utils.sample_uniform( + ranges_vel[:, 0], ranges_vel[:, 1], (len(env_ids), 6), device=env.device + ) + + # 1. Handle Rigid Objects + for cfg in rigid_asset_cfg: + asset = env.scene[cfg.name] + root_states = asset.data.default_root_state[env_ids].clone() + positions = root_states[:, 0:3] + env.scene.env_origins[env_ids] + rand_samples_pose[:, 0:3] + orientations_delta = math_utils.quat_from_euler_xyz( + rand_samples_pose[:, 3], rand_samples_pose[:, 4], rand_samples_pose[:, 5] + ) + orientations = math_utils.quat_mul(root_states[:, 3:7], orientations_delta) + velocities = root_states[:, 7:13] + rand_samples_vel + asset.write_root_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids) + asset.write_root_velocity_to_sim(velocities, env_ids=env_ids) + + # 2. Handle Deformable Objects + for cfg in deformable_asset_cfg: + asset = env.scene[cfg.name] + nodal_state = asset.data.default_nodal_state_w[env_ids].clone() + + nodal_state[..., :3] += rand_samples_pose[:, 0:3].unsqueeze(1) + nodal_state[..., 3:] += rand_samples_vel[:, 0:3].unsqueeze(1) + asset.write_nodal_state_to_sim(nodal_state, env_ids=env_ids) + + +def randomize_cuttable_object_uniform( + env: ManagerBasedRLEnv, + env_ids: torch.Tensor, + attr_name: str, + pose_range: dict[str, tuple[float, float]], +): + """Randomize a cuttable object's pose by moving it via USD transform. + + This works for cuttable objects that are stored as env attributes (e.g., env.cuttable_sausage). + The object is moved using USD Xform operations after being reset/respawned. + + Args: + env: The environment instance. + env_ids: Environment IDs to apply randomization. + attr_name: The attribute name of the CuttableObject on the env. + pose_range: Dict with keys x, y, z, roll, pitch, yaw and (min, max) tuple values. + """ + import omni + from pxr import Gf, UsdGeom + + cuttable_object = getattr(env, attr_name, None) + if cuttable_object is None: + return + + # Sample random offset + range_list = [pose_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] + ranges = torch.tensor(range_list, device=env.device) + rand_samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (len(env_ids), 6), device=env.device) + + stage = omni.usd.get_context().get_stage() + + for i, env_id in enumerate(env_ids): + if env_id >= len(cuttable_object.cuttable_objects): + continue + + single_obj = cuttable_object.cuttable_objects[env_id] + prim = stage.GetPrimAtPath(single_obj._prim_path) + if not prim.IsValid(): + continue + + xformable = UsdGeom.Xformable(prim) + + # Get current transform + current_transform = xformable.ComputeLocalToWorldTransform(0) + current_pos = current_transform.ExtractTranslation() + current_rot = current_transform.ExtractRotationQuat() + + # Apply position offset + offset = rand_samples[i].cpu().numpy() + new_pos = Gf.Vec3d( + current_pos[0] + offset[0], + current_pos[1] + offset[1], + current_pos[2] + offset[2], + ) + + # Apply rotation offset + from scipy.spatial.transform import Rotation as R + + current_quat_wxyz = [current_rot.GetReal(), *current_rot.GetImaginary()] + r_current = R.from_quat( + [current_quat_wxyz[1], current_quat_wxyz[2], current_quat_wxyz[3], current_quat_wxyz[0]] + ) + r_offset = R.from_euler("xyz", [offset[3], offset[4], offset[5]]) + r_new = r_current * r_offset + new_quat_xyzw = r_new.as_quat() + new_quat = Gf.Quatd(new_quat_xyzw[3], Gf.Vec3d(new_quat_xyzw[0], new_quat_xyzw[1], new_quat_xyzw[2])) + + # Build transform matrix and apply + transform_matrix = Gf.Matrix4d() + transform_matrix.SetTranslate(new_pos) + transform_matrix.SetRotateOnly(new_quat) + + # Clear and set new transform using a single transform op + xformable.ClearXformOpOrder() + xform_op = xformable.AddTransformOp() + xform_op.Set(transform_matrix) diff --git a/source/leisaac/leisaac/enhance/envs/mdp/recorders/recorders.py b/source/leisaac/leisaac/enhance/envs/mdp/recorders/recorders.py index 98336f28..10050760 100644 --- a/source/leisaac/leisaac/enhance/envs/mdp/recorders/recorders.py +++ b/source/leisaac/leisaac/enhance/envs/mdp/recorders/recorders.py @@ -22,7 +22,6 @@ class InitialStateWithParticleObjectsRecorder(RecorderTerm): def record_post_reset(self, env_ids: Sequence[int] | None): def extract_env_ids_values(value): - nonlocal env_ids if isinstance(value, dict): return {k: extract_env_ids_values(v) for k, v in value.items()} return value[env_ids] @@ -36,3 +35,25 @@ def extract_env_ids_values(value): state["particle_object"][asset_name] = asset_state return "initial_state", extract_env_ids_values(state) + + +class InitialStateWithCuttableObjectsRecorder(RecorderTerm): + """Records initial state including cuttable objects (sausage, etc.).""" + + def record_post_reset(self, env_ids: Sequence[int] | None): + def extract_env_ids_values(value): + if isinstance(value, dict): + return {k: extract_env_ids_values(v) for k, v in value.items()} + return value[env_ids] + + state = self._env.scene.get_state(is_relative=True) + state["cuttable_object"] = {} + for attr_name in dir(self._env): + if attr_name.startswith("cuttable_"): + cuttable_object = getattr(self._env, attr_name, None) + if cuttable_object is not None and hasattr(cuttable_object, "root_pose_w"): + asset_state = {"root_pose": cuttable_object.root_pose_w.clone()} + asset_state["root_pose"][:, :3] -= self._env.scene.env_origins + state["cuttable_object"][attr_name] = asset_state + + return "initial_state", extract_env_ids_values(state) diff --git a/source/leisaac/leisaac/enhance/envs/mdp/recorders/recorders_cfg.py b/source/leisaac/leisaac/enhance/envs/mdp/recorders/recorders_cfg.py index 1c02df39..81452829 100644 --- a/source/leisaac/leisaac/enhance/envs/mdp/recorders/recorders_cfg.py +++ b/source/leisaac/leisaac/enhance/envs/mdp/recorders/recorders_cfg.py @@ -26,6 +26,13 @@ class InitialStateWithParticleObjectsRecorderCfg(RecorderTermCfg): class_type: type[RecorderTerm] = recorders.InitialStateWithParticleObjectsRecorder +@configclass +class InitialStateWithCuttableObjectsRecorderCfg(RecorderTermCfg): + """Configuration for the initial state with cuttable objects recorder term.""" + + class_type: type[RecorderTerm] = recorders.InitialStateWithCuttableObjectsRecorder + + @configclass class DirectEnvActionStateRecorderManagerCfg(ActionStateRecorderManagerCfg): """Recorder configuration for recording actions and states in direct environment.""" @@ -39,3 +46,10 @@ class DirectEnvActionStateWithParticleObjectsRecorderManagerCfg(DirectEnvActionS """Recorder configuration for recording actions and states with particle objects in direct environment.""" record_initial_state = InitialStateWithParticleObjectsRecorderCfg() + + +@configclass +class DirectEnvActionStateWithCuttableObjectsRecorderManagerCfg(DirectEnvActionStateRecorderManagerCfg): + """Recorder configuration for recording actions and states with cuttable objects in direct environment.""" + + record_initial_state = InitialStateWithCuttableObjectsRecorderCfg() diff --git a/source/leisaac/leisaac/tasks/assemble_hamburger/__init__.py b/source/leisaac/leisaac/tasks/assemble_hamburger/__init__.py index 12684aaa..030589de 100644 --- a/source/leisaac/leisaac/tasks/assemble_hamburger/__init__.py +++ b/source/leisaac/leisaac/tasks/assemble_hamburger/__init__.py @@ -1,19 +1,19 @@ import gymnasium as gym gym.register( - id="LeIsaac-SO101-AssembleHamburger-v0", + id="LeIsaac-SO101-AssembleHamburger-BiArm-v0", entry_point="isaaclab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.assemble_hamhurger_env_cfg:AssembleHamburgerEnvCfg", + "env_cfg_entry_point": f"{__name__}.assemble_hamburger_bi_arm_env_cfg:AssembleHamburgerBiArmEnvCfg", }, ) gym.register( - id="LeIsaac-SO101-AssembleHamburger-BiArm-v0", - entry_point="isaaclab.envs:ManagerBasedRLEnv", + id="LeIsaac-SO101-AssembleHamburger-BiArm-Direct-v0", + entry_point=f"{__name__}.direct.assemble_hamburger_bi_arm_env:AssembleHamburgerBiArmEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.assemble_hamhurger_bi_arm_env_cfg:AssembleHamburgerBiArmEnvCfg", + "env_cfg_entry_point": f"{__name__}.direct.assemble_hamburger_bi_arm_env:AssembleHamburgerBiArmEnvCfg", }, ) diff --git a/source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamburger_bi_arm_env_cfg.py b/source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamburger_bi_arm_env_cfg.py new file mode 100644 index 00000000..192f5311 --- /dev/null +++ b/source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamburger_bi_arm_env_cfg.py @@ -0,0 +1,107 @@ +""" +Assemble Hamburger Bi-Arm Scene Configuration (LeIsaac Pattern) +""" + +import isaaclab.sim as sim_utils +from isaaclab.assets import AssetBaseCfg +from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import TerminationTermCfg as DoneTerm +from isaaclab.sim import SimulationCfg +from isaaclab.utils import configclass +from leisaac.assets.scenes.kitchen import ( + KITCHEN_WITH_HAMBURGER_CFG, + KITCHEN_WITH_HAMBURGER_USD_PATH, +) +from leisaac.utils.domain_randomization import ( + domain_randomization, + randomize_deformable_object_uniform, + randomize_mixed_objects_uniform, +) +from leisaac.utils.general_assets import parse_usd_and_create_subassets + +from ..template import ( + BiArmObservationsCfg, + BiArmTaskEnvCfg, + BiArmTaskSceneCfg, + BiArmTerminationsCfg, +) +from . import mdp + + +@configclass +class AssembleHamburgerBiArmSceneCfg(BiArmTaskSceneCfg): + """Scene configuration - using leisaac automated loading pattern.""" + + scene: AssetBaseCfg = KITCHEN_WITH_HAMBURGER_CFG.replace(prim_path="{ENV_REGEX_NS}/Scene") + + +@configclass +class TerminationsCfg(BiArmTerminationsCfg): + """Termination configuration for the assemble hamburger task.""" + + success = DoneTerm( + func=mdp.burger_assembly_success, + params={ + "beef_cfg": SceneEntityCfg("Burger_Beef_Patties001"), + "plate_cfg": SceneEntityCfg("Burger_Plate"), + }, + ) + + +@configclass +class AssembleHamburgerBiArmEnvCfg(BiArmTaskEnvCfg): + """Configuration for the assemble hamburger environment (Manager-Based).""" + + scene: AssembleHamburgerBiArmSceneCfg = AssembleHamburgerBiArmSceneCfg(env_spacing=4.0) + + observations: BiArmObservationsCfg = BiArmObservationsCfg() + + terminations: TerminationsCfg = TerminationsCfg() + + task_description: str = "Pick the beef patties and place it on the plate" + + # Simulation configuration + render_cfg: sim_utils.RenderCfg = sim_utils.RenderCfg(rendering_mode="quality", antialiasing_mode="FXAA") + sim: SimulationCfg = SimulationCfg(dt=1 / 60, render_interval=1, render=render_cfg, use_fabric=True) + + def __post_init__(self) -> None: + super().__post_init__() + + # Aligned with leisaac kitchen reference system + self.viewer.eye = (2.4, -5.577, 1.52) + self.viewer.lookat = (4.03, -6.41, 0.72) + + # Robot positions - Aligned for burger kitchen + self.scene.left_arm.init_state.pos = (3.4, -5.8, 0.78) + self.scene.left_arm.init_state.rot = (0.707, 0.0, 0.0, 0.707) + + self.scene.right_arm.init_state.pos = (3.4, -6.4, 0.78) + self.scene.right_arm.init_state.rot = (0.707, 0.0, 0.0, 0.707) + + self.decimation = 1 + self.dynamic_reset_gripper_effort_limit = False + + # Add lighting + self.scene.light = AssetBaseCfg( + prim_path="{ENV_REGEX_NS}/Light", + spawn=sim_utils.DomeLightCfg(intensity=1000.0, color=(0.75, 0.75, 0.75)), + ) + + # Automatically parse and add burger components to the scene + parse_usd_and_create_subassets(KITCHEN_WITH_HAMBURGER_USD_PATH, self) + + # Domain randomization - plate, bread, cheese bound together (lehome style) + domain_randomization( + self, + random_options=[ + randomize_mixed_objects_uniform( + rigid_names=["Burger_Plate", "Burger_Bread002"], + deformable_names=["Burger_Cheese001"], + pose_range={"x": (-0.05, 0.05), "y": (-0.05, 0.05), "z": (0.0, 0.0)}, + ), + randomize_deformable_object_uniform( + "Burger_Beef_Patties001", + pose_range={"x": (-0.1, 0.1), "y": (-0.1, 0.1), "z": (0.0, 0.0)}, + ), + ], + ) diff --git a/source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamhurger_bi_arm_env_cfg.py b/source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamhurger_bi_arm_env_cfg.py deleted file mode 100644 index fb31263f..00000000 --- a/source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamhurger_bi_arm_env_cfg.py +++ /dev/null @@ -1,45 +0,0 @@ -from isaaclab.assets import AssetBaseCfg -from isaaclab.utils import configclass -from leisaac.assets.scenes.kitchen import ( - KITCHEN_WITH_HAMBURGER_CFG, - KITCHEN_WITH_HAMBURGER_USD_PATH, -) -from leisaac.utils.general_assets import parse_usd_and_create_subassets - -from ..template import ( - BiArmObservationsCfg, - BiArmTaskEnvCfg, - BiArmTaskSceneCfg, - BiArmTerminationsCfg, -) - - -@configclass -class AssembleHamburgerBiArmSceneCfg(BiArmTaskSceneCfg): - """Scene configuration for the assemble hamburger task using two arms.""" - - scene: AssetBaseCfg = KITCHEN_WITH_HAMBURGER_CFG.replace(prim_path="{ENV_REGEX_NS}/Scene") - - -@configclass -class AssembleHamburgerBiArmEnvCfg(BiArmTaskEnvCfg): - """Configuration for the assemble hamburger environment.""" - - scene: AssembleHamburgerBiArmSceneCfg = AssembleHamburgerBiArmSceneCfg(env_spacing=8.0) - - observations: BiArmObservationsCfg = BiArmObservationsCfg() - - terminations: BiArmTerminationsCfg = BiArmTerminationsCfg() - - task_description: str = "Assemble the hamburger in order." - - def __post_init__(self) -> None: - super().__post_init__() - - self.viewer.eye = (2.5, -1.0, 1.3) - self.viewer.lookat = (3.6, -0.4, 1.0) - - self.scene.left_arm.init_state.pos = (3.4, -0.65, 0.89) - self.scene.right_arm.init_state.pos = (3.8, -0.65, 0.89) - - parse_usd_and_create_subassets(KITCHEN_WITH_HAMBURGER_USD_PATH, self) diff --git a/source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamhurger_env_cfg.py b/source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamhurger_env_cfg.py deleted file mode 100644 index 98ab50a4..00000000 --- a/source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamhurger_env_cfg.py +++ /dev/null @@ -1,44 +0,0 @@ -from isaaclab.assets import AssetBaseCfg -from isaaclab.utils import configclass -from leisaac.assets.scenes.kitchen import ( - KITCHEN_WITH_HAMBURGER_CFG, - KITCHEN_WITH_HAMBURGER_USD_PATH, -) -from leisaac.utils.general_assets import parse_usd_and_create_subassets - -from ..template import ( - SingleArmObservationsCfg, - SingleArmTaskEnvCfg, - SingleArmTaskSceneCfg, - SingleArmTerminationsCfg, -) - - -@configclass -class AssembleHamburgerSceneCfg(SingleArmTaskSceneCfg): - """Scene configuration for the assemble hamburger task.""" - - scene: AssetBaseCfg = KITCHEN_WITH_HAMBURGER_CFG.replace(prim_path="{ENV_REGEX_NS}/Scene") - - -@configclass -class AssembleHamburgerEnvCfg(SingleArmTaskEnvCfg): - """Configuration for the assemble hamburger environment.""" - - scene: AssembleHamburgerSceneCfg = AssembleHamburgerSceneCfg(env_spacing=8.0) - - observations: SingleArmObservationsCfg = SingleArmObservationsCfg() - - terminations: SingleArmTerminationsCfg = SingleArmTerminationsCfg() - - task_description: str = "Assemble the hamburger in order." - - def __post_init__(self) -> None: - super().__post_init__() - - self.viewer.eye = (2.5, -1.0, 1.3) - self.viewer.lookat = (3.6, -0.4, 1.0) - - self.scene.robot.init_state.pos = (3.6, -0.65, 0.89) - - parse_usd_and_create_subassets(KITCHEN_WITH_HAMBURGER_USD_PATH, self) diff --git a/source/leisaac/leisaac/tasks/assemble_hamburger/direct/__init__.py b/source/leisaac/leisaac/tasks/assemble_hamburger/direct/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/source/leisaac/leisaac/tasks/assemble_hamburger/direct/assemble_hamburger_bi_arm_env.py b/source/leisaac/leisaac/tasks/assemble_hamburger/direct/assemble_hamburger_bi_arm_env.py new file mode 100644 index 00000000..b68c35f7 --- /dev/null +++ b/source/leisaac/leisaac/tasks/assemble_hamburger/direct/assemble_hamburger_bi_arm_env.py @@ -0,0 +1,92 @@ +""" +Assemble Hamburger Bi-Arm Direct Environment (LeIsaac Pattern) +""" + +import isaaclab.sim as sim_utils +import torch +from isaaclab.assets import AssetBaseCfg +from isaaclab.managers import SceneEntityCfg +from isaaclab.sim import SimulationCfg +from isaaclab.utils import configclass +from leisaac.utils.domain_randomization import ( + domain_randomization, + randomize_deformable_object_uniform, + randomize_mixed_objects_uniform, +) +from leisaac.utils.general_assets import parse_usd_and_create_subassets + +from ...template import BiArmTaskDirectEnv, BiArmTaskDirectEnvCfg +from .. import mdp +from ..assemble_hamburger_bi_arm_env_cfg import ( + KITCHEN_WITH_HAMBURGER_USD_PATH, + AssembleHamburgerBiArmSceneCfg, +) + + +@configclass +class AssembleHamburgerBiArmEnvCfg(BiArmTaskDirectEnvCfg): + """Direct env configuration.""" + + scene: AssembleHamburgerBiArmSceneCfg = AssembleHamburgerBiArmSceneCfg(env_spacing=4.0) + task_description: str = "Pick the beef patties and place it on the plate" + + # Render configuration - match manager-based env for proper material colors + render_cfg: sim_utils.RenderCfg = sim_utils.RenderCfg(rendering_mode="quality", antialiasing_mode="FXAA") + sim: SimulationCfg = SimulationCfg(dt=1 / 60, render_interval=1, render=render_cfg, use_fabric=True) + + def __post_init__(self) -> None: + super().__post_init__() + + # Aligned with leisaac kitchen + self.viewer.eye = (2.4, -5.577, 1.52) + self.viewer.lookat = (4.03, -6.41, 0.72) + + self.scene.left_arm.init_state.pos = (3.4, -5.8, 0.78) + self.scene.left_arm.init_state.rot = (0.707, 0.0, 0.0, 0.707) + + self.scene.right_arm.init_state.pos = (3.4, -6.4, 0.78) + self.scene.right_arm.init_state.rot = (0.707, 0.0, 0.0, 0.707) + + self.decimation = 1 + + # Add lighting + self.scene.light = AssetBaseCfg( + prim_path="{ENV_REGEX_NS}/Light", + spawn=sim_utils.DomeLightCfg(intensity=1000.0, color=(0.75, 0.75, 0.75)), + ) + + # Parse USD for auto-loading burger components + parse_usd_and_create_subassets(KITCHEN_WITH_HAMBURGER_USD_PATH, self) + + # Domain randomization - plate, bread, cheese bound together (lehome style) + domain_randomization( + self, + random_options=[ + randomize_mixed_objects_uniform( + rigid_names=["Burger_Plate", "Burger_Bread002"], + deformable_names=["Burger_Cheese001"], + pose_range={"x": (-0.05, 0.05), "y": (-0.05, 0.05), "z": (0.0, 0.0)}, + ), + randomize_deformable_object_uniform( + "Burger_Beef_Patties001", + pose_range={"x": (-0.1, 0.1), "y": (-0.1, 0.1), "z": (0.0, 0.0)}, + ), + ], + ) + + +class AssembleHamburgerBiArmEnv(BiArmTaskDirectEnv): + """Direct env for the assemble hamburger task.""" + + cfg: AssembleHamburgerBiArmEnvCfg + + def _check_success(self) -> torch.Tensor: + """Check if the burger assembly task is complete. + + Success: beef is placed on the plate (XY distance < 0.045, Z distance < 0.03). + """ + return mdp.burger_assembly_success( + env=self, + beef_cfg=SceneEntityCfg("Burger_Beef_Patties001"), + plate_cfg=SceneEntityCfg("Burger_Plate"), + ) diff --git a/source/leisaac/leisaac/tasks/assemble_hamburger/mdp/__init__.py b/source/leisaac/leisaac/tasks/assemble_hamburger/mdp/__init__.py new file mode 100644 index 00000000..4149e60b --- /dev/null +++ b/source/leisaac/leisaac/tasks/assemble_hamburger/mdp/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""MDP module for the assemble hamburger task.""" + +from .terminations import burger_assembly_success + +__all__ = ["burger_assembly_success"] diff --git a/source/leisaac/leisaac/tasks/assemble_hamburger/mdp/terminations.py b/source/leisaac/leisaac/tasks/assemble_hamburger/mdp/terminations.py new file mode 100644 index 00000000..056f9be5 --- /dev/null +++ b/source/leisaac/leisaac/tasks/assemble_hamburger/mdp/terminations.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +import torch +from isaaclab.assets import DeformableObject, RigidObject +from isaaclab.envs import DirectRLEnv, ManagerBasedRLEnv +from isaaclab.managers import SceneEntityCfg + + +def burger_assembly_success( + env: ManagerBasedRLEnv | DirectRLEnv, + beef_cfg: SceneEntityCfg, + plate_cfg: SceneEntityCfg, + xy_threshold: float = 0.045, + z_threshold: float = 0.03, +) -> torch.Tensor: + """Check if the burger beef is placed on the plate. + + Success condition: + - XY distance between beef center and plate center < xy_threshold + - Z distance < z_threshold + + Args: + env: The RL environment instance. + beef_cfg: Configuration for the beef deformable object. + plate_cfg: Configuration for the plate rigid object. + xy_threshold: Maximum XY distance for success (default 0.045m). + z_threshold: Maximum Z distance for success (default 0.03m). + + Returns: + A boolean tensor indicating success for each environment. + """ + # Access deformable object from deformable_objects dict + beef: DeformableObject = env.scene.deformable_objects[beef_cfg.name] + # Access rigid object from rigid_objects dict + plate: RigidObject = env.scene.rigid_objects[plate_cfg.name] + + beef_pos = beef.data.root_pos_w # (num_envs, 3) + plate_pos = plate.data.root_pos_w # (num_envs, 3) + + # XY distance + diff_xy = beef_pos[:, :2] - plate_pos[:, :2] + dist_xy = torch.linalg.norm(diff_xy, dim=-1) + + # Z distance + diff_z = torch.abs(beef_pos[:, 2] - plate_pos[:, 2]) + + # Success condition + success_mask = (dist_xy < xy_threshold) & (diff_z < z_threshold) + + # Debug print + if success_mask.any(): + print(f"[BURGER SUCCESS] xy_dist={dist_xy[0]:.4f}, z_dist={diff_z[0]:.4f}") + + return success_mask diff --git a/source/leisaac/leisaac/tasks/sausage_cut/__init__.py b/source/leisaac/leisaac/tasks/sausage_cut/__init__.py new file mode 100644 index 00000000..3af0cfb3 --- /dev/null +++ b/source/leisaac/leisaac/tasks/sausage_cut/__init__.py @@ -0,0 +1,10 @@ +import gymnasium as gym + +gym.register( + id="LeIsaac-SO101-SausageCut-BiArm-Direct-v0", + entry_point=f"{__name__}.direct.sausage_cut_bi_arm_env:SausageCutBiArmEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.direct.sausage_cut_bi_arm_env:SausageCutBiArmEnvCfg", + }, +) diff --git a/source/leisaac/leisaac/tasks/sausage_cut/direct/__init__.py b/source/leisaac/leisaac/tasks/sausage_cut/direct/__init__.py new file mode 100644 index 00000000..ffb0e00c --- /dev/null +++ b/source/leisaac/leisaac/tasks/sausage_cut/direct/__init__.py @@ -0,0 +1 @@ +from .sausage_cut_bi_arm_env import SausageCutBiArmEnv, SausageCutBiArmEnvCfg diff --git a/source/leisaac/leisaac/tasks/sausage_cut/direct/sausage_cut_bi_arm_env.py b/source/leisaac/leisaac/tasks/sausage_cut/direct/sausage_cut_bi_arm_env.py new file mode 100644 index 00000000..2a95f682 --- /dev/null +++ b/source/leisaac/leisaac/tasks/sausage_cut/direct/sausage_cut_bi_arm_env.py @@ -0,0 +1,154 @@ +"""Sausage cutting task with bi-arm robot (Direct environment).""" + +from collections.abc import Sequence + +import isaaclab.sim as sim_utils +import torch +from isaaclab.assets import ArticulationCfg, AssetBaseCfg +from isaaclab.sim import SimulationCfg +from isaaclab.utils import configclass +from leisaac.assets.robots.lerobot import SO101_KINFE_CFG +from leisaac.assets.scenes.kitchen import ( + KITCHEN_WITH_SAUSAGE_CFG, + KITCHEN_WITH_SAUSAGE_USD_PATH, + SAUSAGE_USD_PATH, +) +from leisaac.enhance.assets import CuttableObject, CuttableObjectCfg +from leisaac.enhance.envs.mdp.recorders.recorders_cfg import ( + DirectEnvActionStateWithCuttableObjectsRecorderManagerCfg as CuttableObjectRecordTerm, +) +from leisaac.utils.domain_randomization import ( + domain_randomization, + randomize_cuttable_object_uniform, +) +from leisaac.utils.general_assets import parse_usd_and_create_subassets + +from ...template import BiArmTaskDirectEnv, BiArmTaskDirectEnvCfg, BiArmTaskSceneCfg + +SAUSAGE_BASE_POS = (3.6612, -6.236, 0.84059) +SAUSAGE_BASE_QUAT = (-0.23287, -0.02628, 0.02471, 0.97184) + + +@configclass +class SausageCutBiArmSceneCfg(BiArmTaskSceneCfg): + """Scene configuration for the sausage cut task using two arms.""" + + scene: AssetBaseCfg = KITCHEN_WITH_SAUSAGE_CFG.replace(prim_path="{ENV_REGEX_NS}/Scene") + right_arm: ArticulationCfg = SO101_KINFE_CFG.replace(prim_path="{ENV_REGEX_NS}/Right_Robot") + + +@configclass +class SausageCutBiArmEnvCfg(BiArmTaskDirectEnvCfg): + """Direct env configuration for the sausage cut task.""" + + scene: SausageCutBiArmSceneCfg = SausageCutBiArmSceneCfg(env_spacing=4.0) + task_description: str = "Cut the sausage into pieces." + + # Render configuration with antialiasing enabled + render_cfg: sim_utils.RenderCfg = sim_utils.RenderCfg(rendering_mode="quality", antialiasing_mode="FXAA") + sim: SimulationCfg = SimulationCfg(dt=1 / 60, render_interval=1, render=render_cfg, use_fabric=False) + + # Disable IsaacLab UI window + ui_window_class_type: type | None = None + + def __post_init__(self) -> None: + super().__post_init__() + + self.viewer.eye = (2.15, -6.5, 1.9) + self.viewer.lookat = (3.90, -6.43, 0.92) + + self.scene.left_arm.init_state.pos = (3.4, -5.85, 0.768) + self.scene.left_arm.init_state.rot = (0.707, 0.0, 0.0, 0.707) + self.scene.right_arm.init_state.pos = (3.4, -6.45, 0.768) + self.scene.right_arm.init_state.rot = (0.707, 0.0, 0.0, 0.707) + + # Simulation settings for deformable body cutting (already set in sim config above) + self.decimation = 1 + self.dynamic_reset_gripper_effort_limit = False + + # Recorder for cuttable objects + self.recorders = CuttableObjectRecordTerm() + + self.scene.light = AssetBaseCfg( + prim_path="{ENV_REGEX_NS}/Light", + spawn=sim_utils.DomeLightCfg(intensity=1000.0, color=(0.75, 0.75, 0.75)), + ) + parse_usd_and_create_subassets(KITCHEN_WITH_SAUSAGE_USD_PATH, self, exclude_name_list=["Sausage"]) + + # Domain randomization for the cuttable sausage + domain_randomization( + self, + random_options=[ + randomize_cuttable_object_uniform( + "cuttable_sausage", + pose_range={ + "x": (-0.04, 0.04), + "y": (-0.04, 0.04), + "yaw": (-20 * torch.pi / 180, 20 * torch.pi / 180), + }, + ), + ], + ) + + +class SausageCutBiArmEnv(BiArmTaskDirectEnv): + """Direct env for the sausage cut task.""" + + cfg: SausageCutBiArmEnvCfg + + def _setup_scene(self): + super()._setup_scene() + env_path = self.scene.env_prim_paths[0] + sausage_prim_path = f"{env_path}/Scene/Sausage001" + + # Spawn sausage + cfg = sim_utils.UsdFileCfg(usd_path=SAUSAGE_USD_PATH) + cfg.func(sausage_prim_path, cfg, translation=SAUSAGE_BASE_POS, orientation=SAUSAGE_BASE_QUAT) + + # Initialize cuttable object + self.cuttable_sausage = CuttableObject( + CuttableObjectCfg( + prim_path=sausage_prim_path, + usd_path=SAUSAGE_USD_PATH, + mesh_subfix="Sausage001", + trigger_subfix="Trigger/Cube", + knife_prim_path=f"{env_path}/Right_Robot/gripper/Knife/Knife/Knife002", + base_pos=SAUSAGE_BASE_POS, + base_quat=SAUSAGE_BASE_QUAT, + ) + ) + + def initialize(self): + self.cuttable_sausage.initialize() + + def _apply_action(self) -> None: + super()._apply_action() + self.cuttable_sausage.step() + + def _check_success(self) -> torch.Tensor: + return self.cuttable_sausage.check_success(min_count=2) + + def _reset_idx(self, env_ids: Sequence[int]): + self.cuttable_sausage.reset(list(env_ids)) + super()._reset_idx(env_ids) + + def reset_to( + self, + state: dict[str, dict[str, dict[str, torch.Tensor]]], + env_ids: Sequence[int] | None = None, + seed: int | None = None, + is_relative: bool = False, + ): + if env_ids is None: + env_ids = torch.arange(self.num_envs, dtype=torch.int64, device=self.device) + # Call super first (which calls _reset_idx and resets sausage to default position) + super().reset_to(state, env_ids, seed, is_relative) + # Then override with the position from HDF5 state if available + if "cuttable_object" in state: + for attr_name, asset_state in state["cuttable_object"].items(): + cuttable_object = getattr(self, attr_name, None) + if cuttable_object is not None: + root_pose = asset_state["root_pose"].clone() + if is_relative: + root_pose[:, :3] += self.scene.env_origins[env_ids] + cuttable_object.reset_to(list(env_ids), root_pose[:, :3], root_pose[:, 3:]) diff --git a/source/leisaac/leisaac/utils/collision_checker.py b/source/leisaac/leisaac/utils/collision_checker.py new file mode 100644 index 00000000..3c3f2ecb --- /dev/null +++ b/source/leisaac/leisaac/utils/collision_checker.py @@ -0,0 +1,147 @@ +import numpy as np +import omni +import trimesh +from pxr import Gf, Sdf, Usd, UsdGeom + + +class Collision_Checker: + """ + Real-time collision detection package + See usage example at the bottom + """ + + def __init__( + self, + stage=None, + prim_path0="/World/Scene/Sausage001/Trigger/Cube", + prim_path1="/World/Robot/Right_Robot/gripper/Knife/Knife/Knife002", + apply_world_transform=True, + debug=False, + ): + """ + Parameters: + - stage: pxr Usd stage (defaults to automatically getting the current stage) + - dynamic_vertices: If True, re-read the prim's vertices for each check (for meshes with vertex deformations) + - enable_proximity: If True, attempts to return the closest point and distance on collision (may be slow) + """ + self.stage = stage + self.apply_world_transform = apply_world_transform + self.prim_path0 = prim_path0 + self.prim_path1 = prim_path1 + self.debug = debug + + def get_current_timecode(self): + """ + Returns the Usd.TimeCode corresponding to the current time (usually in frames) + """ + timeline = omni.timeline.get_timeline_interface() + current_time = timeline.get_current_time() + return Usd.TimeCode(current_time) + + def from_str2Usd_Prim(self, prim_path): + path = Sdf.Path(prim_path) + prim = self.stage.GetPrimAtPath(path) + return prim + + def usd_mesh_to_trimesh(self, prim_path): + prim = self.from_str2Usd_Prim(prim_path) + mesh = UsdGeom.Mesh(prim) + points = np.array(mesh.GetPointsAttr().Get(), dtype=np.float64) # (N,3) + face_vertex_indices = np.array(mesh.GetFaceVertexIndicesAttr().Get(), dtype=np.int64) + face_vertex_counts = np.array(mesh.GetFaceVertexCountsAttr().Get(), dtype=np.int64) + faces = [] + index = 0 + for count in face_vertex_counts: + if count == 3: + faces.append(face_vertex_indices[index : index + 3]) + elif count > 3: + # fan triangulation: (v0, v1, v2), (v0, v2, v3), ... + base = face_vertex_indices[index] # first vertex index + for i in range(1, count - 1): + faces.append([ + base, + face_vertex_indices[index + i], + face_vertex_indices[index + i + 1], + ]) + else: + # count < 3: Ignore or report an error (usually should not occur) + # Choose to ignore: + pass + index += count + + trimesh_mesh = trimesh.Trimesh(points, np.array(faces)) + return trimesh_mesh + + def compute_A_position_in_B_space(self, position, primA, primB): + timeline = omni.timeline.get_timeline_interface() + current_time = timeline.get_current_time() + usd_current_time = Usd.TimeCode(current_time) + + body0_xform = UsdGeom.Xformable(primA) + body1_xform = UsdGeom.Xformable(primB) + + Trans0 = body0_xform.ComputeLocalToWorldTransform(usd_current_time) + p0_in_world = Trans0.Transform(position) + + Trans1 = body1_xform.ComputeLocalToWorldTransform(usd_current_time) + p0_in_p1 = Trans1.GetInverse().Transform(p0_in_world) + return p0_in_p1 + + def transfer_A_trimesh_in_B_space(self, mesh: trimesh.Trimesh, primA: str, primB: str): + primA = self.from_str2Usd_Prim(primA) + primB = self.from_str2Usd_Prim(primB) + new_vertices = [] + for vertex in mesh.vertices: + tmp_vertex = self.compute_A_position_in_B_space(Gf.Vec3d(vertex[0], vertex[1], vertex[2]), primA, primB) + vertex = [tmp_vertex[0], tmp_vertex[1], tmp_vertex[2]] + new_vertices.append(vertex) + mesh.vertices = np.array(new_vertices) + return mesh + + def world_aabb_from_world_trimesh(self, mesh: trimesh.Trimesh) -> tuple[np.ndarray, np.ndarray]: + """ + Given a local-space mesh (trimesh) and a Local->World matrix mat, compute the world-space AABB. + Only transform the matrix at the 8 corner points (efficient and accurate). + """ + local_min, local_max = mesh.bounds # (3,), (3,) + corners = np.array( + [ + [local_min[0], local_min[1], local_min[2]], + [local_min[0], local_min[1], local_max[2]], + [local_min[0], local_max[1], local_min[2]], + [local_min[0], local_max[1], local_max[2]], + [local_max[0], local_min[1], local_min[2]], + [local_max[0], local_min[1], local_max[2]], + [local_max[0], local_max[1], local_min[2]], + [local_max[0], local_max[1], local_max[2]], + ], + dtype=np.float64, + ) + h = np.ones((8, 4), dtype=np.float64) + h[:, :3] = corners + transformed = h + transformed_xyz = transformed[:, :3] / transformed[:, 3:4] + world_min = transformed_xyz.min(axis=0) + world_max = transformed_xyz.max(axis=0) + if self.debug: + print(" corners local:", corners) + print(" corners world:", transformed_xyz) + print(" world_min, world_max:", world_min, world_max) + return world_min, world_max + + def aabb_overlap(self, min_a: np.ndarray, max_a: np.ndarray, min_b: np.ndarray, max_b: np.ndarray) -> bool: + return np.all(max_a >= min_b) and np.all(max_b >= min_a) + + def meshes_aabb_collide(self): + mesh_a = self.usd_mesh_to_trimesh(self.prim_path0) + mesh_b = self.usd_mesh_to_trimesh(self.prim_path1) + mesh_a_world = self.transfer_A_trimesh_in_B_space(mesh_a, self.prim_path0, "/World") + mesh_b_world = self.transfer_A_trimesh_in_B_space(mesh_b, self.prim_path1, "/World") + min_a, max_a = self.world_aabb_from_world_trimesh(mesh_a_world) + min_b, max_b = self.world_aabb_from_world_trimesh(mesh_b_world) + collides = self.aabb_overlap(min_a, max_a, min_b, max_b) + if self.debug: + print("A AABB:", min_a, max_a) + print("B AABB:", min_b, max_b) + print("collides:", collides) + return collides, (min_a, max_a), (min_b, max_b) diff --git a/source/leisaac/leisaac/utils/cutMeshNode.py b/source/leisaac/leisaac/utils/cutMeshNode.py new file mode 100644 index 00000000..202644e4 --- /dev/null +++ b/source/leisaac/leisaac/utils/cutMeshNode.py @@ -0,0 +1,391 @@ +import os +import platform +from dataclasses import dataclass + +import numpy as np +import omni +import pip +from pxr import Gf, PhysxSchema, Usd, UsdGeom, UsdPhysics + + +class cutMeshNode: + + @dataclass + class InnerState: + _meshPath: str = "null path" + + @staticmethod + def internal_state(): + return cutMeshNode.InnerState() + + @staticmethod + def compute(db) -> bool: + """Compute the outputs from the current input""" + mesh_path = db.inputs.cut_mesh_path + # state = db.internal_state + stage = omni.usd.get_context().get_stage() + + if db.inputs.cutEventIn: + + knife_path = db.inputs.knife_mesh_path + knife_prim = stage.GetPrimAtPath(knife_path) + plane_center, plane_normal = get_knife_plane_in_world(knife_prim) + + mesh_prim = stage.GetPrimAtPath(mesh_path) + for mesh_child in mesh_prim.GetChildren(): + if mesh_child is not None: + plane_center_in_prim = compute_world_position_in_prim(plane_center, mesh_child) + plane_normal_in_prim = compute_world_dir_in_prim(plane_normal, mesh_child) + cut_prim(stage, mesh_child, plane_center_in_prim, plane_normal_in_prim) + + return True + + +def cut_prim(stage, prim, plane_origin, plane_normal): + try: + import trimesh + except Exception as error: + print("errrrrrrro") + package_list = ("trimesh", "shapely", "rtree", "triangle") # , + for package_name in package_list: + install_package(package_name) + raise RuntimeError("Find package failed") from error + + def get_Trimesh(prim: Usd.Prim) -> trimesh.Trimesh: + mesh = UsdGeom.Mesh(prim) + points = np.array(mesh.GetPointsAttr().Get()) + face_vertex_counts = np.array(mesh.GetFaceVertexCountsAttr().Get()) + face_vertex_indices = np.array(mesh.GetFaceVertexIndicesAttr().Get()) + faces = [] + index = 0 + for count in face_vertex_counts: + if count == 3: + faces.append(face_vertex_indices[index : index + 3]) + elif count > 3: + pass + index += count + trimesh_mesh = trimesh.Trimesh(points, np.array(faces)) + return trimesh_mesh + + def generate_mesh(meshPrim: UsdGeom.Mesh, trimesh: trimesh.Trimesh): + + new_points = [Gf.Vec3f(*v) for v in trimesh.vertices] + face_vertex_counts = [3] * len(trimesh.faces) + face_vertex_indices = trimesh.faces.flatten().tolist() + + meshPrim.CreatePointsAttr(new_points) + meshPrim.CreateFaceVertexCountsAttr(face_vertex_counts) + meshPrim.CreateFaceVertexIndicesAttr(face_vertex_indices) + + trimesh_mesh = get_Trimesh(prim) + + plane_normal_opposite = [-plane_normal[0], -plane_normal[1], -plane_normal[2]] + parent_prim = prim.GetParent() + cut_mesh_name = prim.GetName() + + # try: + sub_trimesh1 = trimesh_mesh.slice_plane(plane_origin, plane_normal, cap=True) + sub_trimesh2 = trimesh_mesh.slice_plane(plane_origin, plane_normal_opposite, cap=True) + + if len(sub_trimesh2.vertices) == 0 or len(sub_trimesh1.vertices) == 0: + return True + + new_prim_name1 = cut_mesh_name + "_sub1" + new_prim_path1 = parent_prim.GetPath().AppendChild(new_prim_name1) + new_mesh1 = UsdGeom.Mesh.Define(stage, new_prim_path1) + generate_mesh(new_mesh1, sub_trimesh1) + new_mesh_prim1 = stage.GetPrimAtPath(new_prim_path1) + if prim.HasAPI(PhysxSchema.PhysxDeformableAPI): + omni.kit.commands.execute( + "AddPhysicsComponent", + usd_prim=new_mesh_prim1, + component="PhysxDeformableBodyAPI", + ) + omni.kit.commands.execute("AddDeformableBodyComponent", skin_mesh_path=new_mesh_prim1.GetPath()) + copy_deformable_properties(prim, new_mesh_prim1) + copy_transform(prim, new_mesh_prim1) + copy_material(prim, new_mesh_prim1) + + elif prim.HasAPI(UsdPhysics.RigidBodyAPI): + omni.kit.commands.execute( + "AddPhysicsComponent", + usd_prim=new_mesh_prim1, + component="PhysicsRigidBodyAPI", + ) + + omni.kit.commands.execute( + "AddPhysicsComponent", + usd_prim=new_mesh_prim1, + component="PhysicsCollisionAPI", + ) + # To copy rigid body properties + copy_transform(prim, new_mesh_prim1) + + new_prim_name2 = cut_mesh_name + "_sub2" + new_prim_path2 = parent_prim.GetPath().AppendChild(new_prim_name2) + new_mesh2 = UsdGeom.Mesh.Define(stage, new_prim_path2) + generate_mesh(new_mesh2, sub_trimesh2) + new_mesh_prim2 = stage.GetPrimAtPath(new_prim_path2) + if prim.HasAPI(PhysxSchema.PhysxDeformableAPI): + omni.kit.commands.execute( + "AddPhysicsComponent", + usd_prim=new_mesh_prim2, + component="PhysxDeformableBodyAPI", + ) + omni.kit.commands.execute("AddDeformableBodyComponent", skin_mesh_path=new_mesh_prim2.GetPath()) + copy_deformable_properties(prim, new_mesh_prim2) + copy_transform(prim, new_mesh_prim2) + copy_material(prim, new_mesh_prim2) + + elif prim.HasAPI(UsdPhysics.RigidBodyAPI): + omni.kit.commands.execute( + "AddPhysicsComponent", + usd_prim=new_mesh_prim2, + component="PhysicsRigidBodyAPI", + ) + + omni.kit.commands.execute( + "AddPhysicsComponent", + usd_prim=new_mesh_prim2, + component="PhysicsCollisionAPI", + ) + # To copy rigid body properties + copy_transform(prim, new_mesh_prim2) + prim.SetActive(False) + # stage.RemovePrim(prim.GetPath()) + # except Exception as error: + # raise RuntimeError(f"Cut Mesh failed{prim.GetName()}") from error + return None + + +def copy_deformable_properties(primA, primB): + deformableStateA = PhysxSchema.PhysxDeformableAPI(primA) + deformableStateB = PhysxSchema.PhysxDeformableAPI(primB) + + deformableEnabled = deformableStateA.GetDeformableEnabledAttr().Get() + deformableStateB.GetDeformableEnabledAttr().Set(deformableEnabled) + + enableCCD = deformableStateA.GetEnableCCDAttr().Get() + deformableStateB.GetEnableCCDAttr().Set(enableCCD) + + maxDepenetrationVelocity = deformableStateA.GetMaxDepenetrationVelocityAttr().Get() + deformableStateB.GetMaxDepenetrationVelocityAttr().Set(maxDepenetrationVelocity) + + velocityDamping = deformableStateA.GetVertexVelocityDampingAttr().Get() + deformableStateB.GetVertexVelocityDampingAttr().Set(velocityDamping) + + selfCollision = deformableStateA.GetSelfCollisionAttr().Get() + deformableStateB.GetSelfCollisionAttr().Set(selfCollision) + + selfCollisionFilterDistance = deformableStateA.GetSelfCollisionFilterDistanceAttr().Get() + deformableStateB.GetSelfCollisionFilterDistanceAttr().Set(selfCollisionFilterDistance) + + settlingThreshold = deformableStateA.GetSettlingThresholdAttr().Get() + deformableStateB.GetSettlingThresholdAttr().Set(settlingThreshold) + + sleepDamping = deformableStateA.GetSleepDampingAttr().Get() + deformableStateB.GetSleepDampingAttr().Set(sleepDamping) + + sleepThreshold = deformableStateA.GetSleepThresholdAttr().Get() + deformableStateB.GetSleepThresholdAttr().Set(sleepThreshold) + + solverPositionIteration = deformableStateA.GetSolverPositionIterationCountAttr().Get() + deformableStateB.GetSolverPositionIterationCountAttr().Set(solverPositionIteration) + + collisionStateA = PhysxSchema.PhysxCollisionAPI(primA) + collisionStateB = PhysxSchema.PhysxCollisionAPI(primB) + + contactOffset = collisionStateA.GetContactOffsetAttr().Get() + collisionStateB.GetContactOffsetAttr().Set(contactOffset) + + restOffset = collisionStateA.GetRestOffsetAttr().Get() + collisionStateB.GetRestOffsetAttr().Set(restOffset) + + if primA.HasAttribute("physxDeformable:simulationHexahedralResolution"): + value = primA.GetAttribute("physxDeformable:simulationHexahedralResolution").Get() + primB.GetAttribute("physxDeformable:simulationHexahedralResolution").Set(value) + + +def copy_transform(primA, primB): + xformableA = UsdGeom.Xformable(primA) + ops = xformableA.GetOrderedXformOps() + + xformableB = UsdGeom.Xformable(primB) + xformableB.ClearXformOpOrder() + + for op in ops: + new_op = xformableB.AddXformOp(op.GetOpType(), op.GetPrecision()) + new_op.Set(op.Get()) + + +def copy_material(primA, primB): + from pxr import UsdShade + + looks_material_path = UsdShade.MaterialBindingAPI(primA).GetDirectBindingRel("").GetTargets()[0] + omni.kit.commands.execute( + "BindMaterial", + material_path=str(looks_material_path), + prim_path=[primB.GetPath()], + strength=["weakerThanDescendants"], + material_purpose="", + ) + + physics_material_path = UsdShade.MaterialBindingAPI(primA).GetDirectBindingRel("physics").GetTargets()[0] + omni.kit.commands.execute( + "BindMaterial", + material_path=str(physics_material_path), + prim_path=[primB.GetPath()], + strength=["weakerThanDescendants"], + material_purpose="physics", + ) + + +def compute_A_position_in_B_space(position, body0_prim, body1_prim): + timeline = omni.timeline.get_timeline_interface() + current_time = timeline.get_current_time() + usd_current_time = Usd.TimeCode(current_time) + body1_xform = UsdGeom.Xformable(body1_prim) + body0_xform = UsdGeom.Xformable(body0_prim) + + Trans0 = body0_xform.ComputeLocalToWorldTransform(usd_current_time) + p0_in_world = Trans0.Transform(position) + + Trans1 = body1_xform.ComputeLocalToWorldTransform(usd_current_time) + p0_in_p1 = Trans1.GetInverse().Transform(p0_in_world) + + return p0_in_p1 + + +def compute_prim_position_in_world(position, prim: Usd.Prim): + timeline = omni.timeline.get_timeline_interface() + current_time = timeline.get_current_time() + usd_current_time = Usd.TimeCode(current_time) + + prim_xform = UsdGeom.Xformable(prim) + Trans0 = prim_xform.ComputeLocalToWorldTransform(usd_current_time) + position_in_world = Trans0.Transform(position) + + return position_in_world + + +def get_knife_plane_in_world(knife_prim): + # p6 p7 + # p4 p5 + # p2 p3 + # p0 p1 + knife_cube = UsdGeom.Cube(knife_prim) + extent = knife_cube.GetExtentAttr().Get() + center = (extent[0] + extent[1]) / 2 + + p0 = extent[0] + p7 = extent[1] + p3 = Gf.Vec3f(extent[1][0], extent[1][1], extent[0][0]) + p1 = Gf.Vec3f(p3[0], extent[0][1], extent[0][2]) + + center_in_world = compute_prim_position_in_world(center, knife_prim) + p0_in_world = compute_prim_position_in_world(p0, knife_prim) + p7_in_world = compute_prim_position_in_world(p7, knife_prim) + p3_in_world = compute_prim_position_in_world(p3, knife_prim) + p1_in_world = compute_prim_position_in_world(p1, knife_prim) + + length13 = (p1_in_world - p3_in_world).GetLength() + length01 = (p0_in_world - p1_in_world).GetLength() + length37 = (p3_in_world - p7_in_world).GetLength() + + knife_normal = None + if length13 < length01 and length13 < length37: + knife_normal = p1_in_world - p3_in_world + elif length01 < length13 and length01 < length37: + knife_normal = p0_in_world - p1_in_world + elif length37 < length13 and length37 < length01: + knife_normal = p3_in_world - p7_in_world + + return center_in_world, knife_normal.GetNormalized() + + +def compute_world_position_in_prim(wolrd_position, prim: Usd.Prim): + timeline = omni.timeline.get_timeline_interface() + current_time = timeline.get_current_time() + usd_current_time = Usd.TimeCode(current_time) + + prim_xform = UsdGeom.Xformable(prim) + Trans0 = prim_xform.ComputeLocalToWorldTransform(usd_current_time) + position_in_prim = Trans0.GetInverse().Transform(wolrd_position) + + return position_in_prim + + +def compute_world_dir_in_prim(wolrd_dir, prim: Usd.Prim): + timeline = omni.timeline.get_timeline_interface() + current_time = timeline.get_current_time() + usd_current_time = Usd.TimeCode(current_time) + + prim_xform = UsdGeom.Xformable(prim) + Trans0 = prim_xform.ComputeLocalToWorldTransform(usd_current_time) + position_in_prim = Trans0.GetInverse().TransformDir(wolrd_dir) + + return position_in_prim.GetNormalized() + + +def install_package(package_name, install_name=None): + """ + Installs a package from a local wheel file or Ali Cloud image. + Parameters: + package_name: The name of the package to import + install_name: The name of the package to install (if different from the import name) + """ + if install_name is None: + install_name = package_name + + print(f"{package_name} module not found. Trying to install...", flush=True) + try: + # Get the current operating system + current_os = platform.system() + + # Local wheel file path + current_dir = os.path.dirname(os.path.realpath(__file__)) + whls_dir = os.path.join(current_dir, "whls") + + # Find the wheel file in the whls directory + package_whl = None + if os.path.exists(whls_dir): + for file in os.listdir(whls_dir): + if file.startswith(install_name) and file.endswith(".whl"): + # Filter wheel files based on operating system + is_windows_match = current_os == "Windows" and ( + "win" in file.lower() or "py3-none-any" in file.lower() + ) + is_linux_match = current_os == "Linux" and ( + "linux" in file.lower() or "manylinux" in file.lower() or "py3-none-any" in file.lower() + ) + if is_windows_match or is_linux_match: + package_whl = os.path.join(whls_dir, file) + break + + if package_whl and os.path.exists(package_whl): + # Install from a local wheel file + print( + f"Install from a local wheel file{install_name}: {package_whl}", + flush=True, + ) + pip.main(["install", package_whl]) + else: + # If the wheel file is not found, install it from the Alibaba Cloud image. + print( + f"No local wheel file found for {current_os}. Installing {install_name} from Alibaba Cloud image...", + flush=True, + ) + pip.main([ + "install", + install_name, + "-i", + "https://mirrors.aliyun.com/pypi/simple/", + "--trusted-host", + "mirrors.aliyun.com", + ]) + + print(f"{install_name} was successfully installed", flush=True) + return True + except Exception as e: + print(f"Installation of {install_name} failed: {str(e)}", flush=True) + return False diff --git a/source/leisaac/leisaac/utils/domain_randomization.py b/source/leisaac/leisaac/utils/domain_randomization.py index 5f3584c1..1903bc78 100644 --- a/source/leisaac/leisaac/utils/domain_randomization.py +++ b/source/leisaac/leisaac/utils/domain_randomization.py @@ -20,6 +20,20 @@ def randomize_object_uniform( ) +def randomize_deformable_object_uniform( + name: str, + pose_range: dict[str, tuple[float, float]], + velocity_range: dict[str, tuple[float, float]] | None = None, +) -> EventTerm: + if velocity_range is None: + velocity_range = {} + return EventTerm( + func=mdp.reset_nodal_state_uniform, + mode="reset", + params={"position_range": pose_range, "velocity_range": velocity_range, "asset_cfg": SceneEntityCfg(name)}, + ) + + def randomize_camera_uniform( name: str, pose_range: dict[str, tuple[float, float]], convention: Literal["ros", "opengl", "world"] = "ros" ) -> EventTerm: @@ -51,3 +65,50 @@ def randomize_particle_object_uniform( def domain_randomization(env_cfg, random_options: list[EventTerm]): for idx, event_item in enumerate(random_options): setattr(env_cfg.events, f"domain_randomize_{idx}", event_item) + + +def randomize_mixed_objects_uniform( + rigid_names: list[str] | None = None, + deformable_names: list[str] | None = None, + pose_range: dict[str, tuple[float, float]] | None = None, + velocity_range: dict[str, tuple[float, float]] | None = None, +) -> EventTerm: + + if rigid_names is None: + rigid_names = [] + if deformable_names is None: + deformable_names = [] + if pose_range is None: + pose_range = {} + if velocity_range is None: + velocity_range = {} + return EventTerm( + func=enhance_mdp.reset_mixed_objects_uniform, + mode="reset", + params={ + "pose_range": pose_range, + "velocity_range": velocity_range, + "rigid_asset_cfg": [SceneEntityCfg(name) for name in rigid_names], + "deformable_asset_cfg": [SceneEntityCfg(name) for name in deformable_names], + }, + ) + + +def randomize_cuttable_object_uniform( + attr_name: str, + pose_range: dict[str, tuple[float, float]], +) -> EventTerm: + """Randomize a cuttable object's pose. + + Args: + attr_name: The attribute name of the CuttableObject on the env (e.g., "cuttable_sausage"). + pose_range: Dict with keys x, y, z, roll, pitch, yaw and (min, max) tuple values. + """ + return EventTerm( + func=enhance_mdp.randomize_cuttable_object_uniform, + mode="reset", + params={ + "attr_name": attr_name, + "pose_range": pose_range, + }, + ) diff --git a/source/leisaac/leisaac/utils/general_assets.py b/source/leisaac/leisaac/utils/general_assets.py index 2ee01fe9..cde004af 100644 --- a/source/leisaac/leisaac/utils/general_assets.py +++ b/source/leisaac/leisaac/utils/general_assets.py @@ -1,4 +1,4 @@ -from pxr import Usd, UsdGeom, UsdPhysics +from pxr import PhysxSchema, Usd, UsdGeom, UsdPhysics def get_all_prims(stage, prim=None, prims_list=None): @@ -17,6 +17,8 @@ def classify_prim(prim): return "Articulation" elif prim.HasAPI(UsdPhysics.RigidBodyAPI): return "RigidBody" + elif prim.HasAPI(PhysxSchema.PhysxDeformableBodyAPI): + return "DeformableBody" else: return "Normal" @@ -29,6 +31,18 @@ def is_rigidbody(prim): return prim.HasAPI(UsdPhysics.RigidBodyAPI) +def is_kinematic_rigidbody(prim): + if not prim.HasAPI(UsdPhysics.RigidBodyAPI): + return False + rigid_body_api = UsdPhysics.RigidBodyAPI(prim) + kinematic_attr = rigid_body_api.GetKinematicEnabledAttr() + return kinematic_attr and kinematic_attr.Get() + + +def is_deformable(prim): + return prim.HasAPI(PhysxSchema.PhysxDeformableBodyAPI) + + def get_all_joints(stage): joints = [] @@ -105,8 +119,12 @@ def get_all_joints_without_fixed(articulation_prim): import isaacsim.core.utils.prims as prim_utils from isaaclab.assets.articulation import ArticulationCfg +from isaaclab.assets.deformable_object import DeformableObjectCfg from isaaclab.assets.rigid_object import RigidObjectCfg -from isaaclab.sim.spawners.spawner_cfg import RigidObjectSpawnerCfg +from isaaclab.sim.spawners.spawner_cfg import ( + DeformableObjectSpawnerCfg, + RigidObjectSpawnerCfg, +) from isaaclab.sim.utils import clone @@ -163,6 +181,9 @@ def parse_usd_and_create_subassets(usd_path, env_cfg, specific_name_list=None, e if is_rigidbody(prim) and match_specific_name(prim.GetPath().pathString, specific_name_list, exclude_name_list): if prim in articulation_sub_prims: continue + # Skip kinematic rigid bodies (e.g., collision meshes for Gaussian scenes) + if is_kinematic_rigidbody(prim): + continue pos, rot = get_prim_pos_rot(prim) orin_prim_path = prim.GetPath().pathString name = orin_prim_path.split("/")[-1] @@ -182,3 +203,29 @@ def parse_usd_and_create_subassets(usd_path, env_cfg, specific_name_list=None, e ), ) setattr(env_cfg.scene, name, rigidcfg) + + for prim in prims: + if is_deformable(prim) and match_specific_name( + prim.GetPath().pathString, specific_name_list, exclude_name_list + ): + if prim in articulation_sub_prims: + continue + pos, rot = get_prim_pos_rot(prim) + orin_prim_path = prim.GetPath().pathString + name = orin_prim_path.split("/")[-1] + if name not in create_attr_record: + create_attr_record[name] = 0 + else: + create_attr_record[name] += 1 + name = f"{name}_{create_attr_record[name]}" + sub_prim_path = orin_prim_path[orin_prim_path.find("/", 1) + 1 :] + prim_path = f"{{ENV_REGEX_NS}}/Scene/{sub_prim_path}" + defcfg = DeformableObjectCfg( + prim_path=prim_path, + spawn=DeformableObjectSpawnerCfg(func=spawn_from_prim_path), + init_state=DeformableObjectCfg.InitialStateCfg( + pos=pos, + rot=rot, + ), + ) + setattr(env_cfg.scene, name, defcfg)