From 8906c1511bcf8ba1799e399a334f96adb712b2b9 Mon Sep 17 00:00:00 2001 From: "zhengtong.zhang" Date: Thu, 15 Jan 2026 17:31:19 +0800 Subject: [PATCH 1/2] feat: add_burger and sausage --- scripts/tutorials/marble_compose.py | 170 +++++--- .../leisaac/leisaac/assets/robots/lerobot.py | 52 +++ .../leisaac/leisaac/assets/scenes/kitchen.py | 11 +- .../leisaac/leisaac/assets/scenes/loftroom.py | 17 + .../leisaac/enhance/envs/mdp/events.py | 48 +++ .../tasks/assemble_hamburger/__init__.py | 10 +- .../assemble_hamburger_bi_arm_env_cfg.py | 99 +++++ .../assemble_hamhurger_bi_arm_env_cfg.py | 45 -- .../assemble_hamhurger_env_cfg.py | 44 -- .../assemble_hamburger/direct/__init__.py | 0 .../direct/assemble_hamburger_bi_arm_env.py | 84 ++++ .../tasks/assemble_hamburger/mdp/__init__.py | 10 + .../assemble_hamburger/mdp/terminations.py | 54 +++ .../leisaac/tasks/sausage_cut/__init__.py | 19 + .../tasks/sausage_cut/direct/__init__.py | 1 + .../direct/sausage_cut_bi_arm_env.py | 215 ++++++++++ .../leisaac/tasks/sausage_cut/mdp/__init__.py | 4 + .../tasks/sausage_cut/mdp/terminations.py | 44 ++ .../sausage_cut/sausage_cut_bi_arm_env_cfg.py | 68 +++ .../leisaac/utils/collision_checker.py | 147 +++++++ source/leisaac/leisaac/utils/cutMeshNode.py | 391 ++++++++++++++++++ .../leisaac/utils/domain_randomization.py | 41 ++ .../leisaac/leisaac/utils/general_assets.py | 40 +- 23 files changed, 1470 insertions(+), 144 deletions(-) create mode 100644 source/leisaac/leisaac/assets/scenes/loftroom.py create mode 100644 source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamburger_bi_arm_env_cfg.py delete mode 100644 source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamhurger_bi_arm_env_cfg.py delete mode 100644 source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamhurger_env_cfg.py create mode 100644 source/leisaac/leisaac/tasks/assemble_hamburger/direct/__init__.py create mode 100644 source/leisaac/leisaac/tasks/assemble_hamburger/direct/assemble_hamburger_bi_arm_env.py create mode 100644 source/leisaac/leisaac/tasks/assemble_hamburger/mdp/__init__.py create mode 100644 source/leisaac/leisaac/tasks/assemble_hamburger/mdp/terminations.py create mode 100644 source/leisaac/leisaac/tasks/sausage_cut/__init__.py create mode 100644 source/leisaac/leisaac/tasks/sausage_cut/direct/__init__.py create mode 100644 source/leisaac/leisaac/tasks/sausage_cut/direct/sausage_cut_bi_arm_env.py create mode 100644 source/leisaac/leisaac/tasks/sausage_cut/mdp/__init__.py create mode 100644 source/leisaac/leisaac/tasks/sausage_cut/mdp/terminations.py create mode 100644 source/leisaac/leisaac/tasks/sausage_cut/sausage_cut_bi_arm_env_cfg.py create mode 100644 source/leisaac/leisaac/utils/collision_checker.py create mode 100644 source/leisaac/leisaac/utils/cutMeshNode.py diff --git a/scripts/tutorials/marble_compose.py b/scripts/tutorials/marble_compose.py index 5f32ca5c..e5a476f1 100644 --- a/scripts/tutorials/marble_compose.py +++ b/scripts/tutorials/marble_compose.py @@ -8,9 +8,11 @@ """ import argparse +import os from pathlib import Path import numpy as np +from scipy.spatial.transform import Rotation as R # Configuration @@ -28,6 +30,12 @@ "cube": { "single": "LeIsaac-SO101-LiftCube-v0", }, + "burger": { + "dual": "LeIsaac-SO101-AssembleHamburger-BiArm-v0", + }, + "sausage": { + "dual": "LeIsaac-SO101-SausageCut-BiArm-v0", + }, } TASK_CONFIG = { @@ -58,7 +66,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 +77,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", + }, } @@ -130,25 +157,33 @@ def compute_scene_transform(orig_pos, orig_quat, target_pos, target_quat): return matrix_to_pose(T_inv) +def compute_lookat(eye, euler_deg, distance=2.0): + """Compute lookat point from eye position and euler angles""" + 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})") + + # USD Utilities 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 +191,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 +208,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 +256,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 +284,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 +293,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 +327,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 +358,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 +410,26 @@ 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)", + ) + # Camera arguments + parser.add_argument("--camera-eye", type=float, nargs=3, metavar=("X", "Y", "Z"), help="Camera eye position") + parser.add_argument( + "--camera-euler", type=float, nargs=3, metavar=("R", "P", "Y"), help="Camera rotation (degrees)" ) args = parser.parse_args() - # Set LEISAAC_ASSETS_ROOT to the provided assets base path - import os + # Handle camera computation + if args.camera_eye and args.camera_euler: + compute_lookat(args.camera_eye, args.camera_euler) + exit(0) + # Set LEISAAC_ASSETS_ROOT to the provided assets base path os.environ["LEISAAC_ASSETS_ROOT"] = args.assets_base print(f"[INFO] Set LEISAAC_ASSETS_ROOT={args.assets_base}") @@ -374,7 +450,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..e65b94e1 100644 --- a/source/leisaac/leisaac/assets/scenes/kitchen.py +++ b/source/leisaac/leisaac/assets/scenes/kitchen.py @@ -15,10 +15,19 @@ ) ) -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, + ) +) 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/envs/mdp/events.py b/source/leisaac/leisaac/enhance/envs/mdp/events.py index 2853083c..daf7ee38 100644 --- a/source/leisaac/leisaac/enhance/envs/mdp/events.py +++ b/source/leisaac/leisaac/enhance/envs/mdp/events.py @@ -122,3 +122,51 @@ 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) 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..e509863c --- /dev/null +++ b/source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamburger_bi_arm_env_cfg.py @@ -0,0 +1,99 @@ +""" +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() + + # Simulation configuration + render_cfg: sim_utils.RenderCfg = sim_utils.RenderCfg(rendering_mode="quality", antialiasing_mode="Off") + sim: SimulationCfg = SimulationCfg(dt=1 / 120, 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.5, -5.0, 1.6) + self.viewer.lookat = (3.7, -6.15, 0.84) + + # 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 = 2 + self.dynamic_reset_gripper_effort_limit = False + + # 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..931f24b5 --- /dev/null +++ b/source/leisaac/leisaac/tasks/assemble_hamburger/direct/assemble_hamburger_bi_arm_env.py @@ -0,0 +1,84 @@ +""" +Assemble Hamburger Bi-Arm Direct Environment (LeIsaac Pattern) +""" + +import isaaclab.sim as sim_utils +import torch +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) + + # Render configuration - match manager-based env for proper material colors + render_cfg: sim_utils.RenderCfg = sim_utils.RenderCfg(rendering_mode="quality", antialiasing_mode="Off") + sim: SimulationCfg = SimulationCfg(dt=1 / 120, 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.5, -5.0, 1.6) + self.viewer.lookat = (3.7, -6.15, 0.84) + + 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 = 2 + + # 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..501a63a1 --- /dev/null +++ b/source/leisaac/leisaac/tasks/sausage_cut/__init__.py @@ -0,0 +1,19 @@ +import gymnasium as gym + +gym.register( + id="LeIsaac-SO101-SausageCut-BiArm-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.sausage_cut_bi_arm_env_cfg:SausageCutBiArmEnvCfg", + }, +) + +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..3fa6e10c --- /dev/null +++ b/source/leisaac/leisaac/tasks/sausage_cut/direct/sausage_cut_bi_arm_env.py @@ -0,0 +1,215 @@ +""" +Sausage Cut Bi-Arm Direct Environment + +A simple Direct environment for the sausage cutting task. +Only overrides cutting-specific logic; robot control and observations +are handled by BiArmTaskDirectEnv base class. +""" + +import os +from pathlib import Path +from types import SimpleNamespace + +import isaaclab.sim as sim_utils +import numpy as np +import omni +import torch +from isaaclab.assets import AssetBaseCfg +from isaaclab.sim import SimulationCfg +from isaaclab.utils import configclass +from isaacsim.core.utils.prims import delete_prim +from leisaac.assets.scenes.kitchen import KITCHEN_WITH_SAUSAGE_USD_PATH +from leisaac.utils.collision_checker import Collision_Checker +from leisaac.utils.constant import ASSETS_ROOT +from leisaac.utils.cutMeshNode import cutMeshNode +from leisaac.utils.general_assets import parse_usd_and_create_subassets +from scipy.spatial.transform import Rotation as R + +from ...template import BiArmTaskDirectEnv, BiArmTaskDirectEnvCfg +from .. import mdp +from ..sausage_cut_bi_arm_env_cfg import SausageCutBiArmSceneCfg + +# Sausage USD path (for dynamic reload after cutting) +SAUSAGE_USD_PATH = str(Path(ASSETS_ROOT) / "scenes/kitchen_with_sausage/objects/Sausage001/Sausage001.usd") + +# Sausage initial pose +SAUSAGE_BASE_T = (3.6612, -6.236, 0.84059) +SAUSAGE_BASE_Q_WXYZ = (-0.23287, -0.02628, 0.02471, 0.97184) + + +@configclass +class SausageCutBiArmEnvCfg(BiArmTaskDirectEnvCfg): + """Direct env configuration for the sausage cut task.""" + + scene: SausageCutBiArmSceneCfg = SausageCutBiArmSceneCfg(env_spacing=4.0) + + # CRITICAL: use_fabric=False is required for mesh cutting + render_cfg: sim_utils.RenderCfg = sim_utils.RenderCfg(rendering_mode="quality", antialiasing_mode="FXAA") + sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=1, render=render_cfg, use_fabric=False) + + def __post_init__(self) -> None: + super().__post_init__() + + self.viewer.eye = (2.5, -5.0, 1.6) + self.viewer.lookat = (3.7, -6.15, 0.84) + + 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) + + self.decimation = 2 + self.dynamic_reset_gripper_effort_limit = False + + 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"] # Handled separately due to cutting + ) + + +class SausageCutBiArmEnv(BiArmTaskDirectEnv): + """Direct env for sausage cutting task.""" + + cfg: SausageCutBiArmEnvCfg + + def __init__(self, cfg: SausageCutBiArmEnvCfg, render_mode: str | None = None, **kwargs): + # Cutting system state + self.dummy_db = SimpleNamespace( + inputs=SimpleNamespace(cut_mesh_path=None, knife_mesh_path=None, cutEventIn=False), + internal_state=cutMeshNode.internal_state(), + ) + self.sausage_count = 1 + self.last_if_collision = False + super().__init__(cfg, render_mode, **kwargs) + + def _setup_scene(self): + """Setup scene and initialize cutting system.""" + super()._setup_scene() + + env_prim_path = self.scene.env_prim_paths[0] + self.sausage_prim_path = f"{env_prim_path}/Scene/Sausage001" + + # Set cutting mesh paths + self.dummy_db.inputs.cut_mesh_path = f"{self.sausage_prim_path}/Sausage001" + self.dummy_db.inputs.knife_mesh_path = f"{env_prim_path}/Right_Robot/gripper/Knife/Knife/Cube" + + # Initialize collision checker + self.stage = omni.usd.get_context().get_stage() + + # Debug: print paths and verify prims exist + print(f"[DEBUG] env_prim_path: {env_prim_path}") + print(f"[DEBUG] sausage_prim_path: {self.sausage_prim_path}") + print(f"[DEBUG] cut_mesh_path: {self.dummy_db.inputs.cut_mesh_path}") + print(f"[DEBUG] knife_mesh_path: {self.dummy_db.inputs.knife_mesh_path}") + + trigger_path = f"{self.sausage_prim_path}/Trigger/Cube" + knife_path = f"{env_prim_path}/Right_Robot/gripper/Knife/Knife/Knife002" + print(f"[DEBUG] trigger_path: {trigger_path}") + print(f"[DEBUG] knife_collision_path: {knife_path}") + + # Verify prims exist + trigger_prim = self.stage.GetPrimAtPath(trigger_path) + knife_prim = self.stage.GetPrimAtPath(knife_path) + print(f"[DEBUG] trigger_prim exists: {trigger_prim.IsValid() if trigger_prim else False}") + print(f"[DEBUG] knife_prim exists: {knife_prim.IsValid() if knife_prim else False}") + + self.collision_checker = Collision_Checker( + stage=self.stage, + prim_path0=trigger_path, + prim_path1=knife_path, + ) + + # Spawn sausage + if os.path.exists(SAUSAGE_USD_PATH): + cfg = sim_utils.UsdFileCfg(usd_path=SAUSAGE_USD_PATH) + cfg.func(self.sausage_prim_path, cfg, translation=SAUSAGE_BASE_T, orientation=SAUSAGE_BASE_Q_WXYZ) + print(f"[DEBUG] Sausage spawned at {self.sausage_prim_path}") + else: + print(f"[ERROR] Sausage USD not found: {SAUSAGE_USD_PATH}") + + def _apply_action(self) -> None: + """Apply robot actions and handle cutting logic.""" + super()._apply_action() + + # Collision detection with debounce (protected from mesh modification errors) + try: + if_collision, _, _ = self.collision_checker.meshes_aabb_collide() + if if_collision and if_collision == self.last_if_collision: + if_collision = False + self.last_if_collision = if_collision + except Exception as e: + # Mesh may be invalid during/after cutting, skip collision check + print(f"[DEBUG] Collision check exception: {e}") + if_collision = False + + # Trigger mesh cutting + self.dummy_db.inputs.cutEventIn = if_collision + try: + cutMeshNode.compute(self.dummy_db) + except Exception as e: + print(f"[ERROR] cutMeshNode.compute failed: {e}") + import traceback + + traceback.print_exc() + + if if_collision: + self.sausage_count = self._count_sausage_meshes() + print(f"[DEBUG] Collision detected, sausage_count: {self.sausage_count}") + + def _check_success(self) -> torch.Tensor: + """Check if sausage has been cut.""" + return mdp.sausage_cut(env=self, min_sausage_count=2, check_rest_pose=False) + + def _count_sausage_meshes(self) -> int: + """Count active sausage mesh pieces.""" + try: + prim = self.stage.GetPrimAtPath(self.sausage_prim_path) + if prim and prim.IsValid(): + return len([p for p in prim.GetChildren() if p.IsActive()]) + except Exception: + pass + return 1 + + def _reset_sausage(self): + """Reset sausage by deleting and recreating with randomized pose.""" + delete_prim(self.sausage_prim_path) + + # Randomize pose (translation: +/-4cm, rotation: +/-20deg on z-axis) + rng = np.random.default_rng() + t_new = np.array(SAUSAGE_BASE_T) + np.array([rng.uniform(-0.04, 0.04), rng.uniform(-0.04, 0.04), 0.0]) + + q_wxyz = np.array(SAUSAGE_BASE_Q_WXYZ) + r_base = R.from_quat([q_wxyz[1], q_wxyz[2], q_wxyz[3], q_wxyz[0]]) + r_new = r_base * R.from_rotvec(np.deg2rad(rng.uniform(-20, 20)) * np.array([0, 0, 1])) + q_new = r_new.as_quat() + q_new_wxyz = np.array([q_new[3], q_new[0], q_new[1], q_new[2]]) + + # Recreate sausage + if os.path.exists(SAUSAGE_USD_PATH): + cfg = sim_utils.UsdFileCfg(usd_path=SAUSAGE_USD_PATH) + cfg.func(self.sausage_prim_path, cfg, translation=t_new, orientation=q_new_wxyz) + + # Reset cutting state + self.dummy_db.inputs.cut_mesh_path = f"{self.sausage_prim_path}/Sausage001" + self.dummy_db.internal_state = cutMeshNode.internal_state() + self.dummy_db.inputs.cutEventIn = False + self.sausage_count = 1 + self.last_if_collision = False + + # Reinitialize collision checker (prim was deleted and recreated) + env_prim_path = self.scene.env_prim_paths[0] + self.collision_checker = Collision_Checker( + stage=self.stage, + prim_path0=f"{self.sausage_prim_path}/Trigger/Cube", + prim_path1=f"{env_prim_path}/Right_Robot/gripper/Knife/Knife/Knife002", + ) + + def _reset_idx(self, env_ids): + """Reset environment.""" + super()._reset_idx(env_ids) + self._reset_sausage() diff --git a/source/leisaac/leisaac/tasks/sausage_cut/mdp/__init__.py b/source/leisaac/leisaac/tasks/sausage_cut/mdp/__init__.py new file mode 100644 index 00000000..c0ac1099 --- /dev/null +++ b/source/leisaac/leisaac/tasks/sausage_cut/mdp/__init__.py @@ -0,0 +1,4 @@ +from isaaclab.envs.mdp import * +from leisaac.enhance.envs.mdp import * + +from .terminations import sausage_cut diff --git a/source/leisaac/leisaac/tasks/sausage_cut/mdp/terminations.py b/source/leisaac/leisaac/tasks/sausage_cut/mdp/terminations.py new file mode 100644 index 00000000..155d2d18 --- /dev/null +++ b/source/leisaac/leisaac/tasks/sausage_cut/mdp/terminations.py @@ -0,0 +1,44 @@ +import torch +from isaaclab.envs import DirectRLEnv, ManagerBasedEnv +from leisaac.utils.robot_utils import is_so101_at_rest_pose + + +def sausage_cut( + env: ManagerBasedEnv | DirectRLEnv, + min_sausage_count: int = 2, + check_rest_pose: bool = True, +) -> torch.Tensor: + """Determine if the sausage cutting task is completed successfully. + + This function evaluates the success conditions for the sausage cutting task: + 1. Sausage has been cut into at least `min_sausage_count` pieces + 2. (Optional) Robot arms return to the rest pose + + Args: + env: The RL environment instance. + min_sausage_count: Minimum number of sausage pieces required for success (default: 2). + check_rest_pose: Whether to check if robots are at rest pose (default: True). + + Returns: + Boolean tensor indicating which environments have completed the task. + """ + done = torch.zeros(env.num_envs, dtype=torch.bool, device=env.device) + + # Check if sausage has been cut (access from environment's custom attribute) + if hasattr(env, "sausage_count"): + sausage_cut_done = env.sausage_count >= min_sausage_count + done = torch.logical_or(done, sausage_cut_done) + + # Optionally check if robots are at rest pose + if check_rest_pose and hasattr(env.scene, "articulations"): + left_arm = env.scene.articulations.get("left_arm") or env.scene.get("left_arm") + right_arm = env.scene.articulations.get("right_arm") or env.scene.get("right_arm") + + if left_arm is not None and right_arm is not None: + is_rest = torch.logical_and( + is_so101_at_rest_pose(left_arm.data.joint_pos, left_arm.data.joint_names), + is_so101_at_rest_pose(right_arm.data.joint_pos, right_arm.data.joint_names), + ) + done = torch.logical_and(done, is_rest) + + return done diff --git a/source/leisaac/leisaac/tasks/sausage_cut/sausage_cut_bi_arm_env_cfg.py b/source/leisaac/leisaac/tasks/sausage_cut/sausage_cut_bi_arm_env_cfg.py new file mode 100644 index 00000000..aee31b41 --- /dev/null +++ b/source/leisaac/leisaac/tasks/sausage_cut/sausage_cut_bi_arm_env_cfg.py @@ -0,0 +1,68 @@ +import isaaclab.sim as sim_utils +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, +) +from leisaac.utils.general_assets import parse_usd_and_create_subassets + +from ..template import ( + BiArmObservationsCfg, + BiArmTaskEnvCfg, + BiArmTaskSceneCfg, + BiArmTerminationsCfg, +) + + +@configclass +class SausageCutBiArmSceneCfg(BiArmTaskSceneCfg): + """Scene configuration for the sausage cut task using two arms.""" + + # Use merged USD that includes sausage + scene: AssetBaseCfg = KITCHEN_WITH_SAUSAGE_CFG.replace(prim_path="{ENV_REGEX_NS}/Scene") + + # Replace right arm with knife-equipped robot (use standard config to match lehome) + right_arm: ArticulationCfg = SO101_KINFE_CFG.replace(prim_path="{ENV_REGEX_NS}/Right_Robot") + + +@configclass +class SausageCutBiArmEnvCfg(BiArmTaskEnvCfg): + """Configuration for the sausage cut environment.""" + + scene: SausageCutBiArmSceneCfg = SausageCutBiArmSceneCfg(env_spacing=4.0) + + observations: BiArmObservationsCfg = BiArmObservationsCfg() + + terminations: BiArmTerminationsCfg = BiArmTerminationsCfg() + + # Simulation configuration - CRITICAL: use_fabric=False is required for DeformableBody/cutting + render_cfg: sim_utils.RenderCfg = sim_utils.RenderCfg(rendering_mode="quality", antialiasing_mode="FXAA") + sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=1, render=render_cfg, use_fabric=False) + + def __post_init__(self) -> None: + super().__post_init__() + + # Camera positions + self.viewer.eye = (2.5, -5.0, 1.6) + self.viewer.lookat = (3.7, -6.15, 0.84) + + # Robot initial positions + 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 + self.decimation = 2 + self.dynamic_reset_gripper_effort_limit = False + + # Parse USD for additional assets (ChoppingBlock, etc.) + parse_usd_and_create_subassets( + KITCHEN_WITH_SAUSAGE_USD_PATH, + self, + exclude_name_list=["Sausage"], # Sausage is handled separately due to cutting + ) 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..f87be3ce 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,30 @@ 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], + }, + ) diff --git a/source/leisaac/leisaac/utils/general_assets.py b/source/leisaac/leisaac/utils/general_assets.py index 2ee01fe9..2fdc63c5 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,10 @@ def is_rigidbody(prim): return prim.HasAPI(UsdPhysics.RigidBodyAPI) +def is_deformable(prim): + return prim.HasAPI(PhysxSchema.PhysxDeformableBodyAPI) + + def get_all_joints(stage): joints = [] @@ -105,8 +111,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 @@ -182,3 +192,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) From dd4c1a2bfe9be0d4d409b6369321ec02b11e7bce Mon Sep 17 00:00:00 2001 From: "zhengtong.zhang" Date: Fri, 16 Jan 2026 18:53:29 +0800 Subject: [PATCH 2/2] feat:modify sausage --- scripts/tutorials/compute_viewpoint.py | 59 ++++ scripts/tutorials/marble_compose.py | 25 +- .../leisaac/leisaac/assets/scenes/kitchen.py | 3 + .../leisaac/enhance/assets/__init__.py | 2 + .../assets/cuttable_object/__init__.py | 4 + .../assets/cuttable_object/cuttable_object.py | 163 +++++++++++ .../cuttable_object/cuttable_object_cfg.py | 26 ++ .../leisaac/enhance/envs/mdp/events.py | 78 ++++++ .../enhance/envs/mdp/recorders/recorders.py | 23 +- .../envs/mdp/recorders/recorders_cfg.py | 14 + .../assemble_hamburger_bi_arm_env_cfg.py | 18 +- .../direct/assemble_hamburger_bi_arm_env.py | 18 +- .../leisaac/tasks/sausage_cut/__init__.py | 9 - .../direct/sausage_cut_bi_arm_env.py | 261 +++++++----------- .../leisaac/tasks/sausage_cut/mdp/__init__.py | 4 - .../tasks/sausage_cut/mdp/terminations.py | 44 --- .../sausage_cut/sausage_cut_bi_arm_env_cfg.py | 68 ----- .../leisaac/utils/domain_randomization.py | 20 ++ .../leisaac/leisaac/utils/general_assets.py | 11 + 19 files changed, 529 insertions(+), 321 deletions(-) create mode 100644 scripts/tutorials/compute_viewpoint.py create mode 100644 source/leisaac/leisaac/enhance/assets/cuttable_object/__init__.py create mode 100644 source/leisaac/leisaac/enhance/assets/cuttable_object/cuttable_object.py create mode 100644 source/leisaac/leisaac/enhance/assets/cuttable_object/cuttable_object_cfg.py delete mode 100644 source/leisaac/leisaac/tasks/sausage_cut/mdp/__init__.py delete mode 100644 source/leisaac/leisaac/tasks/sausage_cut/mdp/terminations.py delete mode 100644 source/leisaac/leisaac/tasks/sausage_cut/sausage_cut_bi_arm_env_cfg.py 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 e5a476f1..9d7dd494 100644 --- a/scripts/tutorials/marble_compose.py +++ b/scripts/tutorials/marble_compose.py @@ -12,7 +12,6 @@ from pathlib import Path import numpy as np -from scipy.spatial.transform import Rotation as R # Configuration @@ -34,7 +33,7 @@ "dual": "LeIsaac-SO101-AssembleHamburger-BiArm-v0", }, "sausage": { - "dual": "LeIsaac-SO101-SausageCut-BiArm-v0", + "dual": "LeIsaac-SO101-SausageCut-BiArm-Direct-v0", }, } @@ -157,18 +156,6 @@ def compute_scene_transform(orig_pos, orig_quat, target_pos, target_quat): return matrix_to_pose(T_inv) -def compute_lookat(eye, euler_deg, distance=2.0): - """Compute lookat point from eye position and euler angles""" - 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})") - - # USD Utilities @@ -416,19 +403,9 @@ def to_relative(abs_path: str) -> str: action="store_true", help="Use dual-arm configuration (uses left_arm as reference)", ) - # Camera arguments - parser.add_argument("--camera-eye", type=float, nargs=3, metavar=("X", "Y", "Z"), help="Camera eye position") - parser.add_argument( - "--camera-euler", type=float, nargs=3, metavar=("R", "P", "Y"), help="Camera rotation (degrees)" - ) args = parser.parse_args() - # Handle camera computation - if args.camera_eye and args.camera_euler: - compute_lookat(args.camera_eye, args.camera_euler) - exit(0) - # Set LEISAAC_ASSETS_ROOT to the provided assets base path os.environ["LEISAAC_ASSETS_ROOT"] = args.assets_base print(f"[INFO] Set LEISAAC_ASSETS_ROOT={args.assets_base}") diff --git a/source/leisaac/leisaac/assets/scenes/kitchen.py b/source/leisaac/leisaac/assets/scenes/kitchen.py index e65b94e1..a8c80f46 100644 --- a/source/leisaac/leisaac/assets/scenes/kitchen.py +++ b/source/leisaac/leisaac/assets/scenes/kitchen.py @@ -31,3 +31,6 @@ 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/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 daf7ee38..261c3fe8 100644 --- a/source/leisaac/leisaac/enhance/envs/mdp/events.py +++ b/source/leisaac/leisaac/enhance/envs/mdp/events.py @@ -170,3 +170,81 @@ def reset_mixed_objects_uniform( 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/assemble_hamburger_bi_arm_env_cfg.py b/source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamburger_bi_arm_env_cfg.py index e509863c..192f5311 100644 --- 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 @@ -58,16 +58,18 @@ class AssembleHamburgerBiArmEnvCfg(BiArmTaskEnvCfg): 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="Off") - sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=1, render=render_cfg, use_fabric=True) + 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.5, -5.0, 1.6) - self.viewer.lookat = (3.7, -6.15, 0.84) + 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) @@ -76,9 +78,15 @@ def __post_init__(self) -> None: 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 = 2 + 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) 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 index 931f24b5..b68c35f7 100644 --- 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 @@ -4,6 +4,7 @@ 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 @@ -27,17 +28,18 @@ 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="Off") - sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=1, render=render_cfg, use_fabric=True) + 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.5, -5.0, 1.6) - self.viewer.lookat = (3.7, -6.15, 0.84) + 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) @@ -45,7 +47,13 @@ def __post_init__(self) -> None: 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 = 2 + 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) diff --git a/source/leisaac/leisaac/tasks/sausage_cut/__init__.py b/source/leisaac/leisaac/tasks/sausage_cut/__init__.py index 501a63a1..3af0cfb3 100644 --- a/source/leisaac/leisaac/tasks/sausage_cut/__init__.py +++ b/source/leisaac/leisaac/tasks/sausage_cut/__init__.py @@ -1,14 +1,5 @@ import gymnasium as gym -gym.register( - id="LeIsaac-SO101-SausageCut-BiArm-v0", - entry_point="isaaclab.envs:ManagerBasedRLEnv", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.sausage_cut_bi_arm_env_cfg:SausageCutBiArmEnvCfg", - }, -) - gym.register( id="LeIsaac-SO101-SausageCut-BiArm-Direct-v0", entry_point=f"{__name__}.direct.sausage_cut_bi_arm_env:SausageCutBiArmEnv", 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 index 3fa6e10c..2a95f682 100644 --- 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 @@ -1,40 +1,40 @@ -""" -Sausage Cut Bi-Arm Direct Environment +"""Sausage cutting task with bi-arm robot (Direct environment).""" -A simple Direct environment for the sausage cutting task. -Only overrides cutting-specific logic; robot control and observations -are handled by BiArmTaskDirectEnv base class. -""" - -import os -from pathlib import Path -from types import SimpleNamespace +from collections.abc import Sequence import isaaclab.sim as sim_utils -import numpy as np -import omni import torch -from isaaclab.assets import AssetBaseCfg +from isaaclab.assets import ArticulationCfg, AssetBaseCfg from isaaclab.sim import SimulationCfg from isaaclab.utils import configclass -from isaacsim.core.utils.prims import delete_prim -from leisaac.assets.scenes.kitchen import KITCHEN_WITH_SAUSAGE_USD_PATH -from leisaac.utils.collision_checker import Collision_Checker -from leisaac.utils.constant import ASSETS_ROOT -from leisaac.utils.cutMeshNode import cutMeshNode +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 scipy.spatial.transform import Rotation as R -from ...template import BiArmTaskDirectEnv, BiArmTaskDirectEnvCfg -from .. import mdp -from ..sausage_cut_bi_arm_env_cfg import SausageCutBiArmSceneCfg +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) -# Sausage USD path (for dynamic reload after cutting) -SAUSAGE_USD_PATH = str(Path(ASSETS_ROOT) / "scenes/kitchen_with_sausage/objects/Sausage001/Sausage001.usd") -# Sausage initial pose -SAUSAGE_BASE_T = (3.6612, -6.236, 0.84059) -SAUSAGE_BASE_Q_WXYZ = (-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 @@ -42,174 +42,113 @@ 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." - # CRITICAL: use_fabric=False is required for mesh cutting + # Render configuration with antialiasing enabled render_cfg: sim_utils.RenderCfg = sim_utils.RenderCfg(rendering_mode="quality", antialiasing_mode="FXAA") - sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=1, render=render_cfg, use_fabric=False) + 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.5, -5.0, 1.6) - self.viewer.lookat = (3.7, -6.15, 0.84) + 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) - self.decimation = 2 + # 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"] # Handled separately due to cutting + 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 sausage cutting task.""" + """Direct env for the sausage cut task.""" cfg: SausageCutBiArmEnvCfg - def __init__(self, cfg: SausageCutBiArmEnvCfg, render_mode: str | None = None, **kwargs): - # Cutting system state - self.dummy_db = SimpleNamespace( - inputs=SimpleNamespace(cut_mesh_path=None, knife_mesh_path=None, cutEventIn=False), - internal_state=cutMeshNode.internal_state(), - ) - self.sausage_count = 1 - self.last_if_collision = False - super().__init__(cfg, render_mode, **kwargs) - def _setup_scene(self): - """Setup scene and initialize cutting system.""" super()._setup_scene() + env_path = self.scene.env_prim_paths[0] + sausage_prim_path = f"{env_path}/Scene/Sausage001" - env_prim_path = self.scene.env_prim_paths[0] - self.sausage_prim_path = f"{env_prim_path}/Scene/Sausage001" - - # Set cutting mesh paths - self.dummy_db.inputs.cut_mesh_path = f"{self.sausage_prim_path}/Sausage001" - self.dummy_db.inputs.knife_mesh_path = f"{env_prim_path}/Right_Robot/gripper/Knife/Knife/Cube" - - # Initialize collision checker - self.stage = omni.usd.get_context().get_stage() - - # Debug: print paths and verify prims exist - print(f"[DEBUG] env_prim_path: {env_prim_path}") - print(f"[DEBUG] sausage_prim_path: {self.sausage_prim_path}") - print(f"[DEBUG] cut_mesh_path: {self.dummy_db.inputs.cut_mesh_path}") - print(f"[DEBUG] knife_mesh_path: {self.dummy_db.inputs.knife_mesh_path}") - - trigger_path = f"{self.sausage_prim_path}/Trigger/Cube" - knife_path = f"{env_prim_path}/Right_Robot/gripper/Knife/Knife/Knife002" - print(f"[DEBUG] trigger_path: {trigger_path}") - print(f"[DEBUG] knife_collision_path: {knife_path}") - - # Verify prims exist - trigger_prim = self.stage.GetPrimAtPath(trigger_path) - knife_prim = self.stage.GetPrimAtPath(knife_path) - print(f"[DEBUG] trigger_prim exists: {trigger_prim.IsValid() if trigger_prim else False}") - print(f"[DEBUG] knife_prim exists: {knife_prim.IsValid() if knife_prim else False}") - - self.collision_checker = Collision_Checker( - stage=self.stage, - prim_path0=trigger_path, - prim_path1=knife_path, + # 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, + ) ) - # Spawn sausage - if os.path.exists(SAUSAGE_USD_PATH): - cfg = sim_utils.UsdFileCfg(usd_path=SAUSAGE_USD_PATH) - cfg.func(self.sausage_prim_path, cfg, translation=SAUSAGE_BASE_T, orientation=SAUSAGE_BASE_Q_WXYZ) - print(f"[DEBUG] Sausage spawned at {self.sausage_prim_path}") - else: - print(f"[ERROR] Sausage USD not found: {SAUSAGE_USD_PATH}") + def initialize(self): + self.cuttable_sausage.initialize() def _apply_action(self) -> None: - """Apply robot actions and handle cutting logic.""" super()._apply_action() - - # Collision detection with debounce (protected from mesh modification errors) - try: - if_collision, _, _ = self.collision_checker.meshes_aabb_collide() - if if_collision and if_collision == self.last_if_collision: - if_collision = False - self.last_if_collision = if_collision - except Exception as e: - # Mesh may be invalid during/after cutting, skip collision check - print(f"[DEBUG] Collision check exception: {e}") - if_collision = False - - # Trigger mesh cutting - self.dummy_db.inputs.cutEventIn = if_collision - try: - cutMeshNode.compute(self.dummy_db) - except Exception as e: - print(f"[ERROR] cutMeshNode.compute failed: {e}") - import traceback - - traceback.print_exc() - - if if_collision: - self.sausage_count = self._count_sausage_meshes() - print(f"[DEBUG] Collision detected, sausage_count: {self.sausage_count}") + self.cuttable_sausage.step() def _check_success(self) -> torch.Tensor: - """Check if sausage has been cut.""" - return mdp.sausage_cut(env=self, min_sausage_count=2, check_rest_pose=False) - - def _count_sausage_meshes(self) -> int: - """Count active sausage mesh pieces.""" - try: - prim = self.stage.GetPrimAtPath(self.sausage_prim_path) - if prim and prim.IsValid(): - return len([p for p in prim.GetChildren() if p.IsActive()]) - except Exception: - pass - return 1 - - def _reset_sausage(self): - """Reset sausage by deleting and recreating with randomized pose.""" - delete_prim(self.sausage_prim_path) - - # Randomize pose (translation: +/-4cm, rotation: +/-20deg on z-axis) - rng = np.random.default_rng() - t_new = np.array(SAUSAGE_BASE_T) + np.array([rng.uniform(-0.04, 0.04), rng.uniform(-0.04, 0.04), 0.0]) - - q_wxyz = np.array(SAUSAGE_BASE_Q_WXYZ) - r_base = R.from_quat([q_wxyz[1], q_wxyz[2], q_wxyz[3], q_wxyz[0]]) - r_new = r_base * R.from_rotvec(np.deg2rad(rng.uniform(-20, 20)) * np.array([0, 0, 1])) - q_new = r_new.as_quat() - q_new_wxyz = np.array([q_new[3], q_new[0], q_new[1], q_new[2]]) - - # Recreate sausage - if os.path.exists(SAUSAGE_USD_PATH): - cfg = sim_utils.UsdFileCfg(usd_path=SAUSAGE_USD_PATH) - cfg.func(self.sausage_prim_path, cfg, translation=t_new, orientation=q_new_wxyz) - - # Reset cutting state - self.dummy_db.inputs.cut_mesh_path = f"{self.sausage_prim_path}/Sausage001" - self.dummy_db.internal_state = cutMeshNode.internal_state() - self.dummy_db.inputs.cutEventIn = False - self.sausage_count = 1 - self.last_if_collision = False - - # Reinitialize collision checker (prim was deleted and recreated) - env_prim_path = self.scene.env_prim_paths[0] - self.collision_checker = Collision_Checker( - stage=self.stage, - prim_path0=f"{self.sausage_prim_path}/Trigger/Cube", - prim_path1=f"{env_prim_path}/Right_Robot/gripper/Knife/Knife/Knife002", - ) + return self.cuttable_sausage.check_success(min_count=2) - def _reset_idx(self, env_ids): - """Reset environment.""" + def _reset_idx(self, env_ids: Sequence[int]): + self.cuttable_sausage.reset(list(env_ids)) super()._reset_idx(env_ids) - self._reset_sausage() + + 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/tasks/sausage_cut/mdp/__init__.py b/source/leisaac/leisaac/tasks/sausage_cut/mdp/__init__.py deleted file mode 100644 index c0ac1099..00000000 --- a/source/leisaac/leisaac/tasks/sausage_cut/mdp/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from isaaclab.envs.mdp import * -from leisaac.enhance.envs.mdp import * - -from .terminations import sausage_cut diff --git a/source/leisaac/leisaac/tasks/sausage_cut/mdp/terminations.py b/source/leisaac/leisaac/tasks/sausage_cut/mdp/terminations.py deleted file mode 100644 index 155d2d18..00000000 --- a/source/leisaac/leisaac/tasks/sausage_cut/mdp/terminations.py +++ /dev/null @@ -1,44 +0,0 @@ -import torch -from isaaclab.envs import DirectRLEnv, ManagerBasedEnv -from leisaac.utils.robot_utils import is_so101_at_rest_pose - - -def sausage_cut( - env: ManagerBasedEnv | DirectRLEnv, - min_sausage_count: int = 2, - check_rest_pose: bool = True, -) -> torch.Tensor: - """Determine if the sausage cutting task is completed successfully. - - This function evaluates the success conditions for the sausage cutting task: - 1. Sausage has been cut into at least `min_sausage_count` pieces - 2. (Optional) Robot arms return to the rest pose - - Args: - env: The RL environment instance. - min_sausage_count: Minimum number of sausage pieces required for success (default: 2). - check_rest_pose: Whether to check if robots are at rest pose (default: True). - - Returns: - Boolean tensor indicating which environments have completed the task. - """ - done = torch.zeros(env.num_envs, dtype=torch.bool, device=env.device) - - # Check if sausage has been cut (access from environment's custom attribute) - if hasattr(env, "sausage_count"): - sausage_cut_done = env.sausage_count >= min_sausage_count - done = torch.logical_or(done, sausage_cut_done) - - # Optionally check if robots are at rest pose - if check_rest_pose and hasattr(env.scene, "articulations"): - left_arm = env.scene.articulations.get("left_arm") or env.scene.get("left_arm") - right_arm = env.scene.articulations.get("right_arm") or env.scene.get("right_arm") - - if left_arm is not None and right_arm is not None: - is_rest = torch.logical_and( - is_so101_at_rest_pose(left_arm.data.joint_pos, left_arm.data.joint_names), - is_so101_at_rest_pose(right_arm.data.joint_pos, right_arm.data.joint_names), - ) - done = torch.logical_and(done, is_rest) - - return done diff --git a/source/leisaac/leisaac/tasks/sausage_cut/sausage_cut_bi_arm_env_cfg.py b/source/leisaac/leisaac/tasks/sausage_cut/sausage_cut_bi_arm_env_cfg.py deleted file mode 100644 index aee31b41..00000000 --- a/source/leisaac/leisaac/tasks/sausage_cut/sausage_cut_bi_arm_env_cfg.py +++ /dev/null @@ -1,68 +0,0 @@ -import isaaclab.sim as sim_utils -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, -) -from leisaac.utils.general_assets import parse_usd_and_create_subassets - -from ..template import ( - BiArmObservationsCfg, - BiArmTaskEnvCfg, - BiArmTaskSceneCfg, - BiArmTerminationsCfg, -) - - -@configclass -class SausageCutBiArmSceneCfg(BiArmTaskSceneCfg): - """Scene configuration for the sausage cut task using two arms.""" - - # Use merged USD that includes sausage - scene: AssetBaseCfg = KITCHEN_WITH_SAUSAGE_CFG.replace(prim_path="{ENV_REGEX_NS}/Scene") - - # Replace right arm with knife-equipped robot (use standard config to match lehome) - right_arm: ArticulationCfg = SO101_KINFE_CFG.replace(prim_path="{ENV_REGEX_NS}/Right_Robot") - - -@configclass -class SausageCutBiArmEnvCfg(BiArmTaskEnvCfg): - """Configuration for the sausage cut environment.""" - - scene: SausageCutBiArmSceneCfg = SausageCutBiArmSceneCfg(env_spacing=4.0) - - observations: BiArmObservationsCfg = BiArmObservationsCfg() - - terminations: BiArmTerminationsCfg = BiArmTerminationsCfg() - - # Simulation configuration - CRITICAL: use_fabric=False is required for DeformableBody/cutting - render_cfg: sim_utils.RenderCfg = sim_utils.RenderCfg(rendering_mode="quality", antialiasing_mode="FXAA") - sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=1, render=render_cfg, use_fabric=False) - - def __post_init__(self) -> None: - super().__post_init__() - - # Camera positions - self.viewer.eye = (2.5, -5.0, 1.6) - self.viewer.lookat = (3.7, -6.15, 0.84) - - # Robot initial positions - 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 - self.decimation = 2 - self.dynamic_reset_gripper_effort_limit = False - - # Parse USD for additional assets (ChoppingBlock, etc.) - parse_usd_and_create_subassets( - KITCHEN_WITH_SAUSAGE_USD_PATH, - self, - exclude_name_list=["Sausage"], # Sausage is handled separately due to cutting - ) diff --git a/source/leisaac/leisaac/utils/domain_randomization.py b/source/leisaac/leisaac/utils/domain_randomization.py index f87be3ce..1903bc78 100644 --- a/source/leisaac/leisaac/utils/domain_randomization.py +++ b/source/leisaac/leisaac/utils/domain_randomization.py @@ -92,3 +92,23 @@ def randomize_mixed_objects_uniform( "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 2fdc63c5..cde004af 100644 --- a/source/leisaac/leisaac/utils/general_assets.py +++ b/source/leisaac/leisaac/utils/general_assets.py @@ -31,6 +31,14 @@ 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) @@ -173,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]