diff --git a/.gitignore b/.gitignore index b3a54b00..6f40f127 100644 --- a/.gitignore +++ b/.gitignore @@ -35,7 +35,9 @@ *.app # IDE stuff -*.vscode +.vscode/* +!.vscode/launch.json +!.vscode/settings.json *.*~ TAGS diff --git a/isaaclab_arena/assets/object_library.py b/isaaclab_arena/assets/object_library.py index 033b65d4..f11da40b 100644 --- a/isaaclab_arena/assets/object_library.py +++ b/isaaclab_arena/assets/object_library.py @@ -14,6 +14,11 @@ from isaaclab_arena.affordances.pressable import Pressable from isaaclab_arena.assets.object import Object from isaaclab_arena.assets.object_base import ObjectType +from isaaclab_arena.assets.object_utils import ( + ASSEMBLY_ARTICULATION_INIT_STATE, + RIGID_BODY_PROPS_HIGH_PRECISION, + RIGID_BODY_PROPS_STANDARD, +) from isaaclab_arena.assets.register import register_asset from isaaclab_arena.utils.pose import Pose @@ -335,3 +340,129 @@ class DexCube(LibraryObject): def __init__(self, prim_path: str | None = None, initial_pose: Pose | None = None): super().__init__(prim_path=prim_path, initial_pose=initial_pose) + + +@register_asset +class Peg(LibraryObject): + """ + A peg. + """ + + name = "peg" + tags = ["object"] + usd_path = f"{ISAACLAB_NUCLEUS_DIR}/Factory/factory_peg_8mm.usd" + object_type = ObjectType.ARTICULATION + scale = (3.0, 3.0, 3.0) + spawn_cfg_addon = { + "rigid_props": RIGID_BODY_PROPS_HIGH_PRECISION, + "mass_props": sim_utils.MassPropertiesCfg(mass=0.019), + "collision_props": sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), + } + asset_cfg_addon = { + "init_state": ASSEMBLY_ARTICULATION_INIT_STATE, + } + + +@register_asset +class Hole(LibraryObject): + """ + A hole. + """ + + name = "hole" + tags = ["object"] + usd_path = f"{ISAACLAB_NUCLEUS_DIR}/Factory/factory_hole_8mm.usd" + object_type = ObjectType.ARTICULATION + scale = (3.0, 3.0, 3.0) + spawn_cfg_addon = { + "rigid_props": RIGID_BODY_PROPS_HIGH_PRECISION, + "mass_props": sim_utils.MassPropertiesCfg(mass=0.05), + "collision_props": sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), + } + asset_cfg_addon = { + "init_state": ASSEMBLY_ARTICULATION_INIT_STATE, + } + + +@register_asset +class SmallGear(LibraryObject): + """ + A small gear. + """ + + name = "small_gear" + tags = ["object"] + usd_path = f"{ISAACLAB_NUCLEUS_DIR}/Factory/factory_gear_small.usd" + object_type = ObjectType.ARTICULATION + scale = (2.0, 2.0, 2.0) + spawn_cfg_addon = { + "rigid_props": RIGID_BODY_PROPS_STANDARD, + "mass_props": sim_utils.MassPropertiesCfg(mass=0.019), + "collision_props": sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), + } + asset_cfg_addon = { + "init_state": ASSEMBLY_ARTICULATION_INIT_STATE, + } + + +@register_asset +class LargeGear(LibraryObject): + """ + A large gear. + """ + + name = "large_gear" + tags = ["object"] + usd_path = f"{ISAACLAB_NUCLEUS_DIR}/Factory/factory_gear_large.usd" + object_type = ObjectType.ARTICULATION + scale = (2.0, 2.0, 2.0) + spawn_cfg_addon = { + "rigid_props": RIGID_BODY_PROPS_STANDARD, + "mass_props": sim_utils.MassPropertiesCfg(mass=0.019), + "collision_props": sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), + } + asset_cfg_addon = { + "init_state": ASSEMBLY_ARTICULATION_INIT_STATE, + } + + +@register_asset +class GearBase(LibraryObject): + """ + Gear base. + """ + + name = "gear_base" + tags = ["object"] + usd_path = f"{ISAACLAB_NUCLEUS_DIR}/Factory/factory_gear_base.usd" + object_type = ObjectType.ARTICULATION + scale = (2.0, 2.0, 2.0) + spawn_cfg_addon = { + "rigid_props": RIGID_BODY_PROPS_STANDARD, + "mass_props": sim_utils.MassPropertiesCfg(mass=0.05), + "collision_props": sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), + } + asset_cfg_addon = { + "init_state": ASSEMBLY_ARTICULATION_INIT_STATE, + } + + +@register_asset +class MediumGear(LibraryObject): + """ + A medium gear. + """ + + name = "medium_gear" + tags = ["object"] + usd_path = f"{ISAACLAB_NUCLEUS_DIR}/Factory/factory_gear_medium.usd" + object_type = ObjectType.ARTICULATION + scale = (2.0, 2.0, 2.0) + spawn_cfg_addon = { + "rigid_props": RIGID_BODY_PROPS_STANDARD, + "mass_props": sim_utils.MassPropertiesCfg(mass=0.019), + "collision_props": sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), + } + asset_cfg_addon = { + "init_state": ASSEMBLY_ARTICULATION_INIT_STATE, + } diff --git a/isaaclab_arena/assets/object_utils.py b/isaaclab_arena/assets/object_utils.py index 7715b78e..29ff9e90 100644 --- a/isaaclab_arena/assets/object_utils.py +++ b/isaaclab_arena/assets/object_utils.py @@ -3,6 +3,8 @@ # # SPDX-License-Identifier: Apache-2.0 +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationCfg from pxr import Usd from isaaclab_arena.assets.object_base import ObjectType @@ -59,3 +61,38 @@ def detect_object_type(usd_path: str | None = None, stage: Usd.Stage | None = No return ObjectType.ARTICULATION else: raise ValueError("This should not happen. There is an unknown USD type in the tree.") + + +# Predefined rigid body property configurations for factory assembly tasks +# High iteration count for precision tasks (peg/hole insertion) +RIGID_BODY_PROPS_HIGH_PRECISION = sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, + max_depenetration_velocity=5.0, + linear_damping=0.0, + angular_damping=0.0, + max_linear_velocity=1000.0, + max_angular_velocity=3666.0, + enable_gyroscopic_forces=True, + solver_position_iteration_count=192, + solver_velocity_iteration_count=1, + max_contact_impulse=1e32, +) + +# Standard iteration count for gear mesh tasks +RIGID_BODY_PROPS_STANDARD = sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, + max_depenetration_velocity=5.0, + linear_damping=0.0, + angular_damping=0.0, + max_linear_velocity=1000.0, + max_angular_velocity=3666.0, + enable_gyroscopic_forces=True, + solver_position_iteration_count=32, + solver_velocity_iteration_count=32, + max_contact_impulse=1e32, +) + +ASSEMBLY_ARTICULATION_INIT_STATE = ArticulationCfg.InitialStateCfg( + joint_pos={}, + joint_vel={}, +) diff --git a/isaaclab_arena/embodiments/franka/franka.py b/isaaclab_arena/embodiments/franka/franka.py index 7ebf8d46..f888fce5 100644 --- a/isaaclab_arena/embodiments/franka/franka.py +++ b/isaaclab_arena/embodiments/franka/franka.py @@ -23,6 +23,9 @@ from isaaclab.sensors.frame_transformer.frame_transformer_cfg import FrameTransformerCfg, OffsetCfg from isaaclab.sim.spawners.from_files.from_files_cfg import UsdFileCfg from isaaclab.utils import configclass + +# Create a custom Franka configuration for assembly tasks +# This is defined at module level to avoid being treated as a config field from isaaclab_assets.robots.franka import FRANKA_PANDA_HIGH_PD_CFG from isaaclab_tasks.manager_based.manipulation.stack.mdp import franka_stack_events from isaaclab_tasks.manager_based.manipulation.stack.mdp.observations import ee_frame_pos, ee_frame_quat @@ -34,6 +37,17 @@ from isaaclab_arena.embodiments.franka.observations import gripper_pos from isaaclab_arena.utils.pose import Pose +FRANKA_PANDA_ASSEMBLY_HIGH_PD_CFG = FRANKA_PANDA_HIGH_PD_CFG.copy() +FRANKA_PANDA_ASSEMBLY_HIGH_PD_CFG.spawn.activate_contact_sensors = True +FRANKA_PANDA_ASSEMBLY_HIGH_PD_CFG.spawn.rigid_props.disable_gravity = True +FRANKA_PANDA_ASSEMBLY_HIGH_PD_CFG.actuators["panda_shoulder"].stiffness = 150.0 +FRANKA_PANDA_ASSEMBLY_HIGH_PD_CFG.actuators["panda_shoulder"].damping = 30.0 +FRANKA_PANDA_ASSEMBLY_HIGH_PD_CFG.actuators["panda_forearm"].stiffness = 150.0 +FRANKA_PANDA_ASSEMBLY_HIGH_PD_CFG.actuators["panda_forearm"].damping = 30.0 +FRANKA_PANDA_ASSEMBLY_HIGH_PD_CFG.actuators["panda_hand"].stiffness = 150.0 +FRANKA_PANDA_ASSEMBLY_HIGH_PD_CFG.actuators["panda_hand"].damping = 30.0 +FRANKA_PANDA_ASSEMBLY_HIGH_PD_CFG.init_state.pos = (0.0, 0.0, 0.0) + @register_asset class FrankaEmbodiment(EmbodimentBase): diff --git a/isaaclab_arena/tasks/assembly_task.py b/isaaclab_arena/tasks/assembly_task.py new file mode 100644 index 00000000..a0a3e617 --- /dev/null +++ b/isaaclab_arena/tasks/assembly_task.py @@ -0,0 +1,191 @@ +# Copyright (c) 2025, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +from dataclasses import MISSING, dataclass + +import isaaclab.envs.mdp as mdp_isaac_lab +from isaaclab.envs.common import ViewerCfg +from isaaclab.envs.mimic_env_cfg import MimicEnvCfg +from isaaclab.managers import EventTermCfg, SceneEntityCfg, TerminationTermCfg +from isaaclab.utils import configclass + +import isaaclab_arena_environments.mdp as mdp +from isaaclab_arena.assets.asset import Asset +from isaaclab_arena.metrics.metric_base import MetricBase +from isaaclab_arena.metrics.object_moved import ObjectMovedRateMetric +from isaaclab_arena.metrics.success_rate import SuccessRateMetric +from isaaclab_arena.tasks.events import randomize_poses_and_align_auxiliary_assets +from isaaclab_arena.tasks.task_base import TaskBase +from isaaclab_arena.tasks.terminations import objects_in_proximity +from isaaclab_arena.utils.cameras import get_viewer_cfg_look_at_object + + +@dataclass +class TerminationConfig: + max_x_separation: float = 0.020 + max_y_separation: float = 0.020 + max_z_separation: float = 0.020 + + +class AssemblyTask(TaskBase): + """ + Assembly task where an object needs to be assembled with a base object, like peg insert, gear mesh, etc. + """ + + def __init__( + self, + fixed_asset: Asset, + held_asset: Asset, + auxiliary_asset_list: list[Asset], + background_scene: Asset, + episode_length_s: float | None = None, + termination_cfg: TerminationConfig = TerminationConfig( + max_x_separation=0.020, max_y_separation=0.020, max_z_separation=0.020 + ), + task_description: str | None = None, + pose_range: dict[str, tuple[float, float]] | None = None, + min_separation: float = 0.10, + randomization_mode: int = 0, + ): + super().__init__(episode_length_s=episode_length_s) + self.fixed_asset = fixed_asset + self.held_asset = held_asset + self.auxiliary_asset_list = auxiliary_asset_list + self.background_scene = background_scene + self.scene_config = None + self.events_cfg = EventsCfg( + pose_range=pose_range, + min_separation=min_separation, + asset_cfgs=[SceneEntityCfg(asset.name) for asset in [self.fixed_asset, self.held_asset]], + fixed_asset_cfg=SceneEntityCfg(self.fixed_asset.name), + auxiliary_asset_cfgs=[SceneEntityCfg(asset.name) for asset in self.auxiliary_asset_list], + randomization_mode=randomization_mode, + ) + self.termination_cfg = self._make_termination_cfg(termination_cfg) + self.task_description = ( + f"Assemble the {self.held_asset.name} with the {self.fixed_asset.name}" + if task_description is None + else task_description + ) + + def get_scene_cfg(self): + """Get scene configuration.""" + return self.scene_config + + def get_termination_cfg(self): + return self.termination_cfg + + def _make_termination_cfg(self, term_cfg: TerminationConfig): + success = TerminationTermCfg( + func=objects_in_proximity, + params={ + "object_cfg": SceneEntityCfg(self.held_asset.name), + "target_object_cfg": SceneEntityCfg(self.fixed_asset.name), + "max_x_separation": term_cfg.max_x_separation, # Tolerance for assembly alignment + "max_y_separation": term_cfg.max_y_separation, + "max_z_separation": term_cfg.max_z_separation, + }, + ) + object_dropped = TerminationTermCfg( + func=mdp_isaac_lab.root_height_below_minimum, + params={ + "minimum_height": self.background_scene.object_min_z, + "asset_cfg": SceneEntityCfg(self.held_asset.name), + }, + ) + return TerminationsCfg( + success=success, + object_dropped=object_dropped, + ) + + def get_events_cfg(self): + """Get events configuration for assembly task.""" + return self.events_cfg + + def get_prompt(self): + raise NotImplementedError("Function not implemented yet.") + + def get_mimic_env_cfg(self, embodiment_name: str): + return FactoryAssemblyMimicEnvCfg() + + def get_metrics(self) -> list[MetricBase]: + return [ + SuccessRateMetric(), + ObjectMovedRateMetric(self.held_asset), + ] + + def get_viewer_cfg(self) -> ViewerCfg: + """Get viewer configuration to look at the held asset. + + Camera is positioned at right-back-top of the object for better view of assembly operations. + """ + return get_viewer_cfg_look_at_object( + lookat_object=self.held_asset, + offset=np.array([1.5, -0.5, 1.0]), # Rotated 180° around z-axis from original view + ) + + +@configclass +class TerminationsCfg: + """Termination terms for the MDP.""" + + time_out: TerminationTermCfg = TerminationTermCfg(func=mdp_isaac_lab.time_out) + + success: TerminationTermCfg = MISSING + + object_dropped: TerminationTermCfg = MISSING + + +@configclass +class EventsCfg: + """ + Configuration for assembly task events. + """ + + reset_all: EventTermCfg = MISSING + randomize_asset_positions: EventTermCfg = MISSING + + def __init__( + self, + pose_range: dict[str, tuple[float, float]], + min_separation: float, + asset_cfgs: list[SceneEntityCfg], + fixed_asset_cfg: SceneEntityCfg, + auxiliary_asset_cfgs: list[SceneEntityCfg], + randomization_mode: int = 0, + ): + self.reset_all = EventTermCfg( + func=mdp.reset_scene_to_default, mode="reset", params={"reset_joint_targets": True} + ) + + self.randomize_asset_positions = EventTermCfg( + func=randomize_poses_and_align_auxiliary_assets, + mode="reset", + params={ + "pose_range": pose_range, + "min_separation": min_separation, + "asset_cfgs": asset_cfgs, + "fixed_asset_cfg": fixed_asset_cfg, + "auxiliary_asset_cfgs": auxiliary_asset_cfgs, + "randomization_mode": randomization_mode, + }, + ) + + +class FactoryAssemblyMimicEnvCfg(MimicEnvCfg): + """ + Isaac Lab Mimic environment config class for assembly task. + + Note: + This is a base configuration class. Specific assembly tasks + (e.g., PegInsert, GearMesh) should create their own subclasses with + appropriate asset names. + """ + + embodiment_name: str = MISSING + fixed_asset_name: str = MISSING + held_asset_name: str = MISSING + assist_asset_list_names: list[str] = MISSING diff --git a/isaaclab_arena/tasks/events.py b/isaaclab_arena/tasks/events.py new file mode 100644 index 00000000..52aa7b1b --- /dev/null +++ b/isaaclab_arena/tasks/events.py @@ -0,0 +1,78 @@ +# Copyright (c) 2025, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +import isaaclab.utils.math as math_utils +from isaaclab.managers import SceneEntityCfg +from isaaclab_tasks.manager_based.manipulation.stack.mdp.franka_stack_events import sample_object_poses + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + +def randomize_poses_and_align_auxiliary_assets( + env: ManagerBasedEnv, + env_ids: torch.Tensor, + asset_cfgs: list[SceneEntityCfg], + min_separation: float = 0.0, + pose_range: dict[str, tuple[float, float]] = {}, + max_sample_tries: int = 5000, + fixed_asset_cfg: SceneEntityCfg | None = None, + auxiliary_asset_cfgs: list[SceneEntityCfg] | None = None, + randomization_mode: int = 0, +): + """ + Randomize object poses and update the poses of related assets accordingly. + + Args: + randomization_mode (int): + 0: Randomize the poses of the fixed and held assets. + 1: Randomize the poses of the fixed, held, and auxiliary assets. The poses + of the auxiliary assets are set relative to the pose of the fixed asset. + """ + if env_ids is None: + return + + # Randomize poses in each environment independently + for cur_env in env_ids.tolist(): + pose_list = sample_object_poses( + num_objects=len(asset_cfgs), + min_separation=min_separation, + pose_range=pose_range, + max_sample_tries=max_sample_tries, + ) + + # Randomize pose for each object + for i in range(len(asset_cfgs)): + asset_cfg = asset_cfgs[i] + asset = env.scene[asset_cfg.name] + + # Write pose to simulation + pose_tensor = torch.tensor([pose_list[i]], device=env.device) + positions = pose_tensor[:, 0:3] + env.scene.env_origins[cur_env, 0:3] + orientations = math_utils.quat_from_euler_xyz(pose_tensor[:, 3], pose_tensor[:, 4], pose_tensor[:, 5]) + + asset.write_root_pose_to_sim( + torch.cat([positions, orientations], dim=-1), env_ids=torch.tensor([cur_env], device=env.device) + ) + asset.write_root_velocity_to_sim( + torch.zeros(1, 6, device=env.device), env_ids=torch.tensor([cur_env], device=env.device) + ) + + if randomization_mode == 1 and auxiliary_asset_cfgs is not None and asset_cfg.name == fixed_asset_cfg.name: + # set the poses of the auxiliary assets according to the pose of the fixed asset + for j in range(len(auxiliary_asset_cfgs)): + rel_asset_cfg = auxiliary_asset_cfgs[j] + rel_asset = env.scene[rel_asset_cfg.name] + rel_asset.write_root_pose_to_sim( + torch.cat([positions, orientations], dim=-1), env_ids=torch.tensor([cur_env], device=env.device) + ) + rel_asset.write_root_velocity_to_sim( + torch.zeros(1, 6, device=env.device), env_ids=torch.tensor([cur_env], device=env.device) + ) diff --git a/isaaclab_arena/tests/test_assembly_task.py b/isaaclab_arena/tests/test_assembly_task.py new file mode 100644 index 00000000..c6a61f8d --- /dev/null +++ b/isaaclab_arena/tests/test_assembly_task.py @@ -0,0 +1,428 @@ +# Copyright (c) 2025, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import gymnasium as gym +import torch + +from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function + +NUM_STEPS = 10 +HEADLESS = True + + +def get_peg_insert_test_environment(num_envs: int, remove_events: bool = False): + """Returns a peg insert environment for testing.""" + import isaaclab.sim as sim_utils + + from isaaclab_arena.assets.asset_registry import AssetRegistry + from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser + from isaaclab_arena.embodiments.franka.franka import FRANKA_PANDA_ASSEMBLY_HIGH_PD_CFG, FrankaEmbodiment + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder + from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment + from isaaclab_arena.scene.scene import Scene + from isaaclab_arena.tasks.assembly_task import AssemblyTask + from isaaclab_arena.utils.pose import Pose + from isaaclab_arena_environments import mdp + + args_parser = get_isaaclab_arena_cli_parser() + args_cli = args_parser.parse_args(["--num_envs", str(num_envs)]) + args_cli.enable_pinocchio = False + + asset_registry = AssetRegistry() + + # Create scene assets + background = asset_registry.get_asset_by_name("table")() + background.set_initial_pose(Pose(position_xyz=(0.55, 0.0, 0.0), rotation_wxyz=(0.707, 0, 0, 0.707))) + + peg = asset_registry.get_asset_by_name("peg")() + peg.set_initial_pose(Pose(position_xyz=(0.45, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + + hole = asset_registry.get_asset_by_name("hole")() + hole.set_initial_pose(Pose(position_xyz=(0.45, 0.1, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + + light_spawner_cfg = sim_utils.DomeLightCfg(color=(0.75, 0.75, 0.75), intensity=1500.0) + light = asset_registry.get_asset_by_name("light")(spawner_cfg=light_spawner_cfg) + + # Create embodiment + embodiment = FrankaEmbodiment() + embodiment.scene_config.robot = FRANKA_PANDA_ASSEMBLY_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + scene = Scene(assets=[background, peg, hole, light]) + + # Create assembly task + task = AssemblyTask( + task_description="Assemble the peg with the hole", + fixed_asset=peg, + held_asset=hole, + auxiliary_asset_list=[], + background_scene=background, + pose_range={"x": (0.2, 0.6), "y": (-0.20, 0.20), "z": (0.0, 0.0), "yaw": (-1.0, 1.0)}, + min_separation=0.1, + ) + + isaaclab_arena_environment = IsaacLabArenaEnvironment( + name="test_peg_insert", + embodiment=embodiment, + scene=scene, + task=task, + env_cfg_callback=mdp.assembly_env_cfg_callback, + ) + + env_builder = ArenaEnvBuilder(isaaclab_arena_environment, args_cli) + name, cfg = env_builder.build_registered() + + if remove_events: + cfg.events.reset_all = None + cfg.events.randomize_asset_positions = None + + env = gym.make(name, cfg=cfg).unwrapped + env.reset() + + return env, peg, hole + + +def get_gear_mesh_test_environment(num_envs: int, remove_events: bool = False): + """Returns a gear mesh environment for testing.""" + import isaaclab.sim as sim_utils + + from isaaclab_arena.assets.asset_registry import AssetRegistry + from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser + from isaaclab_arena.embodiments.franka.franka import FRANKA_PANDA_ASSEMBLY_HIGH_PD_CFG, FrankaEmbodiment + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder + from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment + from isaaclab_arena.scene.scene import Scene + from isaaclab_arena.tasks.assembly_task import AssemblyTask + from isaaclab_arena.utils.pose import Pose + from isaaclab_arena_environments import mdp + + args_parser = get_isaaclab_arena_cli_parser() + args_cli = args_parser.parse_args(["--num_envs", str(num_envs)]) + args_cli.enable_pinocchio = False + + asset_registry = AssetRegistry() + + # Create scene assets + background = asset_registry.get_asset_by_name("table")() + background.set_initial_pose(Pose(position_xyz=(0.55, 0.0, 0.0), rotation_wxyz=(0.707, 0, 0, 0.707))) + + gear_base = asset_registry.get_asset_by_name("gear_base")() + gear_base.set_initial_pose(Pose(position_xyz=(0.6, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + + medium_gear = asset_registry.get_asset_by_name("medium_gear")() + medium_gear.set_initial_pose(Pose(position_xyz=(0.5, 0.2, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + + small_gear = asset_registry.get_asset_by_name("small_gear")() + small_gear.set_initial_pose(Pose(position_xyz=(0.6, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + + large_gear = asset_registry.get_asset_by_name("large_gear")() + large_gear.set_initial_pose(Pose(position_xyz=(0.6, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + + light_spawner_cfg = sim_utils.DomeLightCfg(color=(0.75, 0.75, 0.75), intensity=1500.0) + light = asset_registry.get_asset_by_name("light")(spawner_cfg=light_spawner_cfg) + + # Create embodiment + embodiment = FrankaEmbodiment() + embodiment.scene_config.robot = FRANKA_PANDA_ASSEMBLY_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + scene = Scene(assets=[background, gear_base, medium_gear, small_gear, large_gear, light]) + + # Create gear mesh task + task = AssemblyTask( + task_description="Mesh the medium gear with the gear base", + fixed_asset=gear_base, + held_asset=medium_gear, + auxiliary_asset_list=[small_gear, large_gear], + background_scene=background, + pose_range={"x": (0.25, 0.6), "y": (-0.20, 0.20), "z": (0.0, 0.0), "yaw": (-1.0, 1.0)}, + min_separation=0.18, + randomization_mode=1, + ) + + isaaclab_arena_environment = IsaacLabArenaEnvironment( + name="test_gear_mesh", + embodiment=embodiment, + scene=scene, + task=task, + env_cfg_callback=mdp.assembly_env_cfg_callback, + ) + + env_builder = ArenaEnvBuilder(isaaclab_arena_environment, args_cli) + name, cfg = env_builder.build_registered() + + if remove_events: + cfg.events.reset_all = None + cfg.events.randomize_asset_positions = None + + env = gym.make(name, cfg=cfg).unwrapped + env.reset() + + return env, gear_base, medium_gear, small_gear, large_gear + + +def _test_peg_insert_assembly_single(simulation_app) -> bool: + """Test peg insert assembly with single environment.""" + from isaaclab.envs.manager_based_env import ManagerBasedEnv + + from isaaclab_arena.tests.utils.simulation import step_zeros_and_call + + env, peg, hole = get_peg_insert_test_environment(num_envs=1, remove_events=True) + + def assert_assembled(env: ManagerBasedEnv, terminated: torch.Tensor): + # Check if objects are close together (assembled) + peg_pos = peg.get_object_pose(env, is_relative=True)[:, :3] + hole_pos = hole.get_object_pose(env, is_relative=True)[:, :3] + distance = torch.norm(peg_pos - hole_pos, dim=-1) + + print(f"Distance between peg and hole: {distance.item():.4f}m") + assert distance.item() < 0.05, f"Objects not assembled, distance: {distance.item():.4f}m" + + # Check terminated + assert terminated.shape == torch.Size([1]), "Terminated shape is not correct" + assert terminated.item(), "Task didn't terminate when it should have" + print("✓ Peg insert assembly task completed successfully") + + try: + print("Testing peg insert assembly (single env)...") + + # Manually place hole on peg to simulate successful assembly - use absolute world coordinates + peg_pose = peg.get_object_pose(env, is_relative=False) + env.scene[hole.name].write_root_pose_to_sim(peg_pose, env_ids=torch.tensor([0], device=env.device)) + + step_zeros_and_call(env, NUM_STEPS, assert_assembled) + + except Exception as e: + print(f"Error: {e}") + import traceback + + traceback.print_exc() + return False + finally: + env.close() + + return True + + +def _test_gear_mesh_assembly_single(simulation_app) -> bool: + """Test gear mesh assembly with single environment.""" + from isaaclab.envs.manager_based_env import ManagerBasedEnv + + from isaaclab_arena.tests.utils.simulation import step_zeros_and_call + + env, gear_base, medium_gear, small_gear, large_gear = get_gear_mesh_test_environment(num_envs=1, remove_events=True) + + def assert_assembled(env: ManagerBasedEnv, terminated: torch.Tensor): + # Check if gears are close together (meshed) + base_pos = gear_base.get_object_pose(env, is_relative=True)[:, :3] + medium_pos = medium_gear.get_object_pose(env, is_relative=True)[:, :3] + distance = torch.norm(base_pos - medium_pos, dim=-1) + + print(f"Distance between gear base and medium gear: {distance.item():.4f}m") + assert distance.item() < 0.05, f"Gears not meshed, distance: {distance.item():.4f}m" + + # Check terminated + assert terminated.shape == torch.Size([1]), "Terminated shape is not correct" + assert terminated.item(), "Task didn't terminate when it should have" + print("✓ Gear mesh assembly task completed successfully") + + try: + print("Testing gear mesh assembly (single env)...") + + # Manually place medium gear on gear base to simulate successful assembly - use absolute world coordinates + base_pose = gear_base.get_object_pose(env, is_relative=False) + env.scene[medium_gear.name].write_root_pose_to_sim(base_pose, env_ids=torch.tensor([0], device=env.device)) + + step_zeros_and_call(env, NUM_STEPS, assert_assembled) + + except Exception as e: + print(f"Error: {e}") + import traceback + + traceback.print_exc() + return False + finally: + env.close() + + return True + + +def _test_peg_insert_assembly_multi(simulation_app) -> bool: + """Test peg insert assembly with multiple environments.""" + from isaaclab_arena.tests.utils.simulation import step_zeros_and_call + + env, peg, hole = get_peg_insert_test_environment(num_envs=2, remove_events=True) + + try: + with torch.inference_mode(): + print("Testing peg insert assembly (multi env)...") + + # Assemble in both environments - use absolute world coordinates + peg_poses = peg.get_object_pose(env, is_relative=False) + env.scene[hole.name].write_root_pose_to_sim(peg_poses, env_ids=None) + + step_zeros_and_call(env, NUM_STEPS) + + # Check distances + peg_pos = peg.get_object_pose(env, is_relative=True)[:, :3] + hole_pos = hole.get_object_pose(env, is_relative=True)[:, :3] + distances = torch.norm(peg_pos - hole_pos, dim=-1) + + print(f"Distances in both envs: {distances}") + assert torch.all(distances < 0.05), f"Not all environments assembled: {distances}" + print("✓ Multi-environment peg insert assembly successful") + + except Exception as e: + print(f"Error: {e}") + import traceback + + traceback.print_exc() + return False + finally: + env.close() + + return True + + +def _test_gear_mesh_assembly_multi(simulation_app) -> bool: + """Test gear mesh assembly with multiple environments.""" + from isaaclab_arena.tests.utils.simulation import step_zeros_and_call + + env, gear_base, medium_gear, small_gear, large_gear = get_gear_mesh_test_environment(num_envs=2, remove_events=True) + + try: + with torch.inference_mode(): + print("Testing gear mesh assembly (multi env)...") + + # Assemble in both environments - use absolute world coordinates + base_poses = gear_base.get_object_pose(env, is_relative=False) + env.scene[medium_gear.name].write_root_pose_to_sim(base_poses, env_ids=None) + + step_zeros_and_call(env, NUM_STEPS) + + # Check distances + base_pos = gear_base.get_object_pose(env, is_relative=True)[:, :3] + medium_pos = medium_gear.get_object_pose(env, is_relative=True)[:, :3] + distances = torch.norm(base_pos - medium_pos, dim=-1) + + print(f"Distances in both envs: {distances}") + assert torch.all(distances < 0.05), f"Not all environments assembled: {distances}" + print("✓ Multi-environment gear mesh assembly successful") + + except Exception as e: + print(f"Error: {e}") + import traceback + + traceback.print_exc() + return False + finally: + env.close() + + return True + + +def _test_peg_insert_initialization(simulation_app) -> bool: + """Test that peg insert task initializes correctly with proper asset configurations.""" + from isaaclab_arena.tests.utils.simulation import step_zeros_and_call + + try: + print("Testing peg insert task initialization...") + + # Test peg insert initialization + env, peg, hole = get_peg_insert_test_environment(num_envs=1, remove_events=False) + + # Check that objects exist in scene + assert peg.name in env.scene.keys(), f"Peg '{peg.name}' not found in scene" + assert hole.name in env.scene.keys(), f"Hole '{hole.name}' not found in scene" + + # Run a few steps to ensure stability + step_zeros_and_call(env, NUM_STEPS) + + print("✓ Peg insert task initialized successfully") + env.close() + + except Exception as e: + print(f"Error: {e}") + import traceback + + traceback.print_exc() + return False + + return True + + +def _test_gear_mesh_initialization(simulation_app) -> bool: + """Test that gear mesh task initializes correctly with proper asset configurations.""" + from isaaclab_arena.tests.utils.simulation import step_zeros_and_call + + try: + print("Testing gear mesh task initialization...") + + # Test gear mesh initialization + env, gear_base, medium_gear, small_gear, large_gear = get_gear_mesh_test_environment( + num_envs=1, remove_events=False + ) + + # Check that all gears exist in scene + assert gear_base.name in env.scene.keys(), f"Gear base '{gear_base.name}' not found in scene" + assert medium_gear.name in env.scene.keys(), f"Medium gear '{medium_gear.name}' not found in scene" + assert small_gear.name in env.scene.keys(), f"Small gear '{small_gear.name}' not found in scene" + assert large_gear.name in env.scene.keys(), f"Large gear '{large_gear.name}' not found in scene" + + # Run a few steps to ensure stability + step_zeros_and_call(env, NUM_STEPS) + + print("✓ Gear mesh task initialized successfully") + env.close() + + except Exception as e: + print(f"Error: {e}") + import traceback + + traceback.print_exc() + return False + + return True + + +# Test functions that will be called by pytest +def test_peg_insert_assembly_single(): + result = run_simulation_app_function(_test_peg_insert_assembly_single, headless=HEADLESS, enable_pinocchio=False) + assert result, f"Test {_test_peg_insert_assembly_single.__name__} failed" + + +def test_gear_mesh_assembly_single(): + result = run_simulation_app_function(_test_gear_mesh_assembly_single, headless=HEADLESS) + assert result, f"Test {_test_gear_mesh_assembly_single.__name__} failed" + + +def test_peg_insert_assembly_multi(): + result = run_simulation_app_function(_test_peg_insert_assembly_multi, headless=HEADLESS, enable_pinocchio=False) + assert result, f"Test {_test_peg_insert_assembly_multi.__name__} failed" + + +def test_gear_mesh_assembly_multi(): + result = run_simulation_app_function(_test_gear_mesh_assembly_multi, headless=HEADLESS) + assert result, f"Test {_test_gear_mesh_assembly_multi.__name__} failed" + + +def test_peg_insert_initialization(): + """ + For peg insert task, we need to test the task with pinocchio disabled due to the "peg" and "hole" assets are not compatible with pinocchio. + """ + result = run_simulation_app_function(_test_peg_insert_initialization, headless=HEADLESS, enable_pinocchio=False) + assert result, f"Test {_test_peg_insert_initialization.__name__} failed" + + +def test_gear_mesh_initialization(): + result = run_simulation_app_function(_test_gear_mesh_initialization, headless=HEADLESS) + assert result, f"Test {_test_gear_mesh_initialization.__name__} failed" + + +if __name__ == "__main__": + test_peg_insert_initialization() + test_gear_mesh_initialization() + test_peg_insert_assembly_single() + test_gear_mesh_assembly_single() + test_peg_insert_assembly_multi() + test_gear_mesh_assembly_multi() diff --git a/isaaclab_arena/tests/utils/subprocess.py b/isaaclab_arena/tests/utils/subprocess.py index 081e3ed2..259518aa 100644 --- a/isaaclab_arena/tests/utils/subprocess.py +++ b/isaaclab_arena/tests/utils/subprocess.py @@ -67,7 +67,9 @@ def _close_persistent(): _PERSISTENT_SIM_APP_LAUNCHER.app.close() -def get_persistent_simulation_app(headless: bool, enable_cameras: bool = False) -> SimulationApp: +def get_persistent_simulation_app( + headless: bool, enable_cameras: bool = False, enable_pinocchio: bool = True +) -> SimulationApp: """Create once, reuse forever (until process exit).""" global _PERSISTENT_SIM_APP_LAUNCHER, _PERSISTENT_INIT_ARGS # Create a new simulation app if it doesn't exist @@ -76,6 +78,7 @@ def get_persistent_simulation_app(headless: bool, enable_cameras: bool = False) simulation_app_args = parser.parse_args([]) simulation_app_args.headless = headless simulation_app_args.enable_cameras = enable_cameras + simulation_app_args.enable_pinocchio = enable_pinocchio with _IsolatedArgv([]): app_launcher = get_app_launcher(simulation_app_args) @@ -96,7 +99,11 @@ def get_persistent_simulation_app(headless: bool, enable_cameras: bool = False) def run_simulation_app_function( - function: Callable[..., bool], headless: bool = True, enable_cameras: bool = False, **kwargs + function: Callable[..., bool], + headless: bool = True, + enable_cameras: bool = False, + enable_pinocchio: bool = True, + **kwargs, ) -> bool: """Run a simulation app in a separate process. @@ -115,7 +122,9 @@ def run_simulation_app_function( # Get a persistent simulation app global _AT_LEAST_ONE_TEST_FAILED try: - simulation_app = get_persistent_simulation_app(headless=headless, enable_cameras=enable_cameras) + simulation_app = get_persistent_simulation_app( + headless=headless, enable_cameras=enable_cameras, enable_pinocchio=enable_pinocchio + ) test_result = bool(function(simulation_app, **kwargs)) if not test_result: _AT_LEAST_ONE_TEST_FAILED = True diff --git a/isaaclab_arena_environments/mdp/__init__.py b/isaaclab_arena_environments/mdp/__init__.py new file mode 100644 index 00000000..d9036161 --- /dev/null +++ b/isaaclab_arena_environments/mdp/__init__.py @@ -0,0 +1,17 @@ +# Copyright (c) 2025, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""This sub-module contains the functions that are specific to the environment.""" + +from isaaclab.envs.mdp import * # noqa: F401, F403 +from isaaclab_tasks.manager_based.manipulation.place.mdp import * # noqa: F401, F403 +from isaaclab_tasks.manager_based.manipulation.stack.mdp import * # noqa: F401, F403 + +from .env_callbacks import * # noqa: F401, F403 diff --git a/isaaclab_arena_environments/mdp/env_callbacks.py b/isaaclab_arena_environments/mdp/env_callbacks.py new file mode 100644 index 00000000..656ccef4 --- /dev/null +++ b/isaaclab_arena_environments/mdp/env_callbacks.py @@ -0,0 +1,71 @@ +# Copyright (c) 2025, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Environment configuration callbacks for manipulation tasks.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from isaaclab_arena.environments.isaaclab_arena_manager_based_env import IsaacLabArenaManagerBasedRLEnvCfg + + +def assembly_env_cfg_callback(env_cfg: IsaacLabArenaManagerBasedRLEnvCfg) -> IsaacLabArenaManagerBasedRLEnvCfg: + """ + Environment configuration callback optimized for assembly tasks. + + This callback modifies the simulation settings to provide better stability + and precision for tasks like peg insertion, gear meshing, and other fine + manipulation operations. + + Args: + env_cfg: The environment configuration to modify. + + Returns: + The modified environment configuration. + """ + from isaaclab.sim import PhysxCfg, SimulationCfg + from isaaclab.sim.spawners.materials import RigidBodyMaterialCfg + + # Simulation settings optimized for assembly tasks + env_cfg.sim = SimulationCfg( + dt=1 / 60, # 60Hz - balance between speed and stability + render_interval=2, + physx=PhysxCfg( + solver_type=1, + max_position_iteration_count=192, # Important to avoid interpenetration + max_velocity_iteration_count=1, + bounce_threshold_velocity=0.2, + friction_offset_threshold=0.01, + friction_correlation_distance=0.00625, + gpu_max_rigid_contact_count=2**23, + gpu_max_rigid_patch_count=2**23, + gpu_max_num_partitions=1, # Important for stable simulation + ), + physics_material=RigidBodyMaterialCfg( + static_friction=1.0, + dynamic_friction=1.0, + ), + ) + + # Control frequency = 60Hz / 2 = 30Hz + env_cfg.decimation = 2 + + return env_cfg diff --git a/isaaclab_arena_environments/tabletop_gearmesh_environment.py b/isaaclab_arena_environments/tabletop_gearmesh_environment.py new file mode 100644 index 00000000..539a94a4 --- /dev/null +++ b/isaaclab_arena_environments/tabletop_gearmesh_environment.py @@ -0,0 +1,126 @@ +# Copyright (c) 2025, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse + +from isaaclab_arena_environments.example_environment_base import ExampleEnvironmentBase + + +class GearMeshEnvironment(ExampleEnvironmentBase): + """ + Gear mesh assembly environment with 4 gears: + - gear_base: gear base (to be meshed with) + - medium_gear: medium gear (to be picked and assembled) + - small_gear: small reference gear + - large_gear: large reference gear + """ + + name: str = "gear_mesh" + + def get_env(self, args_cli: argparse.Namespace): # -> IsaacLabArenaEnvironment: + import isaaclab.sim as sim_utils + + from isaaclab_arena.assets.background_library import Table + from isaaclab_arena.assets.object_library import GearBase, LargeGear, Light, MediumGear, SmallGear + from isaaclab_arena.embodiments.franka.franka import FRANKA_PANDA_ASSEMBLY_HIGH_PD_CFG, FrankaEmbodiment + from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment + from isaaclab_arena.scene.scene import Scene + from isaaclab_arena.tasks.assembly_task import AssemblyTask + from isaaclab_arena.utils.pose import Pose + from isaaclab_arena_environments import mdp + + # Get assets from registry + background = self.asset_registry.get_asset_by_name(args_cli.background)() + gear_base = self.asset_registry.get_asset_by_name("gear_base")() + medium_gear = self.asset_registry.get_asset_by_name("medium_gear")() + small_gear = self.asset_registry.get_asset_by_name("small_gear")() + large_gear = self.asset_registry.get_asset_by_name("large_gear")() + light_spawner_cfg = sim_utils.DomeLightCfg(color=(0.75, 0.75, 0.75), intensity=1500.0) + light = self.asset_registry.get_asset_by_name("light")(spawner_cfg=light_spawner_cfg) + embodiment = self.asset_registry.get_asset_by_name(args_cli.embodiment)(enable_cameras=args_cli.enable_cameras) + embodiment.scene_config.robot = FRANKA_PANDA_ASSEMBLY_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + if args_cli.teleop_device is not None: + teleop_device = self.device_registry.get_device_by_name(args_cli.teleop_device)() + else: + teleop_device = None + + background.set_initial_pose(Pose(position_xyz=(0.55, 0.0, 0.0), rotation_wxyz=(0.707, 0, 0, 0.707))) + + # Set initial poses for all 4 gears + gear_base.set_initial_pose( + Pose( + position_xyz=(0.6, 0.0, 0.0), # Gear base position + rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + ) + ) + + medium_gear.set_initial_pose( + Pose( + position_xyz=(0.5, 0.2, 0.0), # Medium gear to be assembled + rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + ) + ) + + small_gear.set_initial_pose( + Pose( + position_xyz=(0.6, 0.0, 0.0), # Small reference gear + rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + ) + ) + + large_gear.set_initial_pose( + Pose( + position_xyz=(0.6, 0.0, 0.0), # Large reference gear + rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + ) + ) + + # Create scene with all 4 gears and background + scene = Scene(assets=[background, gear_base, medium_gear, small_gear, large_gear, light]) + + # Create gear mesh task + task = AssemblyTask( + fixed_asset=gear_base, + held_asset=medium_gear, + auxiliary_asset_list=[small_gear, large_gear], + background_scene=background, + pose_range={"x": (0.25, 0.6), "y": (-0.20, 0.20), "z": (0.0, 0.0), "yaw": (-1.0, 1.0)}, + min_separation=0.18, + randomization_mode=1, + ) + + isaaclab_arena_environment = IsaacLabArenaEnvironment( + name=self.name, + embodiment=embodiment, + scene=scene, + task=task, + teleop_device=teleop_device, + env_cfg_callback=mdp.assembly_env_cfg_callback, + ) + return isaaclab_arena_environment + + @staticmethod + def add_cli_args(parser: argparse.ArgumentParser) -> None: + """Add CLI arguments for gear mesh environment.""" + parser.add_argument("--background", type=str, default="table", help="Background scene (table)") + parser.add_argument("--embodiment", type=str, default="franka", help="Robot embodiment") + parser.add_argument( + "--teleop_device", type=str, default=None, help="Teleoperation device (e.g., keyboard, spacemouse)" + ) diff --git a/isaaclab_arena_environments/tabletop_peginsert_environment.py b/isaaclab_arena_environments/tabletop_peginsert_environment.py new file mode 100644 index 00000000..35d5c1f1 --- /dev/null +++ b/isaaclab_arena_environments/tabletop_peginsert_environment.py @@ -0,0 +1,103 @@ +# Copyright (c) 2025, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse + +from isaaclab_arena_environments.example_environment_base import ExampleEnvironmentBase + + +class PegInsertEnvironment(ExampleEnvironmentBase): + + name: str = "peg_insert" + + def get_env(self, args_cli: argparse.Namespace): # -> IsaacLabArenaEnvironment: + import isaaclab.sim as sim_utils + + from isaaclab_arena.assets.background_library import Table + from isaaclab_arena.assets.object_library import Hole, Light, Peg + from isaaclab_arena.embodiments.franka.franka import FRANKA_PANDA_ASSEMBLY_HIGH_PD_CFG, FrankaEmbodiment + from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment + from isaaclab_arena.scene.scene import Scene + from isaaclab_arena.tasks.assembly_task import AssemblyTask + from isaaclab_arena.utils.pose import Pose + from isaaclab_arena_environments import mdp + + background = self.asset_registry.get_asset_by_name(args_cli.background)() + pick_up_object = self.asset_registry.get_asset_by_name(args_cli.object)() + destination_object = self.asset_registry.get_asset_by_name(args_cli.destination_object)() + light_spawner_cfg = sim_utils.DomeLightCfg(color=(0.75, 0.75, 0.75), intensity=1500.0) + light = self.asset_registry.get_asset_by_name("light")(spawner_cfg=light_spawner_cfg) + embodiment = self.asset_registry.get_asset_by_name(args_cli.embodiment)(enable_cameras=args_cli.enable_cameras) + embodiment.scene_config.robot = FRANKA_PANDA_ASSEMBLY_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + if args_cli.teleop_device is not None: + teleop_device = self.device_registry.get_device_by_name(args_cli.teleop_device)() + else: + teleop_device = None + + background.set_initial_pose(Pose(position_xyz=(0.55, 0.0, 0.0), rotation_wxyz=(0.707, 0, 0, 0.707))) + + pick_up_object.set_initial_pose( + Pose( + position_xyz=(0.45, 0.0, 0.0), + rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + ) + ) + + destination_object.set_initial_pose( + Pose( + position_xyz=(0.45, 0.1, 0.0), + rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + ) + ) + + scene = Scene(assets=[background, pick_up_object, destination_object, light]) + + task = AssemblyTask( + task_description="Assemble the peg with the hole", + fixed_asset=pick_up_object, + held_asset=destination_object, + auxiliary_asset_list=[], + background_scene=background, + pose_range={ + "x": (0.2, 0.6), + "y": (-0.20, 0.20), + "z": (0.0, 0.0), + "yaw": (-1.0, 1.0), + }, + min_separation=0.1, + ) + + isaaclab_arena_environment = IsaacLabArenaEnvironment( + name=self.name, + embodiment=embodiment, + scene=scene, + task=task, + teleop_device=teleop_device, + env_cfg_callback=mdp.assembly_env_cfg_callback, + ) + return isaaclab_arena_environment + + @staticmethod + def add_cli_args(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--object", type=str, default="peg") + parser.add_argument("--destination_object", type=str, default="hole") + parser.add_argument("--background", type=str, default="table") + parser.add_argument("--embodiment", type=str, default="franka") + parser.add_argument("--teleop_device", type=str, default=None)