diff --git a/isaaclab_arena/assets/background_library.py b/isaaclab_arena/assets/background_library.py index 4e5639bcfb..2d8baf697f 100644 --- a/isaaclab_arena/assets/background_library.py +++ b/isaaclab_arena/assets/background_library.py @@ -88,7 +88,10 @@ class PackingTableBackground(LibraryBackground): name = "packing_table" tags = ["background"] usd_path = f"{ARENA_NUCLEUS_DIR}/Arena/assets/background_library/packing_table/packing_table.usd" - initial_pose = Pose(position_xyz=(0.72193, -0.04727, -0.92512), rotation_xyzw=(0.0, 0.0, -0.70711, 0.70711)) + initial_pose = Pose( + position_xyz=(0.72193, -0.04727, -0.92512), + rotation_xyzw=(0.0, 0.0, -0.70711, 0.70711), + ) object_min_z = -0.2 def __init__(self): @@ -204,6 +207,7 @@ class MapleTableRobolab(LibraryBackground): name = "maple_table_robolab" tags = ["background", "robolab"] usd_path = f"{ARENA_NUCLEUS_DIR}/Arena/assets/object_library/srl_robolab_assets/scenes/maple_table.usda" + spawn_cfg_addon = {"rigid_props": sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True)} object_min_z = -0.05 def __init__(self): diff --git a/isaaclab_arena/tasks/gear_assembly/__init__.py b/isaaclab_arena/tasks/gear_assembly/__init__.py new file mode 100644 index 0000000000..c622213604 --- /dev/null +++ b/isaaclab_arena/tasks/gear_assembly/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2026, 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 + +"""Arena Gear Assembly task package.""" + +from .assets import * # noqa: F403 +from .rewards import * # noqa: F403 +from .specs import * # noqa: F403 +from .task import * # noqa: F403 diff --git a/isaaclab_arena/tasks/gear_assembly/assets.py b/isaaclab_arena/tasks/gear_assembly/assets.py new file mode 100644 index 0000000000..33d836e135 --- /dev/null +++ b/isaaclab_arena/tasks/gear_assembly/assets.py @@ -0,0 +1,386 @@ +# Copyright (c) 2026, 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 + +"""Scene assets for the Arena Gear Assembly task.""" + +from __future__ import annotations + +import isaaclab.sim as sim_utils +from isaaclab.assets import RigidObjectCfg +from isaaclab.sim.spawners.from_files.from_files_cfg import UsdFileCfg +from isaaclab.sim.utils import clone +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR + +from isaaclab_arena.assets.object import Object +from isaaclab_arena.assets.object_base import ObjectType +from isaaclab_arena.tasks.gear_assembly.specs import MAPLE_TABLE_TOP_COLLISION_SIZE, MAPLE_TABLE_TOP_COLLISION_THICKNESS +from isaaclab_arena.utils.pose import Pose + +GEAR_ASSET_ROOT = f"{ISAAC_NUCLEUS_DIR}/Props/Factory/gear_assets" +GEAR_GREEN_DIFFUSE_COLOR = (0.0, 0.8, 0.2) +GEAR_GREEN_VISUAL_MATERIAL_PATH = "green_material" +MAPLE_TABLE_TOP_COLLISION_COLOR = (0.43, 0.28, 0.15) +MAPLE_TABLE_LEG_RENDER_COLOR = (0.2, 0.22, 0.24) +NEWTON_GEAR_CONTACT_OFFSET = 0.001 +NEWTON_GEAR_MESH_APPROXIMATION = "convexDecomposition" +NEWTON_GEAR_BASE_MESH_APPROXIMATION = "convexDecomposition" +NEWTON_GEAR_LINEAR_DAMPING = 12.0 +NEWTON_GEAR_ANGULAR_DAMPING = 24.0 +NEWTON_GEAR_MAX_DEPENETRATION_VELOCITY = 1.0 + + +class GearAssemblyRigidObject(Object): + """Rigid object wrapper preserving the source task's contact-sensor setting.""" + + def _generate_rigid_cfg(self) -> RigidObjectCfg: + assert self.object_type == ObjectType.RIGID + object_cfg = RigidObjectCfg( + prim_path=self.prim_path, + spawn=self._get_spawn_cfg(activate_contact_sensors=False), + **self.asset_cfg_addon, + ) + return self._add_initial_pose_to_cfg(object_cfg) + + +@clone +def spawn_newton_mesh_collision_usd( + prim_path: str, + cfg: UsdFileCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + **kwargs, +): + """Spawn a USD while making its mesh collision leaves visible to Newton.""" + from isaaclab.sim import schemas + from isaaclab.sim.utils import ( + bind_visual_material, + create_prim, + get_current_stage, + make_uninstanceable, + select_usd_variants, + ) + from isaaclab.utils.assets import check_file_path, retrieve_file_path + from isaaclab.utils.version import has_kit + + from isaaclab_arena.utils.usd.newton import ensure_newton_valid_rigid_body_inertias_usd + + usd_path = cfg.usd_path + usd_path = ensure_newton_valid_rigid_body_inertias_usd(usd_path) + file_status = check_file_path(usd_path) + if file_status == 0: + raise FileNotFoundError(f"USD file not found at path: {usd_path}") + if file_status == 2: + usd_path = retrieve_file_path(usd_path, force_download=False) + + stage = get_current_stage() + if not stage.GetPrimAtPath(prim_path).IsValid(): + create_prim( + prim_path, + usd_path=usd_path, + translation=translation, + orientation=orientation, + scale=cfg.scale, + stage=stage, + ) + + if cfg.variants is not None: + select_usd_variants(prim_path, cfg.variants) + + make_uninstanceable(prim_path, stage=stage) + _author_newton_mesh_collision_leaves(stage, prim_path) + + if cfg.rigid_props is not None: + schemas.modify_rigid_body_properties(prim_path, cfg.rigid_props) + if cfg.collision_props is not None: + schemas.modify_collision_properties(prim_path, cfg.collision_props) + if cfg.mass_props is not None: + schemas.modify_mass_properties(prim_path, cfg.mass_props) + + if cfg.visual_material is not None and has_kit(): + material_path = ( + f"{prim_path}/{cfg.visual_material_path}" + if not cfg.visual_material_path.startswith("/") + else cfg.visual_material_path + ) + cfg.visual_material.func(material_path, cfg.visual_material) + bind_visual_material(prim_path, material_path, stage=stage) + + return stage.GetPrimAtPath(prim_path) + + +def _author_newton_mesh_collision_leaves(stage, prim_path: str) -> None: + from pxr import Usd, UsdGeom, UsdPhysics + + root = stage.GetPrimAtPath(prim_path) + approximation = NEWTON_GEAR_MESH_APPROXIMATION + if "FactoryGearBase" in prim_path: + approximation = NEWTON_GEAR_BASE_MESH_APPROXIMATION + for prim in Usd.PrimRange(root): + if "/collisions" not in str(prim.GetPath()) or prim.GetTypeName() != "Mesh": + continue + if not prim.HasAPI(UsdPhysics.CollisionAPI): + UsdPhysics.CollisionAPI.Apply(prim) + mesh_api = UsdPhysics.MeshCollisionAPI(prim) + if not prim.HasAPI(UsdPhysics.MeshCollisionAPI): + mesh_api = UsdPhysics.MeshCollisionAPI.Apply(prim) + mesh_api.CreateApproximationAttr().Set(approximation) + + imageable = UsdGeom.Imageable(prim) + if imageable: + imageable.CreatePurposeAttr().Set("default") + + +@clone +def spawn_newton_maple_table_usd( + prim_path: str, + cfg: UsdFileCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + **kwargs, +): + """Spawn the maple table with Newton-readable render colors.""" + from isaaclab.sim.spawners.from_files.from_files import spawn_from_usd + from pxr import Gf, Sdf, UsdShade + + prim = spawn_from_usd(prim_path, cfg, translation=translation, orientation=orientation, **kwargs) + stage = prim.GetStage() + _bind_newton_omnipbr_color( + stage, + f"{prim_path}/table/table_01/top", + f"{prim_path}/Looks/newton_maple_top", + MAPLE_TABLE_TOP_COLLISION_COLOR, + Gf, + Sdf, + UsdShade, + ) + for leg_index in range(4): + _bind_newton_omnipbr_color( + stage, + f"{prim_path}/table/table_01/leg_{leg_index}", + f"{prim_path}/Looks/newton_table_legs", + MAPLE_TABLE_LEG_RENDER_COLOR, + Gf, + Sdf, + UsdShade, + ) + return prim + + +def _bind_newton_omnipbr_color( + stage, + shape_path: str, + material_path: str, + color: tuple[float, float, float], + Gf, + Sdf, + UsdShade, +) -> None: + shape_prim = stage.GetPrimAtPath(shape_path) + if not shape_prim.IsValid(): + return + + material = UsdShade.Material.Define(stage, material_path) + shader = UsdShade.Shader.Define(stage, f"{material_path}/OmniPBRShader") + shader_prim = shader.GetPrim() + shader_prim.CreateAttribute("info:mdl:sourceAsset", Sdf.ValueTypeNames.Asset).Set(Sdf.AssetPath("OmniPBR.mdl")) + shader_prim.CreateAttribute("info:mdl:sourceAsset:subIdentifier", Sdf.ValueTypeNames.Token).Set("OmniPBR") + shader.CreateInput("diffuse_color_constant", Sdf.ValueTypeNames.Color3f).Set(Gf.Vec3f(*color)) + shader.CreateInput("diffuse_tint", Sdf.ValueTypeNames.Color3f).Set(Gf.Vec3f(1.0, 1.0, 1.0)) + UsdShade.MaterialBindingAPI.Apply(shape_prim) + UsdShade.MaterialBindingAPI(shape_prim).Bind(material, bindingStrength=UsdShade.Tokens.strongerThanDescendants) + + +def _author_display_color(stage, prim_path: str, color: tuple[float, float, float]) -> None: + from pxr import Gf, Sdf, UsdGeom + + prim = stage.GetPrimAtPath(prim_path) + if not prim.IsValid(): + return + display_color = UsdGeom.PrimvarsAPI(prim).CreatePrimvar( + "displayColor", Sdf.ValueTypeNames.Color3fArray, UsdGeom.Tokens.constant, 1 + ) + display_color.Set([Gf.Vec3f(*color)]) + + +@clone +def spawn_maple_table_top_collision( + prim_path: str, + cfg: sim_utils.CuboidCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + **kwargs, +): + """Spawn the tabletop collision proxy with a Newton-visible maple color.""" + from isaaclab.sim.spawners.shapes.shapes import spawn_cuboid + + prim = spawn_cuboid(prim_path, cfg, translation=translation, orientation=orientation, **kwargs) + stage = prim.GetStage() + _author_display_color(stage, prim_path, MAPLE_TABLE_TOP_COLLISION_COLOR) + _author_display_color(stage, f"{prim_path}/geometry/mesh", MAPLE_TABLE_TOP_COLLISION_COLOR) + return prim + + +def _gear_spawn_cfg( + kinematic_enabled: bool, + newton_mesh_collisions: bool, + visual_diffuse_color: tuple[float, float, float] | None = None, +) -> dict: + max_depenetration_velocity = 5.0 + linear_damping = 0.0 + angular_damping = 0.0 + contact_offset = 0.02 + if newton_mesh_collisions: + max_depenetration_velocity = NEWTON_GEAR_MAX_DEPENETRATION_VELOCITY + linear_damping = NEWTON_GEAR_LINEAR_DAMPING + angular_damping = NEWTON_GEAR_ANGULAR_DAMPING + contact_offset = NEWTON_GEAR_CONTACT_OFFSET + + cfg = { + "rigid_props": sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, + kinematic_enabled=kinematic_enabled, + max_depenetration_velocity=max_depenetration_velocity, + linear_damping=linear_damping, + angular_damping=angular_damping, + max_linear_velocity=1000.0, + max_angular_velocity=3666.0, + enable_gyroscopic_forces=True, + solver_position_iteration_count=32, + solver_velocity_iteration_count=1, + max_contact_impulse=1e32, + ), + "mass_props": sim_utils.MassPropertiesCfg(mass=None), + "collision_props": sim_utils.CollisionPropertiesCfg(contact_offset=contact_offset, rest_offset=0.0), + } + if visual_diffuse_color is not None: + cfg["visual_material"] = sim_utils.PreviewSurfaceCfg(diffuse_color=visual_diffuse_color, roughness=0.55) + cfg["visual_material_path"] = GEAR_GREEN_VISUAL_MATERIAL_PATH + if newton_mesh_collisions: + cfg["func"] = spawn_newton_mesh_collision_usd + return cfg + + +def make_factory_gear( + name: str, + prim_name: str, + usd_leaf: str, + pose: Pose, + kinematic_enabled: bool = False, + newton_mesh_collisions: bool = False, + visual_diffuse_color: tuple[float, float, float] | None = None, +) -> Object: + """Create one source-parity Factory gear rigid object.""" + gear = GearAssemblyRigidObject( + name=name, + prim_path=f"{{ENV_REGEX_NS}}/{prim_name}", + object_type=ObjectType.RIGID, + usd_path=f"{GEAR_ASSET_ROOT}/{usd_leaf}/{usd_leaf}.usd", + initial_pose=pose, + spawn_cfg_addon=_gear_spawn_cfg( + kinematic_enabled=kinematic_enabled, + newton_mesh_collisions=newton_mesh_collisions, + visual_diffuse_color=visual_diffuse_color, + ), + ) + gear.disable_reset_pose() + return gear + + +def make_factory_gear_base(pose: Pose, newton_mesh_collisions: bool = False) -> Object: + """Create the kinematic source Factory gear base.""" + return make_factory_gear( + name="factory_gear_base", + prim_name="FactoryGearBase", + usd_leaf="factory_gear_base", + pose=pose, + kinematic_enabled=True, + newton_mesh_collisions=newton_mesh_collisions, + ) + + +def make_factory_gear_small(pose: Pose, newton_mesh_collisions: bool = False) -> Object: + """Create the source small Factory gear.""" + return make_factory_gear( + "factory_gear_small", + "FactoryGearSmall", + "factory_gear_small", + pose, + newton_mesh_collisions=newton_mesh_collisions, + visual_diffuse_color=GEAR_GREEN_DIFFUSE_COLOR, + ) + + +def make_factory_gear_medium(pose: Pose, newton_mesh_collisions: bool = False) -> Object: + """Create the source medium Factory gear.""" + return make_factory_gear( + "factory_gear_medium", + "FactoryGearMedium", + "factory_gear_medium", + pose, + newton_mesh_collisions=newton_mesh_collisions, + visual_diffuse_color=GEAR_GREEN_DIFFUSE_COLOR, + ) + + +def make_factory_gear_large(pose: Pose, newton_mesh_collisions: bool = False) -> Object: + """Create the source large Factory gear.""" + return make_factory_gear( + "factory_gear_large", + "FactoryGearLarge", + "factory_gear_large", + pose, + newton_mesh_collisions=newton_mesh_collisions, + visual_diffuse_color=GEAR_GREEN_DIFFUSE_COLOR, + ) + + +def make_ground() -> Object: + """Create the source ground plane.""" + return Object( + name="ground", + prim_path="/World/ground", + object_type=ObjectType.BASE, + spawner_cfg=sim_utils.GroundPlaneCfg(), + initial_pose=Pose(position_xyz=(0.0, 0.0, -1.05)), + ) + + +def make_maple_table_top_collision(pose: Pose) -> Object: + """Create a Newton-safe collision surface for the maple tabletop.""" + collider_pose = Pose( + position_xyz=( + pose.position_xyz[0], + pose.position_xyz[1], + pose.position_xyz[2] - MAPLE_TABLE_TOP_COLLISION_THICKNESS / 2.0, + ), + rotation_xyzw=pose.rotation_xyzw, + ) + return Object( + name="maple_table_top_collision", + prim_path="{ENV_REGEX_NS}/maple_table_top_collision", + object_type=ObjectType.RIGID, + spawner_cfg=sim_utils.CuboidCfg( + func=spawn_maple_table_top_collision, + size=(*MAPLE_TABLE_TOP_COLLISION_SIZE, MAPLE_TABLE_TOP_COLLISION_THICKNESS), + rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True), + collision_props=sim_utils.CollisionPropertiesCfg(contact_offset=NEWTON_GEAR_CONTACT_OFFSET), + visible=True, + ), + initial_pose=collider_pose, + tags=["background", "collision"], + ) + + +def make_stand() -> Object: + """Create the source vertical stand asset.""" + return Object( + name="stand", + prim_path="{ENV_REGEX_NS}/Stand", + object_type=ObjectType.BASE, + spawner_cfg=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/Stand/stand_instanceable.usd", + scale=(2.0, 2.0, 2.0), + ), + ) diff --git a/isaaclab_arena/tasks/gear_assembly/events.py b/isaaclab_arena/tasks/gear_assembly/events.py new file mode 100644 index 0000000000..85b5ec2bbe --- /dev/null +++ b/isaaclab_arena/tasks/gear_assembly/events.py @@ -0,0 +1,156 @@ +# Copyright (c) 2026, 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 + +"""Event terms for Arena Gear Assembly.""" + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +import isaaclab.utils.math as math_utils +from isaaclab.managers import EventTermCfg, ManagerTermBase + +from isaaclab_arena.tasks.gear_assembly.specs import GEAR_TABLETOP_ORIENTATION_XYZW, GEAR_TABLETOP_PARKING_POSITIONS + +if TYPE_CHECKING: + from isaaclab.assets import Articulation, RigidObject + from isaaclab.envs import ManagerBasedEnv + + +class randomize_gears_and_base_pose_with_inactive_gear_parking(ManagerTermBase): + """Randomize Gear Assembly poses while parking inactive gear variants on the tabletop.""" + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): + """Initialize cached gear asset names and type indices.""" + super().__init__(cfg, env) + self.gear_type_map = {"gear_small": 0, "gear_medium": 1, "gear_large": 2} + self.gear_type_indices = torch.zeros(env.num_envs, device=env.device, dtype=torch.long) + self.gear_asset_names = ["factory_gear_small", "factory_gear_medium", "factory_gear_large"] + self.base_asset_name = "factory_gear_base" + + def __call__( + self, + env: ManagerBasedEnv, + env_ids: torch.Tensor, + pose_range: dict | None = None, + velocity_range: dict | None = None, + gear_pos_range: dict | None = None, + parking_positions: dict[str, tuple[float, float, float]] | None = None, + parking_orientation_xyzw: tuple[float, float, float, float] | None = None, + selected_parking_positions: dict[str, tuple[float, float, float]] | None = None, + selected_orientation_xyzw: tuple[float, float, float, float] | None = None, + parking_offsets: dict[str, tuple[float, float, float]] | None = None, + ): + """Randomize the active gear and place inactive gear variants in non-overlapping tabletop positions.""" + if not hasattr(env, "_gear_type_manager"): + raise RuntimeError( + "Gear type manager not initialized. Ensure randomize_gear_type is configured before this event." + ) + + pose_range = pose_range or {} + velocity_range = velocity_range or {} + gear_pos_range = gear_pos_range or {} + parking_positions = parking_positions or GEAR_TABLETOP_PARKING_POSITIONS + parking_orientation_xyzw = parking_orientation_xyzw or GEAR_TABLETOP_ORIENTATION_XYZW + selected_orientation_xyzw = selected_orientation_xyzw or parking_orientation_xyzw + + gear_type_manager = env._gear_type_manager + device = env.device + + pose_keys = ["x", "y", "z", "roll", "pitch", "yaw"] + ranges_pose = torch.tensor([pose_range.get(key, (0.0, 0.0)) for key in pose_keys], device=device) + rand_pose_samples = math_utils.sample_uniform( + ranges_pose[:, 0], ranges_pose[:, 1], (len(env_ids), 6), device=device + ) + orientations_delta = math_utils.quat_from_euler_xyz( + rand_pose_samples[:, 3], rand_pose_samples[:, 4], rand_pose_samples[:, 5] + ) + + ranges_vel = torch.tensor([velocity_range.get(key, (0.0, 0.0)) for key in pose_keys], device=device) + rand_vel_samples = math_utils.sample_uniform( + ranges_vel[:, 0], ranges_vel[:, 1], (len(env_ids), 6), device=device + ) + + positions_by_asset = {} + default_positions_by_asset = {} + default_orientations_by_asset = {} + orientations_by_asset = {} + velocities_by_asset = {} + for asset_name in [self.base_asset_name, *self.gear_asset_names]: + asset: RigidObject | Articulation = env.scene[asset_name] + default_root_pose = asset.data.default_root_pose.torch[env_ids].clone() + default_root_vel = asset.data.default_root_vel.torch[env_ids].clone() + default_positions_by_asset[asset_name] = default_root_pose[:, 0:3] + env.scene.env_origins[env_ids] + default_orientations_by_asset[asset_name] = default_root_pose[:, 3:7] + positions_by_asset[asset_name] = default_positions_by_asset[asset_name] + rand_pose_samples[:, 0:3] + orientations_by_asset[asset_name] = math_utils.quat_mul(default_root_pose[:, 3:7], orientations_delta) + velocities_by_asset[asset_name] = default_root_vel + rand_vel_samples + + ranges_gear = torch.tensor( + [gear_pos_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]], + device=device, + ) + rand_gear_offsets = math_utils.sample_uniform( + ranges_gear[:, 0], ranges_gear[:, 1], (len(env_ids), 3), device=device + ) + + num_reset_envs = len(env_ids) + gear_type_indices = self.gear_type_indices[:num_reset_envs] + gear_type_indices[:] = gear_type_manager.get_all_gear_type_indices()[env_ids] + parking_position_tensor = torch.tensor( + [parking_positions[gear_key] for gear_key in ["gear_small", "gear_medium", "gear_large"]], + device=device, + dtype=torch.float32, + ) + parking_orientation_tensor = torch.tensor(parking_orientation_xyzw, device=device, dtype=torch.float32) + selected_position_tensor = None + if selected_parking_positions is not None: + selected_position_tensor = torch.tensor( + [selected_parking_positions[gear_key] for gear_key in ["gear_small", "gear_medium", "gear_large"]], + device=device, + dtype=torch.float32, + ) + selected_orientation_tensor = torch.tensor(selected_orientation_xyzw, device=device, dtype=torch.float32) + parking_offset_tensor = None + if parking_offsets is not None: + parking_offset_tensor = torch.tensor( + [parking_offsets[gear_key] for gear_key in ["gear_small", "gear_medium", "gear_large"]], + device=device, + dtype=torch.float32, + ) + + for gear_idx, asset_name in enumerate(self.gear_asset_names): + selected_mask = gear_type_indices == gear_idx + if selected_position_tensor is None: + positions_by_asset[asset_name][selected_mask] += rand_gear_offsets[selected_mask] + else: + positions_by_asset[asset_name][selected_mask] = ( + env.scene.env_origins[env_ids][selected_mask] + selected_position_tensor[gear_idx] + ) + orientations_by_asset[asset_name][selected_mask] = selected_orientation_tensor + velocities_by_asset[asset_name][selected_mask] = torch.zeros_like( + velocities_by_asset[asset_name][selected_mask] + ) + if parking_offset_tensor is None: + positions_by_asset[asset_name][~selected_mask] = ( + env.scene.env_origins[env_ids][~selected_mask] + parking_position_tensor[gear_idx] + ) + else: + positions_by_asset[asset_name][~selected_mask] = ( + default_positions_by_asset[asset_name][~selected_mask] + parking_offset_tensor[gear_idx] + ) + orientations_by_asset[asset_name][~selected_mask] = parking_orientation_tensor + velocities_by_asset[asset_name][~selected_mask] = torch.zeros_like( + velocities_by_asset[asset_name][~selected_mask] + ) + + for asset_name, positions in positions_by_asset.items(): + asset = env.scene[asset_name] + asset.write_root_pose_to_sim_index( + root_pose=torch.cat([positions, orientations_by_asset[asset_name]], dim=-1), + env_ids=env_ids, + ) + asset.write_root_velocity_to_sim_index(root_velocity=velocities_by_asset[asset_name], env_ids=env_ids) diff --git a/isaaclab_arena/tasks/gear_assembly/rewards.py b/isaaclab_arena/tasks/gear_assembly/rewards.py new file mode 100644 index 0000000000..c191eb7cdc --- /dev/null +++ b/isaaclab_arena/tasks/gear_assembly/rewards.py @@ -0,0 +1,134 @@ +# Copyright (c) 2026, 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 + +"""Gear Assembly reward terms with Arena-compatible aliases.""" + +from __future__ import annotations + +import torch + +from isaaclab.managers import RewardTermCfg, SceneEntityCfg +from isaaclab_tasks.manager_based.manipulation.deploy.mdp.rewards import ( # noqa: F401 + keypoint_ee_grasp_error, + keypoint_ee_grasp_error_exp, + keypoint_entity_error, + keypoint_entity_error_exp, +) + + +def _normalize_ee_threshold_param(cfg: RewardTermCfg) -> None: + if "ee_gear_threshold" in cfg.params and "ee_grasp_threshold" not in cfg.params: + cfg.params["ee_grasp_threshold"] = cfg.params["ee_gear_threshold"] + + +class keypoint_ee_gear_error(keypoint_ee_grasp_error): + """Compatibility alias for the source grasp-corrected EE/gear keypoint penalty.""" + + def __init__(self, cfg: RewardTermCfg, env): + _normalize_ee_threshold_param(cfg) + super().__init__(cfg, env) + + def __call__( + self, + env, + robot_asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + end_effector_body_name: str = "", + grasp_rot_offset: list | None = None, + gear_offsets_grasp: dict | None = None, + keypoint_scale: float = 1.0, + add_cube_center_kp: bool = True, + weight_ramp_start: float = 0.0, + weight_ramp_steps: int = 1, + ee_grasp_threshold: float = 0.0, + ee_gear_threshold: float | None = None, + ) -> torch.Tensor: + if self.eef_idx is None: + return torch.zeros(env.num_envs, device=env.device) + + eef_pos, eef_quat, gear_grasp_pos, gear_quat_grasp = self._get_grasp_corrected_target(env) + keypoint_dist_sep = self.keypoint_computer.compute( + current_pos=eef_pos, + current_quat=eef_quat, + target_pos=gear_grasp_pos, + target_quat=gear_quat_grasp, + keypoint_scale=keypoint_scale, + ) + mean_kp_error = keypoint_dist_sep.mean(-1) + threshold = ee_grasp_threshold if ee_gear_threshold is None else ee_gear_threshold + is_active = (mean_kp_error > threshold).float() + weight_scale = self._get_weight_scale(env) + scaled_reward = mean_kp_error * weight_scale * is_active + + _log_extra(env, "ee_grasp_kp_error/mean_keypoint_dist", mean_kp_error.mean().item()) + _log_extra(env, "ee_grasp_kp_error/pct_envs_active", is_active.mean().item()) + _log_extra(env, "ee_grasp_kp_error/weight_scale", weight_scale) + return scaled_reward + + +class keypoint_ee_gear_error_exp(keypoint_ee_grasp_error_exp): + """Compatibility alias for the source grasp-corrected exponential EE/gear reward.""" + + def __init__(self, cfg: RewardTermCfg, env): + _normalize_ee_threshold_param(cfg) + super().__init__(cfg, env) + + def __call__( + self, + env, + robot_asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + end_effector_body_name: str = "", + grasp_rot_offset: list | None = None, + gear_offsets_grasp: dict | None = None, + kp_exp_coeffs: list[tuple[float, float]] = [(1.0, 0.1)], + kp_use_sum_of_exps: bool = True, + keypoint_scale: float = 1.0, + add_cube_center_kp: bool = True, + weight_ramp_start: float = 0.0, + weight_ramp_steps: int = 1, + ee_grasp_threshold: float = 0.0, + ee_gear_threshold: float | None = None, + ) -> torch.Tensor: + if self.eef_idx is None: + return torch.zeros(env.num_envs, device=env.device) + + eef_pos, eef_quat, gear_grasp_pos, gear_quat_grasp = self._get_grasp_corrected_target(env) + keypoint_dist_sep = self.keypoint_computer.compute( + current_pos=eef_pos, + current_quat=eef_quat, + target_pos=gear_grasp_pos, + target_quat=gear_quat_grasp, + keypoint_scale=keypoint_scale, + ) + mean_kp_error = keypoint_dist_sep.mean(-1) + threshold = ee_grasp_threshold if ee_gear_threshold is None else ee_gear_threshold + is_active = (mean_kp_error > threshold).float() + + keypoint_reward_exp = torch.zeros_like(keypoint_dist_sep[:, 0]) + if kp_use_sum_of_exps: + for coeff in kp_exp_coeffs: + a, b = coeff + keypoint_reward_exp += ( + 1.0 / (torch.exp(a * keypoint_dist_sep) + b + torch.exp(-a * keypoint_dist_sep)) + ).mean(-1) + else: + kp_dist_mean = keypoint_dist_sep.mean(-1) + for coeff in kp_exp_coeffs: + a, b = coeff + keypoint_reward_exp += 1.0 / (torch.exp(a * kp_dist_mean) + b + torch.exp(-a * kp_dist_mean)) + + weight_scale = self._get_weight_scale(env) + scaled_reward = keypoint_reward_exp * weight_scale * is_active + + _log_extra(env, "ee_grasp_kp_error_exp/mean_keypoint_dist", mean_kp_error.mean().item()) + _log_extra(env, "ee_grasp_kp_error_exp/mean_exp_reward", keypoint_reward_exp.mean().item()) + _log_extra(env, "ee_grasp_kp_error_exp/pct_envs_active", is_active.mean().item()) + _log_extra(env, "ee_grasp_kp_error_exp/weight_scale", weight_scale) + return scaled_reward + + +def _log_extra(env, key: str, value: float) -> None: + if not hasattr(env, "extras"): + env.extras = {} + env.extras.setdefault("log", {})[key] = value diff --git a/isaaclab_arena/tasks/gear_assembly/specs.py b/isaaclab_arena/tasks/gear_assembly/specs.py new file mode 100644 index 0000000000..3a7c00ccd3 --- /dev/null +++ b/isaaclab_arena/tasks/gear_assembly/specs.py @@ -0,0 +1,184 @@ +# Copyright (c) 2026, 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 + +"""Parity constants and Droid-specific setup for Gear Assembly.""" + +from __future__ import annotations + +import math +import torch +from collections.abc import Callable +from dataclasses import dataclass +from typing import Literal + +from isaaclab_arena.utils.pose import Pose + +DroidGearAssemblyEmbodiment = Literal["droid_abs_joint_pos", "droid_rel_joint_pos", "droid_differential_ik"] +GearAssemblyMode = Literal["play", "randomized"] +PhysicsBackend = Literal["newton", "physx"] + +DROID_GEAR_ASSEMBLY_EMBODIMENTS: tuple[DroidGearAssemblyEmbodiment, ...] = ( + "droid_abs_joint_pos", + "droid_rel_joint_pos", + "droid_differential_ik", +) + +GEAR_TYPES = ("gear_small", "gear_medium", "gear_large") +GEAR_OFFSETS = { + "gear_small": [0.076125, 0.0, 0.0], + "gear_medium": [0.030375, 0.0, 0.0], + "gear_large": [-0.045375, 0.0, 0.0], +} +DROID_BASE_GEAR_POSE = Pose(position_xyz=(0.481, -0.073, 0.071), rotation_xyzw=(0.0, 0.0, 0.70711, -0.70711)) +MAPLE_TABLE_TOP_Z = 0.003000684082508087 +MAPLE_TABLE_POSE = Pose(position_xyz=(0.0, 0.0, DROID_BASE_GEAR_POSE.position_xyz[2] - MAPLE_TABLE_TOP_Z)) +MAPLE_TABLE_TOP_COLLISION_SIZE = (0.7, 1.0) +MAPLE_TABLE_TOP_COLLISION_THICKNESS = 0.02 +MAPLE_TABLE_TOP_COLLISION_POSE = Pose( + position_xyz=( + 0.5485909044742584, + 0.02206302247941494, + DROID_BASE_GEAR_POSE.position_xyz[2], + ) +) +GEAR_TABLETOP_PARKING_Z = MAPLE_TABLE_TOP_COLLISION_POSE.position_xyz[2] + 0.043 +GEAR_TABLETOP_PARKING_POSITIONS = { + "gear_small": ( + MAPLE_TABLE_TOP_COLLISION_POSE.position_xyz[0] - 0.20, + MAPLE_TABLE_TOP_COLLISION_POSE.position_xyz[1] + 0.12, + GEAR_TABLETOP_PARKING_Z, + ), + "gear_medium": ( + MAPLE_TABLE_TOP_COLLISION_POSE.position_xyz[0] + 0.20, + MAPLE_TABLE_TOP_COLLISION_POSE.position_xyz[1] + 0.12, + GEAR_TABLETOP_PARKING_Z, + ), + "gear_large": ( + MAPLE_TABLE_TOP_COLLISION_POSE.position_xyz[0], + MAPLE_TABLE_TOP_COLLISION_POSE.position_xyz[1] + 0.32, + GEAR_TABLETOP_PARKING_Z, + ), +} +GEAR_TABLETOP_ORIENTATION_XYZW = (1.0, 0.0, 0.0, 0.0) +GEAR_INACTIVE_TABLE_PARKING_Z = GEAR_TABLETOP_PARKING_Z +GEAR_INACTIVE_PARKING_POSITIONS = GEAR_TABLETOP_PARKING_POSITIONS +GEAR_INACTIVE_TABLETOP_ORIENTATION_XYZW = GEAR_TABLETOP_ORIENTATION_XYZW +GEAR_ASSEMBLED_ROOT_Z_ABOVE_BASE = { + "gear_small": 0.0135, + "gear_medium": 0.03, + "gear_large": 0.0275, +} +GEAR_ASSEMBLED_XY_THRESHOLD = 0.015 +GEAR_ASSEMBLED_Z_THRESHOLD = 0.01 +GEAR_ASSEMBLED_UPRIGHT_AXIS_THRESHOLD_DEG = 15.0 +GEAR_ASSEMBLED_LINEAR_VELOCITY_THRESHOLD = 0.05 +GEAR_ASSEMBLED_ANGULAR_VELOCITY_THRESHOLD = 0.5 +GEAR_ASSEMBLED_SUPPORT_Z_OFFSET = { + "gear_small": -0.017, + "gear_medium": 0.0, + "gear_large": -0.003, +} +GEAR_ASSEMBLED_SUPPORT_Z_THRESHOLD = 0.005 +GEAR_ASSEMBLED_CONSECUTIVE_SUCCESS_STEPS = 10 + +GEAR_POSE_RANGE = { + "x": [-0.1, 0.1], + "y": [-0.25, 0.25], + "z": [0.0, 0.0], + "roll": [0.0, 0.0], + "pitch": [0.0, 0.0], + "yaw": [-math.pi / 6, math.pi / 6], +} +SELECTED_GEAR_POS_RANGE = { + "x": [-0.02, 0.02], + "y": [-0.02, 0.02], + "z": [0.0575, 0.0775], +} + +DROID_ARM_JOINT_NAMES = [ + "panda_joint1", + "panda_joint2", + "panda_joint3", + "panda_joint4", + "panda_joint5", + "panda_joint6", + "panda_joint7", +] + + +@dataclass(frozen=True) +class GearAssemblyRobotSpec: + """Robot-specific Gear Assembly config consumed by the source manager terms.""" + + name: str + joint_names: list[str] + num_arm_joints: int + end_effector_body_name: str + grasp_rot_offset: list[float] + gear_offsets_grasp: dict[str, list[float]] + hand_grasp_width: dict[str, float] + hand_close_width: dict[str, float] + gripper_joint_setter_func: Callable[[torch.Tensor, list[int], list[int], float], None] + state_space: int + observation_space: int + startup_materials: dict[str, tuple[float, float, float]] + reset_randomizes_robot: bool + set_grasp_pos_randomization_range: dict[str, list[float]] + + +def get_droid_robot_spec() -> GearAssemblyRobotSpec: + """Return the Gear Assembly spec for Arena's Droid Franka/Robotiq embodiment.""" + return GearAssemblyRobotSpec( + name="droid", + joint_names=list(DROID_ARM_JOINT_NAMES), + num_arm_joints=7, + end_effector_body_name="base_link", + grasp_rot_offset=[math.sqrt(2.0) / 2.0, math.sqrt(2.0) / 2.0, 0.0, 0.0], + gear_offsets_grasp={ + "gear_small": [0.0, GEAR_OFFSETS["gear_small"][0], -0.19], + "gear_medium": [0.0, GEAR_OFFSETS["gear_medium"][0], -0.19], + "gear_large": [0.0, GEAR_OFFSETS["gear_large"][0], -0.19], + }, + hand_grasp_width={"gear_small": 0.64, "gear_medium": 0.46, "gear_large": 0.4}, + hand_close_width={"gear_small": 0.69, "gear_medium": 0.51, "gear_large": 0.45}, + gripper_joint_setter_func=_set_droid_gripper_joint_pos, + state_space=28, + observation_space=21, + startup_materials={ + "factory_gear_small": (0.75, 0.75, 0.0), + "factory_gear_medium": (0.75, 0.75, 0.0), + "factory_gear_large": (0.75, 0.75, 0.0), + "factory_gear_base": (0.75, 0.75, 0.0), + "robot": (0.75, 0.75, 0.0), + }, + reset_randomizes_robot=False, + set_grasp_pos_randomization_range={ + "x": [-0.0, 0.0], + "y": [-0.005, 0.005], + "z": [-0.003, 0.003], + }, + ) + + +def gear_pose_for_mode(mode: GearAssemblyMode) -> Pose: + """Return the construction gear/base pose for Droid Gear Assembly.""" + return DROID_BASE_GEAR_POSE + + +def _set_droid_gripper_joint_pos( + joint_pos: torch.Tensor, + reset_ind_joint_pos: list[int], + gripper_joints: list[int], + joint_position: float, +) -> None: + """Set Droid's Robotiq gripper joints for the source reset term.""" + assert len(gripper_joints) >= 6, f"Droid gripper requires at least 6 gripper joints, got {len(gripper_joints)}" + for idx in reset_ind_joint_pos: + joint_pos[idx, gripper_joints[0]] = joint_position + joint_pos[idx, gripper_joints[1]] = joint_position + joint_pos[idx, gripper_joints[2]] = -joint_position + joint_pos[idx, gripper_joints[3]] = joint_position + joint_pos[idx, gripper_joints[4]] = -joint_position + joint_pos[idx, gripper_joints[5]] = -joint_position diff --git a/isaaclab_arena/tasks/gear_assembly/task.py b/isaaclab_arena/tasks/gear_assembly/task.py new file mode 100644 index 0000000000..14dcf0ddce --- /dev/null +++ b/isaaclab_arena/tasks/gear_assembly/task.py @@ -0,0 +1,382 @@ +# Copyright (c) 2026, 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 + +"""Arena task config for source-parity Gear Assembly.""" + +from __future__ import annotations + +from dataclasses import MISSING +from typing import Any + +import isaaclab.envs.mdp as mdp +from isaaclab.managers import ( + EventTermCfg, + ObservationGroupCfg, + ObservationTermCfg, + RewardTermCfg, + SceneEntityCfg, + TerminationTermCfg, +) +from isaaclab.utils.configclass import configclass +from isaaclab.utils.noise import UniformNoiseCfg +from isaaclab_tasks.manager_based.manipulation.deploy.mdp.events import randomize_gear_type, set_robot_to_grasp_pose +from isaaclab_tasks.manager_based.manipulation.deploy.mdp.noise_models import ( + ResetSampledConstantNoiseModelCfg, + ResetSampledQuaternionNoiseModelCfg, +) +from isaaclab_tasks.manager_based.manipulation.deploy.mdp.observations import ( + gear_pos_w, + gear_quat_w, + gear_shaft_pos_w, + gear_shaft_quat_w, +) +from isaaclab_tasks.manager_based.manipulation.deploy.mdp.terminations import ( + reset_when_gear_dropped, + reset_when_gear_orientation_exceeds_threshold, +) + +from isaaclab_arena.assets.register import register_task +from isaaclab_arena.embodiments.droid.observations import gripper_pos as droid_gripper_pos +from isaaclab_arena.metrics.metric_base import MetricBase +from isaaclab_arena.tasks.gear_assembly import rewards as gear_rewards +from isaaclab_arena.tasks.gear_assembly.events import randomize_gears_and_base_pose_with_inactive_gear_parking +from isaaclab_arena.tasks.gear_assembly.specs import ( + GEAR_ASSEMBLED_ANGULAR_VELOCITY_THRESHOLD, + GEAR_ASSEMBLED_CONSECUTIVE_SUCCESS_STEPS, + GEAR_ASSEMBLED_LINEAR_VELOCITY_THRESHOLD, + GEAR_ASSEMBLED_ROOT_Z_ABOVE_BASE, + GEAR_ASSEMBLED_SUPPORT_Z_OFFSET, + GEAR_ASSEMBLED_SUPPORT_Z_THRESHOLD, + GEAR_ASSEMBLED_UPRIGHT_AXIS_THRESHOLD_DEG, + GEAR_ASSEMBLED_XY_THRESHOLD, + GEAR_ASSEMBLED_Z_THRESHOLD, + GEAR_OFFSETS, + GEAR_POSE_RANGE, + GEAR_TABLETOP_ORIENTATION_XYZW, + GEAR_TABLETOP_PARKING_POSITIONS, + GEAR_TYPES, + SELECTED_GEAR_POS_RANGE, + GearAssemblyMode, + GearAssemblyRobotSpec, +) +from isaaclab_arena.tasks.gear_assembly.terminations import selected_gear_on_base +from isaaclab_arena.tasks.task_base import TaskBase + + +@register_task +class GearAssemblyTask(TaskBase): + """Source-parity Gear Assembly task implemented through Arena composition.""" + + DEFAULT_EPISODE_LENGTH_S = 6.66 + + def __init__(self, robot_spec: GearAssemblyRobotSpec, mode: GearAssemblyMode = "play"): + super().__init__(episode_length_s=self.DEFAULT_EPISODE_LENGTH_S, task_description="Assemble the gear.") + self.robot_spec = robot_spec + self.mode = mode + self.observation_cfg = ObservationsCfg(robot_spec=robot_spec, mode=mode) + self.events_cfg = EventsCfg(robot_spec=robot_spec, mode=mode) + self.rewards_cfg = RewardsCfg(robot_spec=robot_spec) + self.termination_cfg = TerminationsCfg(robot_spec=robot_spec, mode=mode) + + def get_scene_cfg(self) -> Any: + return None + + def get_termination_cfg(self) -> Any: + return self.termination_cfg + + def get_events_cfg(self) -> Any: + return self.events_cfg + + def get_observation_cfg(self) -> Any: + return self.observation_cfg + + def get_rewards_cfg(self) -> Any: + return self.rewards_cfg + + def get_mimic_env_cfg(self, arm_mode) -> Any: + return None + + def get_metrics(self) -> list[MetricBase]: + return [] + + def runtime_env_attrs(self) -> dict[str, Any]: + """Attributes required by source Gear Assembly manager terms at env construction.""" + return { + "gear_offsets": GEAR_OFFSETS, + "gear_offsets_grasp": self.robot_spec.gear_offsets_grasp, + "hand_grasp_width": self.robot_spec.hand_grasp_width, + "hand_close_width": self.robot_spec.hand_close_width, + "end_effector_body_name": self.robot_spec.end_effector_body_name, + "num_arm_joints": self.robot_spec.num_arm_joints, + "grasp_rot_offset": self.robot_spec.grasp_rot_offset, + "gripper_joint_setter_func": self.robot_spec.gripper_joint_setter_func, + "joint_action_scale": 0.025, + } + + +@configclass +class ObservationsCfg: + """Observation terms for Gear Assembly.""" + + @configclass + class PolicyCfg(ObservationGroupCfg): + joint_pos: ObservationTermCfg = MISSING + gripper_pos: ObservationTermCfg | None = None + joint_vel: ObservationTermCfg = MISSING + gear_shaft_pos: ObservationTermCfg = MISSING + gear_shaft_quat: ObservationTermCfg = MISSING + + def __post_init__(self): + self.enable_corruption = True + self.concatenate_terms = True + + @configclass + class CriticCfg(ObservationGroupCfg): + joint_pos: ObservationTermCfg = MISSING + joint_vel: ObservationTermCfg = MISSING + gear_shaft_pos: ObservationTermCfg = MISSING + gear_shaft_quat: ObservationTermCfg = MISSING + gear_pos: ObservationTermCfg = MISSING + gear_quat: ObservationTermCfg = MISSING + + policy: PolicyCfg = MISSING + critic: CriticCfg = MISSING + + def __init__(self, robot_spec: GearAssemblyRobotSpec, mode: GearAssemblyMode): + joint_cfg = SceneEntityCfg("robot", joint_names=robot_spec.joint_names) + self.policy = self.PolicyCfg( + joint_pos=ObservationTermCfg(func=mdp.joint_pos, params={"asset_cfg": joint_cfg}), + joint_vel=ObservationTermCfg(func=mdp.joint_vel, params={"asset_cfg": joint_cfg}), + gear_shaft_pos=ObservationTermCfg( + func=gear_shaft_pos_w, + params={"gear_offsets": GEAR_OFFSETS}, + noise=ResetSampledConstantNoiseModelCfg( + noise_cfg=UniformNoiseCfg(n_min=-0.01, n_max=0.01, operation="add") + ), + ), + gear_shaft_quat=ObservationTermCfg( + func=gear_shaft_quat_w, + noise=ResetSampledQuaternionNoiseModelCfg( + roll_range=(-0.03491, 0.03491), + pitch_range=(-0.03491, 0.03491), + yaw_range=(-0.03491, 0.03491), + ), + ), + ) + if mode == "play": + self.policy.enable_corruption = False + self.policy.concatenate_terms = False + self.policy.gripper_pos = ObservationTermCfg(func=droid_gripper_pos) + self.critic = self.CriticCfg( + joint_pos=ObservationTermCfg(func=mdp.joint_pos, params={"asset_cfg": joint_cfg}), + joint_vel=ObservationTermCfg(func=mdp.joint_vel, params={"asset_cfg": joint_cfg}), + gear_shaft_pos=ObservationTermCfg(func=gear_shaft_pos_w, params={"gear_offsets": GEAR_OFFSETS}), + gear_shaft_quat=ObservationTermCfg(func=gear_shaft_quat_w), + gear_pos=ObservationTermCfg(func=gear_pos_w), + gear_quat=ObservationTermCfg(func=gear_quat_w), + ) + + +@configclass +class EventsCfg: + """Reset and startup events for Gear Assembly.""" + + robot_joint_stiffness_and_damping: EventTermCfg | None = None + joint_friction: EventTermCfg | None = None + small_gear_physics_material: EventTermCfg = MISSING + medium_gear_physics_material: EventTermCfg = MISSING + large_gear_physics_material: EventTermCfg = MISSING + gear_base_physics_material: EventTermCfg = MISSING + robot_physics_material: EventTermCfg = MISSING + randomize_gear_type: EventTermCfg = MISSING + reset_all: EventTermCfg = MISSING + randomize_gears_and_base_pose: EventTermCfg = MISSING + set_robot_to_grasp_pose: EventTermCfg = MISSING + + def __init__(self, robot_spec: GearAssemblyRobotSpec, mode: GearAssemblyMode): + self.robot_joint_stiffness_and_damping = None + self.joint_friction = None + if robot_spec.reset_randomizes_robot: + self.robot_joint_stiffness_and_damping = EventTermCfg( + func=mdp.randomize_actuator_gains, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg("robot", joint_names=["shoulder_.*", "elbow_.*", "wrist_.*"]), + "stiffness_distribution_params": (0.75, 1.5), + "damping_distribution_params": (0.3, 3.0), + "operation": "scale", + "distribution": "log_uniform", + }, + ) + self.joint_friction = EventTermCfg( + func=mdp.randomize_joint_parameters, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg("robot", joint_names=["shoulder_.*", "elbow_.*", "wrist_.*"]), + "friction_distribution_params": (0.3, 0.7), + "operation": "add", + "distribution": "uniform", + }, + ) + + self.small_gear_physics_material = _material_event("factory_gear_small", ".*", robot_spec.startup_materials) + self.medium_gear_physics_material = _material_event("factory_gear_medium", ".*", robot_spec.startup_materials) + self.large_gear_physics_material = _material_event("factory_gear_large", ".*", robot_spec.startup_materials) + self.gear_base_physics_material = _material_event("factory_gear_base", ".*", robot_spec.startup_materials) + self.robot_physics_material = _material_event("robot", ".*finger.*", robot_spec.startup_materials) + + gear_types = list(GEAR_TYPES) + pose_range = dict(GEAR_POSE_RANGE) + selected_gear_pos_range = dict(SELECTED_GEAR_POS_RANGE) + + self.randomize_gear_type = EventTermCfg( + func=randomize_gear_type, + mode="reset", + params={"gear_types": gear_types}, + ) + self.reset_all = EventTermCfg(func=mdp.reset_scene_to_default, mode="reset") + self.randomize_gears_and_base_pose = EventTermCfg( + func=randomize_gears_and_base_pose_with_inactive_gear_parking, + mode="reset", + params={ + "pose_range": pose_range, + "gear_pos_range": selected_gear_pos_range, + "parking_positions": GEAR_TABLETOP_PARKING_POSITIONS, + "parking_orientation_xyzw": GEAR_TABLETOP_ORIENTATION_XYZW, + "velocity_range": {}, + }, + ) + self.set_robot_to_grasp_pose = EventTermCfg( + func=set_robot_to_grasp_pose, + mode="reset", + params={ + "robot_asset_cfg": SceneEntityCfg("robot"), + "pos_randomization_range": robot_spec.set_grasp_pos_randomization_range, + "gear_offsets_grasp": robot_spec.gear_offsets_grasp, + "end_effector_body_name": robot_spec.end_effector_body_name, + "num_arm_joints": robot_spec.num_arm_joints, + "grasp_rot_offset": robot_spec.grasp_rot_offset, + "gripper_joint_setter_func": robot_spec.gripper_joint_setter_func, + }, + ) + + +@configclass +class RewardsCfg: + """Reward terms for Gear Assembly.""" + + end_effector_gear_keypoint_tracking: RewardTermCfg = MISSING + end_effector_gear_keypoint_tracking_exp: RewardTermCfg = MISSING + end_effector_base_keypoint_tracking: RewardTermCfg = MISSING + end_effector_base_keypoint_tracking_exp: RewardTermCfg = MISSING + action_rate: RewardTermCfg = MISSING + + def __init__(self, robot_spec: GearAssemblyRobotSpec): + self.end_effector_gear_keypoint_tracking = RewardTermCfg( + func=gear_rewards.keypoint_entity_error, + weight=-1.5, + params={"asset_cfg_1": SceneEntityCfg("factory_gear_base"), "keypoint_scale": 0.15}, + ) + self.end_effector_gear_keypoint_tracking_exp = RewardTermCfg( + func=gear_rewards.keypoint_entity_error_exp, + weight=1.5, + params={ + "asset_cfg_1": SceneEntityCfg("factory_gear_base"), + "kp_exp_coeffs": [(50, 0.0001), (300, 0.0001)], + "kp_use_sum_of_exps": False, + "keypoint_scale": 0.15, + }, + ) + ee_params = { + "robot_asset_cfg": SceneEntityCfg("robot"), + "keypoint_scale": 0.15, + "ee_gear_threshold": 0.0, + "weight_ramp_start": 0.0, + "weight_ramp_steps": 512_000, + "end_effector_body_name": robot_spec.end_effector_body_name, + "grasp_rot_offset": robot_spec.grasp_rot_offset, + "gear_offsets_grasp": robot_spec.gear_offsets_grasp, + } + self.end_effector_base_keypoint_tracking = RewardTermCfg( + func=gear_rewards.keypoint_ee_gear_error, + weight=-0.5, + params=dict(ee_params), + ) + self.end_effector_base_keypoint_tracking_exp = RewardTermCfg( + func=gear_rewards.keypoint_ee_gear_error_exp, + weight=0.5, + params={ + **ee_params, + "kp_exp_coeffs": [(50, 0.0001), (300, 0.0001)], + "kp_use_sum_of_exps": False, + }, + ) + self.action_rate = RewardTermCfg(func=mdp.action_rate_l2, weight=-5.0e-06) + + +@configclass +class TerminationsCfg: + """Termination terms for Gear Assembly.""" + + time_out: TerminationTermCfg = TerminationTermCfg(func=mdp.time_out, time_out=True) + success: TerminationTermCfg | None = None + gear_dropped: TerminationTermCfg = MISSING + gear_orientation_exceeded: TerminationTermCfg = MISSING + + def __init__(self, robot_spec: GearAssemblyRobotSpec, mode: GearAssemblyMode): + self.time_out = TerminationTermCfg(func=mdp.time_out, time_out=True) + self.success = None + if mode == "play": + self.success = TerminationTermCfg( + func=selected_gear_on_base, + params={ + "base_asset_cfg": SceneEntityCfg("factory_gear_base"), + "root_z_above_base": GEAR_ASSEMBLED_ROOT_Z_ABOVE_BASE, + "xy_threshold": GEAR_ASSEMBLED_XY_THRESHOLD, + "z_threshold": GEAR_ASSEMBLED_Z_THRESHOLD, + "upright_axis_threshold_deg": GEAR_ASSEMBLED_UPRIGHT_AXIS_THRESHOLD_DEG, + "linear_velocity_threshold": GEAR_ASSEMBLED_LINEAR_VELOCITY_THRESHOLD, + "angular_velocity_threshold": GEAR_ASSEMBLED_ANGULAR_VELOCITY_THRESHOLD, + "support_z_offset": GEAR_ASSEMBLED_SUPPORT_Z_OFFSET, + "support_z_threshold": GEAR_ASSEMBLED_SUPPORT_Z_THRESHOLD, + "consecutive_success_steps": GEAR_ASSEMBLED_CONSECUTIVE_SUCCESS_STEPS, + }, + ) + self.gear_dropped = TerminationTermCfg( + func=reset_when_gear_dropped, + params={ + "distance_threshold": 0.15, + "robot_asset_cfg": SceneEntityCfg("robot"), + "gear_offsets_grasp": robot_spec.gear_offsets_grasp, + "end_effector_body_name": robot_spec.end_effector_body_name, + "grasp_rot_offset": robot_spec.grasp_rot_offset, + }, + ) + self.gear_orientation_exceeded = TerminationTermCfg( + func=reset_when_gear_orientation_exceeds_threshold, + params={ + "roll_threshold_deg": 15.0, + "pitch_threshold_deg": 15.0, + "yaw_threshold_deg": 180.0, + "robot_asset_cfg": SceneEntityCfg("robot"), + "end_effector_body_name": robot_spec.end_effector_body_name, + "grasp_rot_offset": robot_spec.grasp_rot_offset, + }, + ) + + +def _material_event(asset_name: str, body_names: str, materials: dict[str, tuple[float, float, float]]) -> EventTermCfg: + static_friction, dynamic_friction, restitution = materials[asset_name] + return EventTermCfg( + func=mdp.randomize_rigid_body_material, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg(asset_name, body_names=body_names), + "static_friction_range": (static_friction, static_friction), + "dynamic_friction_range": (dynamic_friction, dynamic_friction), + "restitution_range": (restitution, restitution), + "num_buckets": 16, + }, + ) diff --git a/isaaclab_arena/tasks/gear_assembly/terminations.py b/isaaclab_arena/tasks/gear_assembly/terminations.py new file mode 100644 index 0000000000..994537c6e4 --- /dev/null +++ b/isaaclab_arena/tasks/gear_assembly/terminations.py @@ -0,0 +1,222 @@ +# Copyright (c) 2026, 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 + +"""Termination terms for Arena Gear Assembly.""" + +from __future__ import annotations + +import torch +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import isaaclab.sim as sim_utils +import isaaclab.utils.math as math_utils +from isaaclab.managers import ManagerTermBase, SceneEntityCfg, TerminationTermCfg + +if TYPE_CHECKING: + from isaaclab.assets import RigidObject + from isaaclab.envs import ManagerBasedEnv + from pxr import Usd + + +class selected_gear_on_base(ManagerTermBase): + """Terminate when the active gear is seated and settled on the gear base.""" + + def __init__(self, cfg: TerminationTermCfg, env: ManagerBasedEnv): + """Cache the active gear assets and base asset.""" + super().__init__(cfg, env) + self.base_asset_cfg: SceneEntityCfg = cfg.params.get("base_asset_cfg", SceneEntityCfg("factory_gear_base")) + self.base_asset = env.scene[self.base_asset_cfg.name] + self.gear_assets = { + "gear_small": env.scene["factory_gear_small"], + "gear_medium": env.scene["factory_gear_medium"], + "gear_large": env.scene["factory_gear_large"], + } + self.gear_names = ["gear_small", "gear_medium", "gear_large"] + self.env_indices = torch.arange(env.num_envs, device=env.device) + self.up_axis = torch.tensor([[0.0, 0.0, 1.0]], device=env.device, dtype=torch.float32).repeat(env.num_envs, 1) + self.consecutive_success_count = torch.zeros(env.num_envs, device=env.device, dtype=torch.int32) + self.base_collision_corners = self._collision_corners(self.base_asset, env.device) + self.gear_collision_corners = { + gear_name: self._collision_corners(asset, env.device) for gear_name, asset in self.gear_assets.items() + } + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + """Reset the consecutive success counter.""" + if env_ids is None: + env_ids = slice(None) + self.consecutive_success_count[env_ids] = 0 + + @staticmethod + def _collision_corners(asset: RigidObject, device: str) -> torch.Tensor: + from pxr import Usd, UsdGeom, UsdPhysics + + root_prims = sim_utils.find_matching_prims(asset.cfg.prim_path) + assert root_prims, f"{asset.cfg.prim_path} has no matching prims" + root_prim = root_prims[0] + rigid_prim = selected_gear_on_base._rigid_body_prim(root_prim) + assert rigid_prim is not None, f"{asset.cfg.prim_path} has no rigid-body prim" + + bbox_cache = UsdGeom.BBoxCache(0, [UsdGeom.Tokens.default_], useExtentsHint=True) + corners = [] + for prim in Usd.PrimRange(root_prim): + if not prim.HasAPI(UsdPhysics.CollisionAPI) or not prim.IsA(UsdGeom.Boundable): + continue + local_box = bbox_cache.ComputeRelativeBound(prim, rigid_prim).ComputeAlignedBox() + box_min = local_box.GetMin() + box_max = local_box.GetMax() + corners.extend( + [x, y, z] + for x in (box_min[0], box_max[0]) + for y in (box_min[1], box_max[1]) + for z in (box_min[2], box_max[2]) + ) + assert corners, f"{asset.cfg.prim_path} has no collision geometry" + return torch.tensor(corners, device=device, dtype=torch.float32) + + @staticmethod + def _rigid_body_prim(root_prim: Usd.Prim) -> Usd.Prim | None: + from pxr import Usd, UsdPhysics + + for prim in Usd.PrimRange(root_prim): + if prim.HasAPI(UsdPhysics.RigidBodyAPI): + return prim + return None + + @staticmethod + def _world_collision_z_bounds( + local_corners: torch.Tensor, root_pos: torch.Tensor, root_quat: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + num_envs = root_pos.shape[0] + num_corners = local_corners.shape[0] + corners = local_corners.unsqueeze(0).expand(num_envs, num_corners, 3).reshape(-1, 3) + quats = root_quat.unsqueeze(1).expand(num_envs, num_corners, 4).reshape(-1, 4) + positions = root_pos.unsqueeze(1).expand(num_envs, num_corners, 3).reshape(-1, 3) + world_z = (positions + math_utils.quat_apply(quats, corners))[:, 2].reshape(num_envs, num_corners) + return world_z.min(dim=1).values, world_z.max(dim=1).values + + def _selected_param( + self, value: float | dict[str, float], gear_type_indices: torch.Tensor, env: ManagerBasedEnv + ) -> torch.Tensor: + if isinstance(value, dict): + values = torch.tensor( + [value[gear_name] for gear_name in self.gear_names], device=env.device, dtype=torch.float32 + ) + return values[gear_type_indices] + return torch.full((env.num_envs,), float(value), device=env.device, dtype=torch.float32) + + def __call__( + self, + env: ManagerBasedEnv, + base_asset_cfg: SceneEntityCfg = SceneEntityCfg("factory_gear_base"), + root_z_above_base: float | dict[str, float] = 0.03, + xy_threshold: float = 0.015, + z_threshold: float = 0.01, + upright_axis_threshold_deg: float = 15.0, + linear_velocity_threshold: float = 0.05, + angular_velocity_threshold: float = 0.5, + support_z_offset: float | dict[str, float] = 0.0, + support_z_threshold: float = 0.005, + consecutive_success_steps: int = 10, + ) -> torch.Tensor: + """Return true when the selected gear is aligned with and settled on the base. + + Args: + env: Environment instance. + base_asset_cfg: Configuration of the gear-base asset. + root_z_above_base: Expected selected gear root height above the base root. + xy_threshold: Maximum allowed root XY error relative to the base root. + z_threshold: Maximum allowed root height error relative to ``root_z_above_base``. + upright_axis_threshold_deg: Maximum allowed local-Z axis angle between gear and base. + linear_velocity_threshold: Maximum selected gear linear speed. + angular_velocity_threshold: Maximum selected gear angular speed. + support_z_offset: Expected gear-bottom height relative to the base-top collision surface. + support_z_threshold: Maximum allowed error around ``support_z_offset``. + consecutive_success_steps: Number of consecutive checks required before terminating. + + Returns: + Boolean tensor indicating which environments have completed the assembly. + """ + if not hasattr(env, "_gear_type_manager"): + self.consecutive_success_count.zero_() + return torch.zeros(env.num_envs, dtype=torch.bool, device=env.device) + assert ( + base_asset_cfg.name == self.base_asset_cfg.name + ), "selected_gear_on_base does not support changing base_asset_cfg after initialization" + + gear_type_indices = env._gear_type_manager.get_all_gear_type_indices() + + all_gear_pos = torch.stack( + [ + self.gear_assets["gear_small"].data.root_link_pos_w.torch, + self.gear_assets["gear_medium"].data.root_link_pos_w.torch, + self.gear_assets["gear_large"].data.root_link_pos_w.torch, + ], + dim=1, + ) + all_gear_quat = torch.stack( + [ + self.gear_assets["gear_small"].data.root_link_quat_w.torch, + self.gear_assets["gear_medium"].data.root_link_quat_w.torch, + self.gear_assets["gear_large"].data.root_link_quat_w.torch, + ], + dim=1, + ) + all_gear_vel = torch.stack( + [ + self.gear_assets["gear_small"].data.root_com_vel_w.torch, + self.gear_assets["gear_medium"].data.root_com_vel_w.torch, + self.gear_assets["gear_large"].data.root_com_vel_w.torch, + ], + dim=1, + ) + + gear_pos = all_gear_pos[self.env_indices, gear_type_indices] + gear_quat = all_gear_quat[self.env_indices, gear_type_indices] + gear_vel = all_gear_vel[self.env_indices, gear_type_indices] + + base_pos = self.base_asset.data.root_link_pos_w.torch + base_quat = self.base_asset.data.root_link_quat_w.torch + + xy_error = torch.linalg.norm(gear_pos[:, :2] - base_pos[:, :2], dim=-1) + root_z_targets = self._selected_param(root_z_above_base, gear_type_indices, env) + z_error = torch.abs((gear_pos[:, 2] - base_pos[:, 2]) - root_z_targets) + _, base_top_z = self._world_collision_z_bounds(self.base_collision_corners, base_pos, base_quat) + gear_bottom_z = torch.empty(env.num_envs, dtype=torch.float32, device=env.device) + for gear_idx, gear_name in enumerate(self.gear_names): + mask = gear_type_indices == gear_idx + if not mask.any(): + continue + selected_gear_bottom_z, _ = self._world_collision_z_bounds( + self.gear_collision_corners[gear_name], + self.gear_assets[gear_name].data.root_link_pos_w.torch, + self.gear_assets[gear_name].data.root_link_quat_w.torch, + ) + gear_bottom_z[mask] = selected_gear_bottom_z[mask] + support_targets = self._selected_param(support_z_offset, gear_type_indices, env) + support_error = torch.abs((gear_bottom_z - base_top_z) - support_targets) + + gear_up = math_utils.quat_apply(gear_quat, self.up_axis) + base_up = math_utils.quat_apply(base_quat, self.up_axis) + min_upright_cos = torch.cos(torch.deg2rad(torch.tensor(upright_axis_threshold_deg, device=env.device))) + upright = torch.sum(gear_up * base_up, dim=-1) >= min_upright_cos + + linear_speed = torch.linalg.norm(gear_vel[:, :3], dim=-1) + angular_speed = torch.linalg.norm(gear_vel[:, 3:], dim=-1) + + success_now = ( + (xy_error <= xy_threshold) + & (z_error <= z_threshold) + & (support_error <= support_z_threshold) + & upright + & (linear_speed <= linear_velocity_threshold) + & (angular_speed <= angular_velocity_threshold) + ) + self.consecutive_success_count = torch.where( + success_now, + self.consecutive_success_count + 1, + torch.zeros_like(self.consecutive_success_count), + ) + return self.consecutive_success_count >= consecutive_success_steps diff --git a/isaaclab_arena/tasks/task_library.py b/isaaclab_arena/tasks/task_library.py index 1273b69964..b5648f8a3a 100644 --- a/isaaclab_arena/tasks/task_library.py +++ b/isaaclab_arena/tasks/task_library.py @@ -23,3 +23,4 @@ sorting_task, turn_knob_task, ) +from isaaclab_arena.tasks.gear_assembly import task as gear_assembly_task # noqa: F401 diff --git a/isaaclab_arena/tests/test_gear_assembly_environment.py b/isaaclab_arena/tests/test_gear_assembly_environment.py new file mode 100644 index 0000000000..dc59380348 --- /dev/null +++ b/isaaclab_arena/tests/test_gear_assembly_environment.py @@ -0,0 +1,681 @@ +# Copyright (c) 2026, 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 + +"""Regression checks for the Arena Gear Assembly scene setup.""" + +from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function + + +def _test_gear_assembly_scene_and_newton_cfg(simulation_app): + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder + from isaaclab_arena.environments.arena_env_builder_cfg import ArenaEnvBuilderCfg + from isaaclab_arena.tasks.gear_assembly.assets import ( + GEAR_GREEN_DIFFUSE_COLOR, + GEAR_GREEN_VISUAL_MATERIAL_PATH, + NEWTON_GEAR_ANGULAR_DAMPING, + NEWTON_GEAR_BASE_MESH_APPROXIMATION, + NEWTON_GEAR_CONTACT_OFFSET, + NEWTON_GEAR_LINEAR_DAMPING, + NEWTON_GEAR_MAX_DEPENETRATION_VELOCITY, + NEWTON_GEAR_MESH_APPROXIMATION, + spawn_maple_table_top_collision, + spawn_newton_maple_table_usd, + spawn_newton_mesh_collision_usd, + ) + from isaaclab_arena.tasks.gear_assembly.events import randomize_gears_and_base_pose_with_inactive_gear_parking + from isaaclab_arena.tasks.gear_assembly.specs import ( + DROID_BASE_GEAR_POSE, + GEAR_ASSEMBLED_ANGULAR_VELOCITY_THRESHOLD, + GEAR_ASSEMBLED_CONSECUTIVE_SUCCESS_STEPS, + GEAR_ASSEMBLED_LINEAR_VELOCITY_THRESHOLD, + GEAR_ASSEMBLED_ROOT_Z_ABOVE_BASE, + GEAR_ASSEMBLED_SUPPORT_Z_OFFSET, + GEAR_ASSEMBLED_SUPPORT_Z_THRESHOLD, + GEAR_ASSEMBLED_UPRIGHT_AXIS_THRESHOLD_DEG, + GEAR_ASSEMBLED_XY_THRESHOLD, + GEAR_ASSEMBLED_Z_THRESHOLD, + GEAR_POSE_RANGE, + GEAR_TABLETOP_ORIENTATION_XYZW, + GEAR_TABLETOP_PARKING_POSITIONS, + MAPLE_TABLE_POSE, + MAPLE_TABLE_TOP_COLLISION_POSE, + MAPLE_TABLE_TOP_COLLISION_SIZE, + MAPLE_TABLE_TOP_COLLISION_THICKNESS, + SELECTED_GEAR_POS_RANGE, + ) + from isaaclab_arena.tasks.gear_assembly.terminations import selected_gear_on_base + from isaaclab_arena_environments.gear_assembly_environment import ( + GearAssemblyEnvironment, + GearAssemblyEnvironmentCfg, + ) + + arena_env = GearAssemblyEnvironment().build(GearAssemblyEnvironmentCfg()) + + assert arena_env.name == "gear_assembly" + assert arena_env.embodiment.name == "droid_abs_joint_pos" + assert arena_env.rl_framework_entry_point is None + assert arena_env.rl_policy_cfg is None + + assets = arena_env.scene.assets + assert list(assets) == [ + "ground", + "maple_table_robolab", + "maple_table_top_collision", + "factory_gear_base", + "factory_gear_small", + "factory_gear_medium", + "factory_gear_large", + "light", + "table", + ] + assert assets["maple_table_robolab"].get_initial_pose() == MAPLE_TABLE_POSE + assert assets["maple_table_robolab"].object_cfg.spawn.rigid_props.kinematic_enabled is True + assert assets["maple_table_robolab"].object_cfg.spawn.func == spawn_newton_maple_table_usd + tabletop_collider_pose = assets["maple_table_top_collision"].get_initial_pose() + assert tabletop_collider_pose.position_xyz[:2] == MAPLE_TABLE_TOP_COLLISION_POSE.position_xyz[:2] + assert ( + abs( + tabletop_collider_pose.position_xyz[2] + - (MAPLE_TABLE_TOP_COLLISION_POSE.position_xyz[2] - MAPLE_TABLE_TOP_COLLISION_THICKNESS / 2.0) + ) + < 1e-6 + ) + assert assets["maple_table_top_collision"].object_cfg.spawn.size == ( + *MAPLE_TABLE_TOP_COLLISION_SIZE, + MAPLE_TABLE_TOP_COLLISION_THICKNESS, + ) + assert assets["maple_table_top_collision"].object_cfg.spawn.rigid_props.kinematic_enabled is True + assert assets["maple_table_top_collision"].object_cfg.spawn.visible is True + assert assets["maple_table_top_collision"].object_cfg.spawn.func == spawn_maple_table_top_collision + assert assets["maple_table_top_collision"].object_cfg.spawn.visual_material is None + assert assets["table"].parent_asset is assets["maple_table_robolab"] + assert abs(assets["table"].get_initial_pose().position_xyz[2] - DROID_BASE_GEAR_POSE.position_xyz[2]) < 1e-6 + assert assets["factory_gear_base"].object_cfg.spawn.activate_contact_sensors is False + assert assets["factory_gear_base"].object_cfg.spawn.rigid_props.kinematic_enabled is True + assert assets["factory_gear_base"].object_cfg.spawn.visual_material is None + assert assets["factory_gear_small"].object_cfg.spawn.activate_contact_sensors is False + assert assets["factory_gear_small"].object_cfg.spawn.rigid_props.kinematic_enabled is False + assert assets["factory_gear_base"].object_cfg.spawn.func == spawn_newton_mesh_collision_usd + assert assets["factory_gear_small"].object_cfg.spawn.func == spawn_newton_mesh_collision_usd + for gear_name in ( + "factory_gear_small", + "factory_gear_medium", + "factory_gear_large", + ): + gear_spawn = assets[gear_name].object_cfg.spawn + assert gear_spawn.visual_material.diffuse_color == GEAR_GREEN_DIFFUSE_COLOR + assert gear_spawn.visual_material_path == GEAR_GREEN_VISUAL_MATERIAL_PATH + assert assets["factory_gear_small"].object_cfg.spawn.rigid_props.linear_damping == NEWTON_GEAR_LINEAR_DAMPING + assert assets["factory_gear_small"].object_cfg.spawn.rigid_props.angular_damping == NEWTON_GEAR_ANGULAR_DAMPING + assert ( + assets["factory_gear_small"].object_cfg.spawn.rigid_props.max_depenetration_velocity + == NEWTON_GEAR_MAX_DEPENETRATION_VELOCITY + ) + assert assets["factory_gear_small"].object_cfg.spawn.collision_props.contact_offset == NEWTON_GEAR_CONTACT_OFFSET + assert NEWTON_GEAR_MESH_APPROXIMATION == "convexDecomposition" + assert NEWTON_GEAR_BASE_MESH_APPROXIMATION == "convexDecomposition" + assert GEAR_POSE_RANGE["z"] == [0.0, 0.0] + assert GEAR_POSE_RANGE["roll"] == [0.0, 0.0] + assert GEAR_POSE_RANGE["pitch"] == [0.0, 0.0] + assert SELECTED_GEAR_POS_RANGE == { + "x": [-0.02, 0.02], + "y": [-0.02, 0.02], + "z": [0.0575, 0.0775], + } + assert ( + arena_env.task.events_cfg.randomize_gears_and_base_pose.func + == randomize_gears_and_base_pose_with_inactive_gear_parking + ) + assert ( + arena_env.task.events_cfg.randomize_gears_and_base_pose.params["parking_positions"] + == GEAR_TABLETOP_PARKING_POSITIONS + ) + assert ( + arena_env.task.events_cfg.randomize_gears_and_base_pose.params["parking_orientation_xyzw"] + == GEAR_TABLETOP_ORIENTATION_XYZW + ) + assert "selected_parking_positions" not in arena_env.task.events_cfg.randomize_gears_and_base_pose.params + assert "selected_orientation_xyzw" not in arena_env.task.events_cfg.randomize_gears_and_base_pose.params + + builder = ArenaEnvBuilder(arena_env, ArenaEnvBuilderCfg(num_envs=1)) + env_cfg, _ = builder.compose_manager_cfg() + solver_cfg = env_cfg.sim.physics.solver_cfg + + assert type(env_cfg.sim.physics).__name__ == "NewtonCfg" + assert type(solver_cfg).__name__ == "MJWarpSolverCfg" + assert solver_cfg.solver == "newton" + assert solver_cfg.integrator == "implicitfast" + assert solver_cfg.use_mujoco_contacts is False + assert env_cfg.sim.physics.default_shape_cfg.gap == 0.0 + assert env_cfg.scene.robot.spawn.usd_path.endswith("_newton_inertia.usd") + assert env_cfg.scene.replicate_physics is True + assert ( + env_cfg.events.randomize_gears_and_base_pose.params["selected_parking_positions"] + == GEAR_TABLETOP_PARKING_POSITIONS + ) + assert ( + env_cfg.events.randomize_gears_and_base_pose.params["selected_orientation_xyzw"] + == GEAR_TABLETOP_ORIENTATION_XYZW + ) + assert env_cfg.sim.dt == 1.0 / 120.0 + assert env_cfg.decimation == 4 + assert env_cfg.episode_length_s == 6.66 + assert env_cfg.observations.policy.concatenate_terms is False + assert env_cfg.observations.policy.enable_corruption is False + assert env_cfg.observations.policy.gripper_pos is not None + assert env_cfg.terminations.success is not None + assert env_cfg.terminations.success.func == selected_gear_on_base + assert env_cfg.terminations.success.params["root_z_above_base"] == GEAR_ASSEMBLED_ROOT_Z_ABOVE_BASE + assert env_cfg.terminations.success.params["xy_threshold"] == GEAR_ASSEMBLED_XY_THRESHOLD + assert env_cfg.terminations.success.params["z_threshold"] == GEAR_ASSEMBLED_Z_THRESHOLD + assert ( + env_cfg.terminations.success.params["upright_axis_threshold_deg"] == GEAR_ASSEMBLED_UPRIGHT_AXIS_THRESHOLD_DEG + ) + assert env_cfg.terminations.success.params["linear_velocity_threshold"] == GEAR_ASSEMBLED_LINEAR_VELOCITY_THRESHOLD + assert ( + env_cfg.terminations.success.params["angular_velocity_threshold"] == GEAR_ASSEMBLED_ANGULAR_VELOCITY_THRESHOLD + ) + assert env_cfg.terminations.success.params["support_z_offset"] == GEAR_ASSEMBLED_SUPPORT_Z_OFFSET + assert env_cfg.terminations.success.params["support_z_threshold"] == GEAR_ASSEMBLED_SUPPORT_Z_THRESHOLD + assert env_cfg.terminations.success.params["consecutive_success_steps"] == GEAR_ASSEMBLED_CONSECUTIVE_SUCCESS_STEPS + assert env_cfg.terminations.time_out is None + assert env_cfg.terminations.gear_dropped is None + assert env_cfg.terminations.gear_orientation_exceeded is None + + randomized_env = GearAssemblyEnvironment().build(GearAssemblyEnvironmentCfg(mode="randomized")) + assert randomized_env.task.observation_cfg.policy.concatenate_terms is True + assert randomized_env.task.observation_cfg.policy.gripper_pos is None + assert randomized_env.task.termination_cfg.time_out is not None + assert randomized_env.task.termination_cfg.gear_dropped is not None + assert randomized_env.task.termination_cfg.gear_orientation_exceeded is not None + return True + + +def test_gear_assembly_scene_and_newton_cfg(): + assert run_simulation_app_function(_test_gear_assembly_scene_and_newton_cfg, headless=True) + + +def _test_gear_assembly_newton_gears_settle(simulation_app): # noqa: C901 + import torch + + import warp as wp + from isaaclab.sim.utils import get_current_stage + from isaaclab.utils.math import quat_apply + from isaaclab_newton.physics.newton_manager import NewtonManager + from pxr import Usd, UsdGeom, UsdPhysics, UsdShade + + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder + from isaaclab_arena.environments.arena_env_builder_cfg import ArenaEnvBuilderCfg + from isaaclab_arena.tasks.gear_assembly.assets import ( + GEAR_GREEN_DIFFUSE_COLOR, + GEAR_GREEN_VISUAL_MATERIAL_PATH, + MAPLE_TABLE_TOP_COLLISION_COLOR, + NEWTON_GEAR_BASE_MESH_APPROXIMATION, + NEWTON_GEAR_MESH_APPROXIMATION, + ) + from isaaclab_arena.tasks.gear_assembly.specs import ( + GEAR_TABLETOP_PARKING_POSITIONS, + GEAR_TABLETOP_PARKING_Z, + MAPLE_TABLE_TOP_COLLISION_POSE, + MAPLE_TABLE_TOP_COLLISION_SIZE, + MAPLE_TABLE_TOP_COLLISION_THICKNESS, + ) + from isaaclab_arena_environments.gear_assembly_environment import ( + GearAssemblyEnvironment, + GearAssemblyEnvironmentCfg, + ) + + arena_env = GearAssemblyEnvironment().build( + GearAssemblyEnvironmentCfg(enable_cameras=False, embodiment="droid_rel_joint_pos") + ) + arena_env.name = "gear_assembly_newton_settle_regression" + builder = ArenaEnvBuilder(arena_env, ArenaEnvBuilderCfg(num_envs=1)) + env_cfg, env_kwargs = builder.compose_manager_cfg() + + env_cfg.scene.robot.init_state.pos = (10.0, 10.0, 0.0) + env_cfg.episode_length_s = 30.0 + for name in ( + "init_franka_arm_pose", + "randomize_franka_joint_state", + "set_robot_to_grasp_pose", + ): + if hasattr(env_cfg.events, name): + setattr(env_cfg.events, name, None) + for name in ("time_out", "gear_dropped", "gear_orientation_exceeded"): + if hasattr(env_cfg.terminations, name): + setattr(env_cfg.terminations, name, None) + + env = builder.make_registered(env_cfg, env_kwargs) + uenv = env.unwrapped + env.reset() + device = uenv.device + env_id = torch.tensor([0], device=device, dtype=torch.long) + zero_velocity = torch.zeros((1, 6), device=device) + far_pose = torch.tensor([[8.0, 8.0, 0.3, 1.0, 0.0, 0.0, 0.0]], device=device) + gear_names = ("factory_gear_small", "factory_gear_medium", "factory_gear_large") + action = torch.zeros(env.action_space.shape, dtype=torch.float32, device=device) + + table_center = torch.tensor(MAPLE_TABLE_TOP_COLLISION_POSE.position_xyz, device=device) + table_size = torch.tensor(MAPLE_TABLE_TOP_COLLISION_SIZE, device=device) + step_dt = uenv.step_dt + stage = get_current_stage() + bbox_cache = UsdGeom.BBoxCache(0, [UsdGeom.Tokens.default_], useExtentsHint=True) + collision_local_corners = {} + + asset_stage_paths = { + "factory_gear_base": ( + "/World/envs/env_0/FactoryGearBase", + "/World/envs/env_0/FactoryGearBase/factory_gear_base", + NEWTON_GEAR_BASE_MESH_APPROXIMATION, + ), + "factory_gear_small": ( + "/World/envs/env_0/FactoryGearSmall", + "/World/envs/env_0/FactoryGearSmall/factory_gear_small", + NEWTON_GEAR_MESH_APPROXIMATION, + ), + "factory_gear_medium": ( + "/World/envs/env_0/FactoryGearMedium", + "/World/envs/env_0/FactoryGearMedium/factory_gear_medium", + NEWTON_GEAR_MESH_APPROXIMATION, + ), + "factory_gear_large": ( + "/World/envs/env_0/FactoryGearLarge", + "/World/envs/env_0/FactoryGearLarge/factory_gear_large", + NEWTON_GEAR_MESH_APPROXIMATION, + ), + } + + def _box_corners(box) -> list[list[float]]: + box_min = box.GetMin() + box_max = box.GetMax() + return [ + [x, y, z] + for x in (box_min[0], box_max[0]) + for y in (box_min[1], box_max[1]) + for z in (box_min[2], box_max[2]) + ] + + def _assert_newton_collision_meshes( + asset_name: str, prim_path: str, rigid_prim_path: str, approximation: str + ) -> None: + root = stage.GetPrimAtPath(prim_path) + rigid_root = stage.GetPrimAtPath(rigid_prim_path) + assert root.IsValid(), f"{asset_name} prim is missing from the stage" + assert rigid_root.IsValid(), f"{asset_name} rigid-body prim is missing from the stage" + collision_meshes = [ + prim + for prim in Usd.PrimRange(root) + if prim.GetTypeName() == "Mesh" and "/collisions" in str(prim.GetPath()) + ] + assert collision_meshes, f"{asset_name} has no concrete collision mesh leaves" + local_corners = [] + for prim in collision_meshes: + assert prim.HasAPI(UsdPhysics.CollisionAPI), f"{prim.GetPath()} is missing CollisionAPI" + assert prim.HasAPI(UsdPhysics.MeshCollisionAPI), f"{prim.GetPath()} is missing MeshCollisionAPI" + authored_approximation = UsdPhysics.MeshCollisionAPI(prim).GetApproximationAttr().Get() + assert authored_approximation == approximation + local_corners.extend(_box_corners(bbox_cache.ComputeRelativeBound(prim, rigid_root).ComputeAlignedBox())) + collision_local_corners[asset_name] = torch.tensor(local_corners, device=device, dtype=torch.float32) + + def _world_collision_z_bounds(asset_name: str, root_pose: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + corners = collision_local_corners[asset_name] + world_corners = root_pose[:3] + quat_apply(root_pose[3:7].repeat(corners.shape[0], 1), corners) + return world_corners[:, 2].min(), world_corners[:, 2].max() + + def _shape_indices(labels: list[str], asset_fragment: str) -> list[int]: + matches = [ + index + for index, label in enumerate(labels) + if asset_fragment in label + and "/collisions/" in label + and "visuals" not in label + and not label.endswith("_visual") + ] + assert matches, f"Expected Newton collision shapes for {asset_fragment}" + return matches + + def _assert_newton_contact_pairs() -> None: + model = NewtonManager._model + labels = list(model.shape_label) + pairs = wp.to_torch(model.shape_contact_pairs).cpu().to(torch.long) + pair_set = {tuple(sorted(pair.tolist())) for pair in pairs} + + base_shapes = _shape_indices(labels, "FactoryGearBase") + gear_shapes = { + "factory_gear_small": _shape_indices(labels, "FactoryGearSmall"), + "factory_gear_medium": _shape_indices(labels, "FactoryGearMedium"), + "factory_gear_large": _shape_indices(labels, "FactoryGearLarge"), + } + for asset_name, asset_shapes in gear_shapes.items(): + for gear_shape in asset_shapes: + assert any( + tuple(sorted((base_shape, gear_shape))) in pair_set for base_shape in base_shapes + ), f"{asset_name} shape {gear_shape} is not paired with gear base" + + table_top_shapes = [index for index, label in enumerate(labels) if "/maple_table_top_collision/" in label] + assert table_top_shapes, "Newton model is missing the finite maple tabletop collision shape" + table_top_shape = table_top_shapes[0] + for asset_name, asset_shapes in { + "factory_gear_base": base_shapes, + **gear_shapes, + }.items(): + for shape in asset_shapes: + assert ( + tuple(sorted((table_top_shape, shape))) in pair_set + ), f"{asset_name} shape {shape} is not paired with table top" + + def _assert_finite_tabletop_proxy() -> None: + proxy_mesh = stage.GetPrimAtPath("/World/envs/env_0/maple_table_top_collision/geometry/mesh") + assert proxy_mesh.IsValid(), "Finite tabletop proxy mesh is missing from the stage" + assert proxy_mesh.GetTypeName() == "Cube" + proxy_box = bbox_cache.ComputeWorldBound(proxy_mesh).ComputeAlignedBox() + proxy_extent = proxy_box.GetSize() + assert abs(proxy_extent[0] - MAPLE_TABLE_TOP_COLLISION_SIZE[0]) < 1e-4 + assert abs(proxy_extent[1] - MAPLE_TABLE_TOP_COLLISION_SIZE[1]) < 1e-4 + assert abs(proxy_extent[2] - MAPLE_TABLE_TOP_COLLISION_THICKNESS) < 1e-4 + display_color = UsdGeom.PrimvarsAPI(proxy_mesh).GetPrimvar("displayColor").Get() + assert display_color is not None + authored_color = tuple(display_color[0]) + assert all( + abs(actual - expected) < 1e-6 for actual, expected in zip(authored_color, MAPLE_TABLE_TOP_COLLISION_COLOR) + ) + + def _assert_maple_table_top_newton_material() -> None: + table_top = stage.GetPrimAtPath("/World/envs/env_0/maple_table_robolab/table/table_01/top") + assert table_top.IsValid(), "Maple table top is missing from the stage" + bound_material, _ = UsdShade.MaterialBindingAPI(table_top).ComputeBoundMaterial() + assert ( + bound_material and bound_material.GetPrim().IsValid() + ), "Maple table top has no bound Newton-readable material" + material_path = str(bound_material.GetPrim().GetPath()) + assert material_path.endswith("/Looks/newton_maple_top") + shader_prim = stage.GetPrimAtPath(f"{material_path}/OmniPBRShader") + assert shader_prim.IsValid(), "Maple table top OmniPBR shader is missing" + assert str(shader_prim.GetAttribute("info:mdl:sourceAsset:subIdentifier").Get()) == "OmniPBR" + authored_color = tuple(shader_prim.GetAttribute("inputs:diffuse_color_constant").Get()) + assert all( + abs(actual - expected) < 1e-6 for actual, expected in zip(authored_color, MAPLE_TABLE_TOP_COLLISION_COLOR) + ) + + def _assert_green_visual_material(prim_path: str) -> None: + root = stage.GetPrimAtPath(prim_path) + bound_material, _ = UsdShade.MaterialBindingAPI(root).ComputeBoundMaterial() + assert bound_material and bound_material.GetPrim().IsValid(), f"{prim_path} has no bound visual material" + assert str(bound_material.GetPrim().GetPath()) == f"{prim_path}/{GEAR_GREEN_VISUAL_MATERIAL_PATH}" + shader_prim = stage.GetPrimAtPath(f"{prim_path}/{GEAR_GREEN_VISUAL_MATERIAL_PATH}/Shader") + assert shader_prim.IsValid(), f"{prim_path} green preview shader is missing" + authored_color = tuple(shader_prim.GetAttribute("inputs:diffuseColor").Get()) + assert all(abs(actual - expected) < 1e-6 for actual, expected in zip(authored_color, GEAR_GREEN_DIFFUSE_COLOR)) + + def _park_except(active_name: str, include_base: bool) -> None: + names = list(gear_names) + if include_base: + names.append("factory_gear_base") + for asset_name in names: + if asset_name == active_name: + continue + asset = uenv.scene[asset_name] + asset.write_root_pose_to_sim_index(root_pose=far_pose, env_ids=env_id) + asset.write_root_velocity_to_sim_index(root_velocity=zero_velocity, env_ids=env_id) + + def _root_motion( + asset_name: str, steps: int, sample_from_step: int + ) -> tuple[torch.Tensor, float, float, float, float, float]: + asset = uenv.scene[asset_name] + recent_poses = [] + for step in range(steps + 1): + root_pose = asset.data.root_link_pose_w.torch[0] + root_velocity = asset.data.root_com_vel_w.torch[0] + assert torch.isfinite(root_pose).all(), f"{asset_name} produced non-finite root pose at step {step}" + assert torch.isfinite(root_velocity).all(), f"{asset_name} produced non-finite root velocity at step {step}" + if step >= sample_from_step: + recent_poses.append(root_pose.detach().clone()) + env.step(action) + recent_poses = torch.stack(recent_poses) + recent_positions = recent_poses[:, :3] + linear_speeds, angular_speeds = _finite_difference_speeds(recent_poses) + excursion = torch.linalg.norm(recent_positions - recent_positions.mean(dim=0), dim=1).max().item() + return ( + asset.data.root_link_pose_w.torch[0].detach().clone(), + excursion, + linear_speeds.max().item(), + angular_speeds.max().item(), + linear_speeds[-1].item(), + angular_speeds[-1].item(), + ) + + def _finite_difference_speeds( + poses: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + position_delta = poses[1:, :3] - poses[:-1, :3] + linear_speeds = torch.linalg.norm(position_delta, dim=1) / step_dt + quat_dot = torch.abs(torch.sum(poses[1:, 3:7] * poses[:-1, 3:7], dim=1)) + angular_speeds = 2.0 * torch.acos(torch.clamp(quat_dot, max=1.0)) / step_dt + return linear_speeds, angular_speeds + + def _recent_pose_metrics( + poses: list[torch.Tensor], + ) -> tuple[float, float, float, float, float]: + recent_poses = torch.stack(poses) + linear_speeds, angular_speeds = _finite_difference_speeds(recent_poses) + excursion = ( + torch.linalg.norm( + recent_poses[:, :3] - recent_poses[:, :3].mean(dim=0), + dim=1, + ) + .max() + .item() + ) + return ( + excursion, + linear_speeds.max().item(), + angular_speeds.max().item(), + linear_speeds[-1].item(), + angular_speeds[-1].item(), + ) + + for asset_name, ( + prim_path, + rigid_prim_path, + approximation, + ) in asset_stage_paths.items(): + _assert_newton_collision_meshes(asset_name, prim_path, rigid_prim_path, approximation) + if asset_name in gear_names: + _assert_green_visual_material(prim_path) + _assert_finite_tabletop_proxy() + _assert_maple_table_top_newton_material() + _assert_newton_contact_pairs() + + for gear_name in gear_names: + env.reset() + _park_except(gear_name, include_base=True) + gear = uenv.scene[gear_name] + table_settle_pose = torch.tensor( + [[ + table_center[0].item(), + table_center[1].item(), + GEAR_TABLETOP_PARKING_Z, + 1.0, + 0.0, + 0.0, + 0.0, + ]], + device=device, + ) + gear.write_root_pose_to_sim_index(root_pose=table_settle_pose, env_ids=env_id) + gear.write_root_velocity_to_sim_index(root_velocity=zero_velocity, env_ids=env_id) + uenv.sim.forward() + + ( + table_state, + table_excursion, + table_max_linear, + table_max_angular, + table_final_linear, + table_final_angular, + ) = _root_motion(gear_name, steps=600, sample_from_step=540) + table_xy_error = torch.abs(table_state[:2] - table_center[:2]) + table_collision_min_z, _ = _world_collision_z_bounds(gear_name, table_state) + assert (table_xy_error <= table_size / 2).all(), f"{gear_name} slid off the tabletop: {table_state[:3]}" + assert ( + abs(table_collision_min_z.item() - table_center[2].item()) < 0.006 + ), f"{gear_name} collision mesh did not settle on the tabletop: {table_collision_min_z}" + assert table_excursion < 0.003, f"{gear_name} did not settle on the tabletop; max excursion={table_excursion}" + assert table_max_linear < 0.08, f"{gear_name} tabletop linear velocity is high" + assert table_max_angular < 0.2, f"{gear_name} tabletop angular velocity is high" + assert table_final_linear < 0.02, f"{gear_name} final tabletop linear velocity is high" + assert table_final_angular < 0.2, f"{gear_name} final tabletop angular velocity is high" + + for gear_key in ("gear_small", "gear_medium", "gear_large"): + gear_type_cfg = uenv.event_manager.get_term_cfg("randomize_gear_type") + gear_type_cfg.params["gear_types"] = [gear_key] + uenv.event_manager.set_term_cfg("randomize_gear_type", gear_type_cfg) + env.reset() + + all_positions = {asset_name: [] for asset_name in gear_names} + recent_poses = {asset_name: [] for asset_name in gear_names} + for step in range(601): + for asset_name in gear_names: + asset = uenv.scene[asset_name] + root_pose = asset.data.root_link_pose_w.torch[0] + root_velocity = asset.data.root_com_vel_w.torch[0] + assert torch.isfinite(root_pose).all(), f"{asset_name} reset pose became non-finite at step {step}" + assert torch.isfinite( + root_velocity + ).all(), f"{asset_name} reset velocity became non-finite at step {step}" + all_positions[asset_name].append(root_pose[:3].detach().clone()) + if step >= 540: + recent_poses[asset_name].append(root_pose.detach().clone()) + env.step(action) + + base_pose_w = uenv.scene["factory_gear_base"].data.root_link_pose_w.torch[0] + base_velocity_w = uenv.scene["factory_gear_base"].data.root_com_vel_w.torch[0] + base_collision_min_z, _ = _world_collision_z_bounds("factory_gear_base", base_pose_w) + assert torch.isfinite(base_pose_w).all(), "factory_gear_base reset pose became non-finite" + assert torch.isfinite(base_velocity_w).all(), "factory_gear_base reset velocity became non-finite" + assert (torch.abs(base_pose_w[:2] - table_center[:2]) <= table_size / 2).all() + assert abs(base_collision_min_z.item() - table_center[2].item()) < 0.005 + assert torch.linalg.norm(base_velocity_w).item() < 1e-6 + for asset_name in gear_names: + asset = uenv.scene[asset_name] + root_pose = asset.data.root_link_pose_w.torch[0] + trajectory = torch.stack(all_positions[asset_name]) + ( + excursion, + max_linear_speed, + max_angular_speed, + final_linear_speed, + final_angular_speed, + ) = _recent_pose_metrics(recent_poses[asset_name]) + xy_error = torch.abs(root_pose[:2] - table_center[:2]) + trajectory_xy_error = torch.abs(trajectory[:, :2] - table_center[:2]) + collision_min_z, collision_max_z = _world_collision_z_bounds(asset_name, root_pose) + + assert (xy_error <= table_size / 2).all(), f"{asset_name} left the tabletop in reset: {root_pose[:3]}" + assert ( + trajectory_xy_error <= table_size / 2 + ).all(), f"{asset_name} left the tabletop during reset: {trajectory[-1]}" + assert collision_min_z > table_center[2] - 0.006, f"{asset_name} fell through the tabletop: {root_pose[:3]}" + assert max_linear_speed < 0.15, f"{asset_name} reset linear velocity is high" + assert max_angular_speed < 5.0, f"{asset_name} reset angular velocity is high" + assert final_linear_speed < 0.1, f"{asset_name} final reset linear velocity is high" + assert final_angular_speed < 2.0, f"{asset_name} final reset angular velocity is high" + assert ( + collision_max_z < table_center[2] + 0.11 + ), f"{asset_name} is floating above the tabletop: {root_pose[:3]}" + assert excursion < 0.005, f"{asset_name} reset excursion is high: {excursion}" + parking_key = asset_name.removeprefix("factory_") + expected_position = torch.tensor(GEAR_TABLETOP_PARKING_POSITIONS[parking_key], device=device) + assert torch.linalg.norm(root_pose[:2] - expected_position[:2]).item() < 0.03 + + env.close() + return True + + +def test_gear_assembly_newton_gears_settle(): + assert run_simulation_app_function(_test_gear_assembly_newton_gears_settle, headless=True) + + +def _test_gear_assembly_newton_success_termination(simulation_app): + import torch + + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder + from isaaclab_arena.environments.arena_env_builder_cfg import ArenaEnvBuilderCfg + from isaaclab_arena.tasks.gear_assembly.specs import ( + GEAR_ASSEMBLED_CONSECUTIVE_SUCCESS_STEPS, + GEAR_ASSEMBLED_ROOT_Z_ABOVE_BASE, + ) + from isaaclab_arena_environments.gear_assembly_environment import ( + GearAssemblyEnvironment, + GearAssemblyEnvironmentCfg, + ) + + arena_env = GearAssemblyEnvironment().build( + GearAssemblyEnvironmentCfg(enable_cameras=False, embodiment="droid_rel_joint_pos") + ) + arena_env.name = "gear_assembly_newton_success_termination_regression" + builder = ArenaEnvBuilder(arena_env, ArenaEnvBuilderCfg(num_envs=1)) + env_cfg, env_kwargs = builder.compose_manager_cfg() + + env_cfg.scene.robot.init_state.pos = (10.0, 10.0, 0.0) + for name in ( + "init_franka_arm_pose", + "randomize_franka_joint_state", + "set_robot_to_grasp_pose", + ): + if hasattr(env_cfg.events, name): + setattr(env_cfg.events, name, None) + + env = builder.make_registered(env_cfg, env_kwargs) + uenv = env.unwrapped + assert uenv.termination_manager.active_terms == ["success"] + + action = torch.zeros(env.action_space.shape, dtype=torch.float32, device=uenv.device) + env_id = torch.tensor([0], device=uenv.device, dtype=torch.long) + zero_velocity = torch.zeros((1, 6), device=uenv.device) + gear_names = ("factory_gear_small", "factory_gear_medium", "factory_gear_large") + release_root_z_above_base = max(GEAR_ASSEMBLED_ROOT_Z_ABOVE_BASE.values()) + + for gear_key, active_gear_name in zip(("gear_small", "gear_medium", "gear_large"), gear_names): + gear_type_cfg = uenv.event_manager.get_term_cfg("randomize_gear_type") + gear_type_cfg.params["gear_types"] = [gear_key] + uenv.event_manager.set_term_cfg("randomize_gear_type", gear_type_cfg) + + _, _ = env.reset() + for step in range(180): + _, _, terminated, truncated, _ = env.step(action) + assert not terminated.item(), f"{gear_key} succeeded before reaching the base at step {step}" + assert not truncated.item() + assert not uenv.termination_manager.get_term("success").item() + + active_gear = uenv.scene[active_gear_name] + base_pose = uenv.scene["factory_gear_base"].data.root_link_pose_w.torch[0].detach().clone() + assembled_pose = base_pose.clone().unsqueeze(0) + assembled_pose[:, 2] += release_root_z_above_base + + active_gear.write_root_pose_to_sim_index(root_pose=assembled_pose, env_ids=env_id) + active_gear.write_root_velocity_to_sim_index(root_velocity=zero_velocity, env_ids=env_id) + uenv.sim.forward() + + immediate_termination = uenv.termination_manager.compute() + assert immediate_termination.tolist() == [False] + assert uenv.termination_manager.get_term("success").tolist() == [False] + + terminated_on_step = None + for step in range(GEAR_ASSEMBLED_CONSECUTIVE_SUCCESS_STEPS + 60): + _, _, terminated, truncated, _ = env.step(action) + assert not truncated.item() + if terminated.item(): + terminated_on_step = step + break + assert terminated_on_step is not None, f"{gear_key} did not terminate after settling on the base" + assert terminated_on_step >= GEAR_ASSEMBLED_CONSECUTIVE_SUCCESS_STEPS - 2 + + env.close() + return True + + +def test_gear_assembly_newton_success_termination(): + assert run_simulation_app_function(_test_gear_assembly_newton_success_termination, headless=True) diff --git a/isaaclab_arena/utils/usd/newton.py b/isaaclab_arena/utils/usd/newton.py new file mode 100644 index 0000000000..71fc2df1f4 --- /dev/null +++ b/isaaclab_arena/utils/usd/newton.py @@ -0,0 +1,85 @@ +# Copyright (c) 2026, 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 + +"""USD compatibility helpers for Newton/MuJoCo-Warp.""" + +from __future__ import annotations + +import shutil +from pathlib import Path + + +def ensure_newton_valid_rigid_body_inertias_usd( + usd_path: str, + min_mass: float = 0.02, + min_diagonal_inertia: float = 1.0e-5, +) -> str: + """Return a cached USD copy with positive mass and inertia on rigid bodies.""" + from isaaclab.utils.assets import retrieve_file_path + + source = Path(retrieve_file_path(usd_path, force_download=False)) + assert source.is_file(), f"USD path must resolve to a local file: {usd_path}" + + target = source.with_name(f"{source.stem}_newton_inertia{source.suffix}") + if target.exists() and target.stat().st_mtime >= source.stat().st_mtime and _has_valid_rigid_body_inertias(target): + return str(target) + + shutil.copy2(source, target) + _author_minimum_rigid_body_inertias(target, min_mass=min_mass, min_diagonal_inertia=min_diagonal_inertia) + return str(target) + + +def _has_valid_rigid_body_inertias(usd_path: Path) -> bool: + from pxr import Usd, UsdPhysics + + stage = Usd.Stage.Open(str(usd_path)) + assert stage is not None, f"Could not open USD: {usd_path}" + return all(not (prim.HasAPI(UsdPhysics.RigidBodyAPI) and _needs_minimum_inertia(prim)) for prim in stage.Traverse()) + + +def _author_minimum_rigid_body_inertias( + usd_path: Path, + min_mass: float, + min_diagonal_inertia: float, +) -> None: + from pxr import Gf, Usd, UsdPhysics + + stage = Usd.Stage.Open(str(usd_path)) + assert stage is not None, f"Could not open USD: {usd_path}" + + diagonal_inertia = Gf.Vec3f(min_diagonal_inertia, min_diagonal_inertia, min_diagonal_inertia) + for prim in stage.Traverse(): + if not prim.HasAPI(UsdPhysics.RigidBodyAPI) or not _needs_minimum_inertia(prim): + continue + + mass_api = UsdPhysics.MassAPI(prim) + if not prim.HasAPI(UsdPhysics.MassAPI): + mass_api = UsdPhysics.MassAPI.Apply(prim) + + if _invalid_positive_value(mass_api.GetMassAttr().Get()): + mass_api.CreateMassAttr().Set(min_mass) + if _invalid_diagonal_inertia(mass_api.GetDiagonalInertiaAttr().Get()): + mass_api.CreateDiagonalInertiaAttr().Set(diagonal_inertia) + + stage.GetRootLayer().Save() + + +def _needs_minimum_inertia(prim) -> bool: + from pxr import UsdPhysics + + if not prim.HasAPI(UsdPhysics.MassAPI): + return True + mass_api = UsdPhysics.MassAPI(prim) + return _invalid_positive_value(mass_api.GetMassAttr().Get()) or _invalid_diagonal_inertia( + mass_api.GetDiagonalInertiaAttr().Get() + ) + + +def _invalid_positive_value(value) -> bool: + return value is None or float(value) <= 0.0 + + +def _invalid_diagonal_inertia(value) -> bool: + return value is None or any(float(component) <= 0.0 for component in value) diff --git a/isaaclab_arena_environments/__init__.py b/isaaclab_arena_environments/__init__.py index 69104f473c..29a2b9c11e 100644 --- a/isaaclab_arena_environments/__init__.py +++ b/isaaclab_arena_environments/__init__.py @@ -6,10 +6,17 @@ import importlib import pkgutil -import isaaclab_arena_environments - _NON_ENVIRONMENT_MODULES = {"cli", "example_environment_base"} +_ENVIRONMENTS_REGISTERED = False + + +def register_environments() -> None: + """Import all first-party environment modules so their decorators register them.""" + global _ENVIRONMENTS_REGISTERED + if _ENVIRONMENTS_REGISTERED: + return -for _importer, _modname, _ispkg in pkgutil.iter_modules(isaaclab_arena_environments.__path__): - if not _ispkg and _modname not in _NON_ENVIRONMENT_MODULES: - importlib.import_module(f"isaaclab_arena_environments.{_modname}") + for _importer, modname, ispkg in pkgutil.iter_modules(__path__): + if not ispkg and modname not in _NON_ENVIRONMENT_MODULES: + importlib.import_module(f"{__name__}.{modname}") + _ENVIRONMENTS_REGISTERED = True diff --git a/isaaclab_arena_environments/cli.py b/isaaclab_arena_environments/cli.py index 387601b248..1de214f9c6 100644 --- a/isaaclab_arena_environments/cli.py +++ b/isaaclab_arena_environments/cli.py @@ -25,11 +25,12 @@ def ensure_environments_registered(): """Trigger registration of all environments in the ``isaaclab_arena_environments`` package. - Importing the package fires the ``@register_environment`` decorator on each - environment module, which handles registration. The import is cached by - Python, so subsequent calls are free. + Environment modules are imported explicitly here so importing this CLI module + does not load USD/pxr before SimulationApp starts. """ - import isaaclab_arena_environments # noqa: F401 + import isaaclab_arena_environments + + isaaclab_arena_environments.register_environments() # Legacy argparse compatibility diff --git a/isaaclab_arena_environments/gear_assembly_environment.py b/isaaclab_arena_environments/gear_assembly_environment.py new file mode 100644 index 0000000000..6647f17aa3 --- /dev/null +++ b/isaaclab_arena_environments/gear_assembly_environment.py @@ -0,0 +1,172 @@ +# Copyright (c) 2026, 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 + +"""Gear Assembly scene using Arena's existing Droid embodiment.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from isaaclab_arena.assets.register import register_environment +from isaaclab_arena.environments.arena_environment_factory import ArenaEnvironmentCfg, ArenaEnvironmentFactory +from isaaclab_arena.tasks.gear_assembly.specs import ( + DROID_GEAR_ASSEMBLY_EMBODIMENTS, + GEAR_TABLETOP_ORIENTATION_XYZW, + GEAR_TABLETOP_PARKING_POSITIONS, + MAPLE_TABLE_POSE, + MAPLE_TABLE_TOP_COLLISION_POSE, + gear_pose_for_mode, + get_droid_robot_spec, +) + +if TYPE_CHECKING: + from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment + from isaaclab_arena.environments.isaaclab_arena_manager_based_env_cfg import IsaacLabArenaManagerBasedRLEnvCfg + from isaaclab_arena.tasks.gear_assembly.task import GearAssemblyTask + + +@dataclass +class GearAssemblyEnvironmentCfg(ArenaEnvironmentCfg): + """Configure the Arena Droid Gear Assembly environment.""" + + embodiment: str = "droid_abs_joint_pos" + mode: str = "play" + physics_backend: str = "newton" + + def __post_init__(self) -> None: + assert ( + self.embodiment in DROID_GEAR_ASSEMBLY_EMBODIMENTS + ), f"Gear Assembly is Droid-only; got embodiment={self.embodiment!r}" + assert self.mode in {"play", "randomized"}, f"Unsupported Gear Assembly mode: {self.mode!r}" + assert self.physics_backend in { + "newton", + "physx", + }, f"Unsupported Gear Assembly physics backend: {self.physics_backend!r}" + + +@register_environment +class GearAssemblyEnvironment(ArenaEnvironmentFactory[GearAssemblyEnvironmentCfg]): + """Arena Gear Assembly factory using the existing Droid robot.""" + + name = "gear_assembly" + _legacy_argparse_cfg_type = GearAssemblyEnvironmentCfg + + def build(self, cfg: GearAssemblyEnvironmentCfg) -> IsaacLabArenaEnvironment: + import isaaclab.sim as sim_utils + + from isaaclab_arena.assets.object_base import ObjectType + from isaaclab_arena.assets.object_library import DomeLight + from isaaclab_arena.assets.object_reference import ObjectReference + from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment + from isaaclab_arena.relations.relations import IsAnchor + from isaaclab_arena.scene.scene import Scene + from isaaclab_arena.tasks.gear_assembly.assets import ( + make_factory_gear_base, + make_factory_gear_large, + make_factory_gear_medium, + make_factory_gear_small, + make_ground, + make_maple_table_top_collision, + spawn_newton_maple_table_usd, + ) + from isaaclab_arena.tasks.gear_assembly.task import GearAssemblyTask + + embodiment = self.asset_registry.get_asset_by_name(cfg.embodiment)(enable_cameras=cfg.enable_cameras) + if cfg.physics_backend == "newton": + from isaaclab_arena.utils.usd.newton import ensure_newton_valid_rigid_body_inertias_usd + + embodiment.scene_config.robot.spawn.usd_path = ensure_newton_valid_rigid_body_inertias_usd( + embodiment.scene_config.robot.spawn.usd_path + ) + embodiment.observation_config = None + robot_spec = get_droid_robot_spec() + gear_pose = gear_pose_for_mode(cfg.mode) + newton_mesh_collisions = cfg.physics_backend == "newton" + maple_table = self.asset_registry.get_asset_by_name("maple_table_robolab")() + maple_table.set_initial_pose(MAPLE_TABLE_POSE) + if newton_mesh_collisions: + maple_table.object_cfg.spawn.func = spawn_newton_maple_table_usd + table_reference = ObjectReference( + name="table", + prim_path="{ENV_REGEX_NS}/maple_table_robolab/table", + parent_asset=maple_table, + object_type=ObjectType.RIGID, + ) + table_reference.add_relation(IsAnchor()) + + assets = [make_ground(), maple_table] + if newton_mesh_collisions: + assets.append(make_maple_table_top_collision(MAPLE_TABLE_TOP_COLLISION_POSE)) + assets += [ + make_factory_gear_base(gear_pose, newton_mesh_collisions=newton_mesh_collisions), + make_factory_gear_small(gear_pose, newton_mesh_collisions=newton_mesh_collisions), + make_factory_gear_medium(gear_pose, newton_mesh_collisions=newton_mesh_collisions), + make_factory_gear_large(gear_pose, newton_mesh_collisions=newton_mesh_collisions), + table_reference, + DomeLight( + instance_name="light", + prim_path="/World/light", + spawner_cfg=sim_utils.DomeLightCfg(color=(0.75, 0.75, 0.75), intensity=2500.0), + ), + ] + + task = GearAssemblyTask(robot_spec=robot_spec, mode=cfg.mode) + return IsaacLabArenaEnvironment( + name=self.name, + embodiment=embodiment, + scene=Scene(assets=assets), + task=task, + env_cfg_callback=_make_env_cfg_callback(cfg, task), + ) + + +def _make_env_cfg_callback(cfg: GearAssemblyEnvironmentCfg, task: GearAssemblyTask): + def gear_assembly_env_cfg_callback( + env_cfg: IsaacLabArenaManagerBasedRLEnvCfg, + ) -> IsaacLabArenaManagerBasedRLEnvCfg: + from isaaclab_physx.physics import PhysxCfg + + from isaaclab_arena.environments.isaaclab_arena_manager_based_env_cfg import ArenaPhysicsCfg + + env_cfg.episode_length_s = 6.66 + env_cfg.viewer.eye = (1.6, -1.2, 1.0) + env_cfg.viewer.lookat = (0.55, 0.05, 0.08) + env_cfg.decimation = 4 + env_cfg.sim.render_interval = 4 + env_cfg.sim.dt = 1.0 / 120.0 + + if cfg.physics_backend == "newton": + env_cfg.sim.physics = ArenaPhysicsCfg().newton + env_cfg.sim.physics.default_shape_cfg.gap = 0.0 + env_cfg.scene.replicate_physics = True + elif cfg.physics_backend == "physx": + env_cfg.sim.physics = PhysxCfg( + gpu_collision_stack_size=2**30, + gpu_max_rigid_contact_count=2**23, + gpu_max_rigid_patch_count=2**23, + ) + env_cfg.scene.replicate_physics = False + else: + raise ValueError(f"Unsupported Gear Assembly physics backend: {cfg.physics_backend}") + + for attr_name, value in task.runtime_env_attrs().items(): + setattr(env_cfg, attr_name, value) + if cfg.physics_backend == "newton": + env_cfg.events.randomize_gears_and_base_pose.params["selected_parking_positions"] = ( + GEAR_TABLETOP_PARKING_POSITIONS + ) + env_cfg.events.randomize_gears_and_base_pose.params["selected_orientation_xyzw"] = ( + GEAR_TABLETOP_ORIENTATION_XYZW + ) + if cfg.mode == "play": + # Evaluation starts the selected gear on the table. The source failure terms are + # gripper-relative, so play mode terminates only on the assembled success term. + env_cfg.terminations.gear_dropped = None + env_cfg.terminations.gear_orientation_exceeded = None + env_cfg.terminations.time_out = None + return env_cfg + + return gear_assembly_env_cfg_callback