diff --git a/isaaclab_arena/affordances/placeable.py b/isaaclab_arena/affordances/placeable.py new file mode 100644 index 00000000..71aa4b94 --- /dev/null +++ b/isaaclab_arena/affordances/placeable.py @@ -0,0 +1,275 @@ +# Copyright (c) 2025, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import math +import torch +from typing import Literal + +import isaaclab.utils.math as math_utils +from isaaclab.assets import RigidObject +from isaaclab.envs.manager_based_env import ManagerBasedEnv +from isaaclab.managers import SceneEntityCfg + +from isaaclab_arena.affordances.affordance_base import AffordanceBase + + +class Placeable(AffordanceBase): + """Interface for placeable objects. + Placeable objects are objects that can be placed upright in a scene. + They are characterized by their upright axis and orientation threshold, which are used to determine if the object is placed upright. + - The upright axis is the axis in the object's local frame that is used to determine if the object is placed upright. + - The orientation threshold is the threshold for the angle between the upright axis and the world +Z direction, in radians. + """ + + def __init__( + self, upright_axis_name: Literal["x", "y", "z"] = "z", orientation_threshold: float = math.pi / 18.0, **kwargs + ): + super().__init__(**kwargs) + self.upright_axis_name = upright_axis_name + assert upright_axis_name in ["x", "y", "z"], "Upright axis must be one of x, y, or z" + self.orientation_threshold = orientation_threshold + + def is_placed_upright( + self, + env: ManagerBasedEnv, + asset_cfg: SceneEntityCfg | None = None, + orientation_threshold: float | torch.Tensor | None = None, + ) -> torch.Tensor: + """Returns a boolean tensor of whether the object is placeable. + + Args: + env: The environment instance. + asset_cfg: Asset configuration. If None, uses the object's name. + orientation_threshold: Orientation threshold in radians. If None, uses the object's orientation threshold. + If a tensor, must have shape [num_envs]. + + Returns: + A boolean tensor of whether the object is placed upright. Shape: [num_envs]. + """ + if asset_cfg is None: + asset_cfg = SceneEntityCfg(self.name) + # We allow for overriding the object-level threshold by passing an argument to this + # function explicitly. Otherwise we use the object-level threshold. + if orientation_threshold is None: + orientation_threshold = self.orientation_threshold + object_entity: RigidObject = env.scene[asset_cfg.name] + object_quat = object_entity.data.root_quat_w + + upright_axis_world = get_object_axis_in_world_frame(object_quat, self.upright_axis_name) + + world_up = torch.zeros_like(upright_axis_world) + world_up[..., 2] = 1.0 + + cos_angle = torch.sum(upright_axis_world * world_up, dim=-1).clamp(-1.0, 1.0) + angle_error = torch.acos(cos_angle) + + orientation_threshold_tensor = torch.as_tensor( + orientation_threshold, device=object_quat.device, dtype=object_quat.dtype + ) + + success = angle_error < orientation_threshold_tensor + + return success + + def place_upright( + self, + env: ManagerBasedEnv, + env_ids: torch.Tensor | None, + asset_cfg: SceneEntityCfg | None = None, + upright_percentage: float | torch.Tensor = 1.0, + ): + """Place the object upright (in all the environments). + + This function places the object upright by rotating its upright axis towards the world +Z direction. + The upright percentage is a value in [0, 1] that describes the absolute target angle of the object_upright_axis: + 0.0 results in a perpendicular orientation (90 degrees) in the plane spanned by the object_upright_axis and world +Z, + and 1.0 fully aligns the axis with +Z (0 degrees). Other target angle in (0, 1) rotates in the plane spanned by the object_upright_axis and world +Z as well. + + Args: + env: The environment instance. + env_ids: Environment indices to apply the operation to. If None, applies to all environments. + asset_cfg: Asset configuration. If None, uses the object's name. + upright_percentage: Target upright percentage(s) in [0, 1]. Can be: + - A scalar float applied to all specified environments + - A tensor with shape matching len(env_ids), with per-environment values + + Raises: + ValueError: If upright_percentage is a tensor with size not matching len(env_ids). + """ + if asset_cfg is None: + asset_cfg = SceneEntityCfg(self.name) + + set_normalized_object_pose( + env=env, + asset_cfg=asset_cfg, + upright_percentage=upright_percentage, + env_ids=env_ids, + upright_axis_name=self.upright_axis_name, + ) + + +def get_object_axis_in_world_frame(object_quat: torch.tensor, upright_axis_name: str) -> torch.tensor: + """Get the object axis in the world frame. + Args: + object_quat: The quaternion of the object. Shape: [num_envs, 4] (w, x, y, z). + upright_axis_name: The name of the object axis to get in the world frame. Can be "x", "y", or "z". + + Returns: + The object axis in the world frame. Shape: [num_envs, 3]. + """ + axis_index = {"x": 0, "y": 1, "z": 2}[upright_axis_name] + rotation_mats = math_utils.matrix_from_quat(object_quat) + object_axis_world = rotation_mats[:, :, axis_index] + return object_axis_world + + +def set_normalized_object_pose( + env: ManagerBasedEnv, + asset_cfg: SceneEntityCfg, + upright_percentage: float | torch.Tensor = 1.0, + env_ids: torch.Tensor | None = None, + upright_axis_name: str = "z", +) -> None: + """Rotate a rigid object towards the upright orientation directly through its root pose. + + Values of ``upright_percentage`` in [0, 1] describe the absolute target orientation of the object_upright_axis. + + Target orientation diagram (only rotates in the plane spanned by the object_upright_axis and world +Z): + + +Z (world up) + ↑ target upright axis (upright_percentage=1.0) + | + | + | + | + | + | / object_upright_axis + | / + | / + | / θ + ●------→ target upright axis (upright_percentage=0.0, 90° to +Z) + object (the direction with the smallest angle to the object_upright_axis) + + Target orientation: only rotates in the plane spanned by the object_upright_axis and world +Z. + - upright_percentage = 1.0: the target upright axis aligns with world +Z (0°) + - upright_percentage = 0.0: the target upright axis is perpendicular to world +Z (90°), and the direction with the smallest angle to the object_upright_axis + + Args: + env: The environment instance. + asset_cfg: Configuration for the rigid object asset. + upright_percentage: Target upright percentage(s) in [0, 1]. Can be: + - A scalar float applied to all specified environments + - A tensor with shape matching len(env_ids), with per-environment values + env_ids: Environment indices to apply the operation to. If None, applies to all environments. + upright_axis_name: Name of the object's local axis to align ("x", "y", or "z"). + + Raises: + ValueError: If upright_percentage is a tensor with size not matching len(env_ids). + """ + object_entity: RigidObject = env.scene[asset_cfg.name] + device = env.device + dtype = object_entity.data.root_quat_w.dtype + + if env_ids is not None: + env_ids = env_ids.to(env.device) + else: + env_ids = torch.arange(object_entity.data.root_quat_w.shape[0], device=env.device) + + # Validate upright_percentage shape if it's a tensor + if isinstance(upright_percentage, torch.Tensor): + num_envs = len(env_ids) + if upright_percentage.numel() != num_envs: + raise ValueError( + f"upright_percentage tensor must have {num_envs} elements to match env_ids, " + f"but got {upright_percentage.numel()} elements" + ) + + object_quat = object_entity.data.root_quat_w[env_ids] + object_pos = object_entity.data.root_pos_w[env_ids] + + target_quat = _compute_target_quaternions( + object_quat=object_quat, + upright_percentage=upright_percentage, + upright_axis_name=upright_axis_name, + ) + + pose_tensor = torch.cat([object_pos, target_quat], dim=-1) + object_entity.write_root_pose_to_sim(pose_tensor, env_ids=env_ids) + zero_velocity = torch.zeros((env_ids.numel(), 6), device=device, dtype=dtype) + object_entity.write_root_velocity_to_sim(zero_velocity, env_ids=env_ids) + + +def _compute_target_quaternions( + object_quat: torch.Tensor, + upright_percentage: float | torch.Tensor, + upright_axis_name: str, +) -> torch.Tensor: + """Compute the target quaternions for the object given the current orientation and the upright percentage. + + Args: + object_quat: Current quaternion orientations of the objects. Shape: [num_envs, 4] (w, x, y, z). + upright_percentage: Target upright percentage(s) in [0, 1]. Can be: + - A scalar float applied to all environments + - A tensor with num_envs elements (shape [num_envs] or [num_envs, 1]) + Value of 0.0 creates perpendicular orientation (90° to world +Z) in the plane spanned by the object_upright_axis and world +Z, + value of 1.0 fully aligns with world +Z (0°). + upright_axis_name: Name of the object's local axis to align ("x", "y", or "z"). + + Returns: + Target quaternions for the object. Shape: [num_envs, 4] (w, x, y, z). + """ + + upright_percentage_t = torch.as_tensor(upright_percentage, device=object_quat.device, dtype=object_quat.dtype) + if upright_percentage_t.dim() == 0: + upright_percentage_t = upright_percentage_t.expand(object_quat.shape[0]) + + current_axis = get_object_axis_in_world_frame(object_quat, upright_axis_name) + + world_up = torch.zeros_like(current_axis) + world_up[:, 2] = 1.0 + + horizontal = current_axis - torch.sum(current_axis * world_up, dim=-1, keepdim=True) * world_up + horizontal_norm = horizontal.norm(dim=-1, keepdim=True) + fallback_axis = torch.zeros_like(horizontal) + fallback_axis[:, 0] = 1.0 + + # if current_axis is parallel to world up, use the fallback axis to compute the rotation axis + needs_fallback = horizontal_norm.squeeze(-1) < 1e-6 + if needs_fallback.any(): + horizontal[needs_fallback] = fallback_axis[needs_fallback] + horizontal_norm[needs_fallback] = horizontal[needs_fallback].norm(dim=-1, keepdim=True) + horizontal_dir = horizontal / horizontal_norm.clamp(min=1e-6) + + target_angles = (1.0 - upright_percentage_t).view(-1, 1) * (math.pi / 2.0) + target_axis = torch.sin(target_angles) * horizontal_dir + torch.cos(target_angles) * world_up + target_axis = torch.nn.functional.normalize(target_axis, dim=-1) + + dot_product = torch.sum(current_axis * target_axis, dim=-1).clamp(-1.0, 1.0) + rotation_angle = torch.acos(dot_product) + + rotation_axis = torch.cross(current_axis, target_axis, dim=-1) + axis_norm = rotation_axis.norm(dim=-1, keepdim=True) + + # if current_axis is parallel to target_axis, use the fallback axis to compute the rotation axis + needs_fallback = axis_norm.squeeze(-1) < 1e-6 + if needs_fallback.any(): + current_axis_fb = current_axis[needs_fallback] + fallback_axis_fb = fallback_axis[needs_fallback] + near_parallel_x = torch.abs(current_axis_fb[:, 0]) > 0.9 + if near_parallel_x.any(): + fallback_axis_fb[near_parallel_x] = torch.tensor( + [0.0, 1.0, 0.0], device=object_quat.device, dtype=object_quat.dtype + ) + rotation_axis[needs_fallback] = torch.cross(current_axis_fb, fallback_axis_fb, dim=-1) + axis_norm[needs_fallback] = rotation_axis[needs_fallback].norm(dim=-1, keepdim=True) + + axis_norm = axis_norm.clamp(min=1e-6) + rotation_axis_unit = rotation_axis / axis_norm + + align_quat = math_utils.quat_from_angle_axis(rotation_angle, rotation_axis_unit) + + new_quat = math_utils.quat_mul(align_quat, object_quat) + new_quat = torch.nn.functional.normalize(new_quat, dim=-1) + return new_quat diff --git a/isaaclab_arena/assets/background_library.py b/isaaclab_arena/assets/background_library.py index 9ced4b37..1b080ab4 100644 --- a/isaaclab_arena/assets/background_library.py +++ b/isaaclab_arena/assets/background_library.py @@ -3,6 +3,8 @@ # # SPDX-License-Identifier: Apache-2.0 +from typing import Any + from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR from isaaclab_arena.assets.background import Background @@ -21,6 +23,8 @@ class LibraryBackground(Background): usd_path: str initial_pose: Pose | None = None object_min_z: float + spawn_cfg_addon: dict[str, Any] = {} + asset_cfg_addon: dict[str, Any] = {} def __init__(self, **kwargs): super().__init__( @@ -29,6 +33,8 @@ def __init__(self, **kwargs): usd_path=self.usd_path, initial_pose=self.initial_pose, object_min_z=self.object_min_z, + spawn_cfg_addon=self.spawn_cfg_addon, + asset_cfg_addon=self.asset_cfg_addon, **kwargs, ) diff --git a/isaaclab_arena/assets/object.py b/isaaclab_arena/assets/object.py index 30322177..69b7f6e6 100644 --- a/isaaclab_arena/assets/object.py +++ b/isaaclab_arena/assets/object.py @@ -3,6 +3,8 @@ # # SPDX-License-Identifier: Apache-2.0 +from typing import Any + from isaaclab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg from isaaclab.sim.spawners.from_files.from_files_cfg import UsdFileCfg @@ -27,6 +29,9 @@ def __init__( initial_pose: Pose | None = None, **kwargs, ): + # Pull out addons (and remove them from kwargs before passing to super) + spawn_cfg_addon: dict[str, Any] = kwargs.pop("spawn_cfg_addon", {}) or {} + asset_cfg_addon: dict[str, Any] = kwargs.pop("asset_cfg_addon", {}) or {} if object_type is not ObjectType.SPAWNER: assert usd_path is not None # Detect object type if not provided @@ -36,6 +41,8 @@ def __init__( self.usd_path = usd_path self.scale = scale self.initial_pose = initial_pose + self.spawn_cfg_addon = spawn_cfg_addon + self.asset_cfg_addon = asset_cfg_addon self.object_cfg = self._init_object_cfg() def set_initial_pose(self, pose: Pose) -> None: @@ -56,7 +63,9 @@ def _generate_rigid_cfg(self) -> RigidObjectCfg: usd_path=self.usd_path, scale=self.scale, activate_contact_sensors=True, + **self.spawn_cfg_addon, ), + **self.asset_cfg_addon, ) object_cfg = self._add_initial_pose_to_cfg(object_cfg) return object_cfg @@ -69,7 +78,9 @@ def _generate_articulation_cfg(self) -> ArticulationCfg: usd_path=self.usd_path, scale=self.scale, activate_contact_sensors=True, + **self.spawn_cfg_addon, ), + **self.asset_cfg_addon, actuators={}, ) object_cfg = self._add_initial_pose_to_cfg(object_cfg) @@ -82,7 +93,12 @@ def _generate_base_cfg(self) -> AssetBaseCfg: print("WARNING: Base object has lights, this may cause issues when using with multiple environments.") object_cfg = AssetBaseCfg( prim_path="{ENV_REGEX_NS}/" + self.name, - spawn=UsdFileCfg(usd_path=self.usd_path, scale=self.scale), + spawn=UsdFileCfg( + usd_path=self.usd_path, + scale=self.scale, + **self.spawn_cfg_addon, + ), + **self.asset_cfg_addon, ) object_cfg = self._add_initial_pose_to_cfg(object_cfg) return object_cfg @@ -92,6 +108,7 @@ def _generate_spawner_cfg(self) -> AssetBaseCfg: object_cfg = AssetBaseCfg( prim_path=self.prim_path, spawn=self.spawner_cfg, + **self.asset_cfg_addon, ) object_cfg = self._add_initial_pose_to_cfg(object_cfg) return object_cfg diff --git a/isaaclab_arena/assets/object_library.py b/isaaclab_arena/assets/object_library.py index 8fe360fd..033b65d4 100644 --- a/isaaclab_arena/assets/object_library.py +++ b/isaaclab_arena/assets/object_library.py @@ -3,11 +3,14 @@ # # SPDX-License-Identifier: Apache-2.0 +from typing import Any + import isaaclab.sim as sim_utils from isaaclab.sim.spawners.from_files.from_files_cfg import GroundPlaneCfg from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR from isaaclab_arena.affordances.openable import Openable +from isaaclab_arena.affordances.placeable import Placeable from isaaclab_arena.affordances.pressable import Pressable from isaaclab_arena.assets.object import Object from isaaclab_arena.assets.object_base import ObjectType @@ -26,6 +29,8 @@ class LibraryObject(Object): usd_path: str | None = None object_type: ObjectType = ObjectType.RIGID scale: tuple[float, float, float] = (1.0, 1.0, 1.0) + spawn_cfg_addon: dict[str, Any] = {} + asset_cfg_addon: dict[str, Any] = {} def __init__(self, prim_path: str | None = None, initial_pose: Pose | None = None, **kwargs): super().__init__( @@ -36,6 +41,8 @@ def __init__(self, prim_path: str | None = None, initial_pose: Pose | None = Non object_type=self.object_type, scale=self.scale, initial_pose=initial_pose, + spawn_cfg_addon=self.spawn_cfg_addon, + asset_cfg_addon=self.asset_cfg_addon, **kwargs, ) @@ -232,6 +239,42 @@ def __init__(self, prim_path: str | None = None, initial_pose: Pose | None = Non super().__init__(prim_path=prim_path, initial_pose=initial_pose) +@register_asset +class Mug(LibraryObject, Placeable): + """ + A mug. + """ + + name = "mug" + tags = ["object"] + usd_path = f"{ISAACLAB_NUCLEUS_DIR}/Objects/Mug/mug.usd" + object_type = ObjectType.RIGID + default_prim_path = "{ENV_REGEX_NS}/Mug" + scale = (1.0, 1.0, 1.0) + + # Placeable affordance parameters + upright_axis_name = "z" + orientation_threshold = 0.5 + + def __init__(self, prim_path: str | None = None, initial_pose: Pose | None = None): + super().__init__( + prim_path=prim_path, + initial_pose=initial_pose, + upright_axis_name=self.upright_axis_name, + orientation_threshold=self.orientation_threshold, + ) + RIGID_BODY_PROPS = sim_utils.RigidBodyPropertiesCfg( + solver_position_iteration_count=16, + solver_velocity_iteration_count=1, + max_angular_velocity=1000.0, + max_linear_velocity=1000.0, + max_depenetration_velocity=5.0, + disable_gravity=False, + ) + self.object_cfg.spawn.rigid_props = RIGID_BODY_PROPS + self.object_cfg.spawn.mass_props = sim_utils.MassPropertiesCfg(mass=0.25) + + @register_asset class GroundPlane(LibraryObject): """ diff --git a/isaaclab_arena/assets/retargeter_library.py b/isaaclab_arena/assets/retargeter_library.py index cfceaff6..b86a088c 100644 --- a/isaaclab_arena/assets/retargeter_library.py +++ b/isaaclab_arena/assets/retargeter_library.py @@ -83,3 +83,17 @@ def get_retargeter_cfg( self, franka_embodiment, sim_device: str, enable_visualization: bool = False ) -> RetargeterCfg | None: return None + + +@register_retargeter +class AgibotKeyboardRetargeter(RetargetterBase): + device = "keyboard" + embodiment = "agibot" + + def __init__(self): + pass + + def get_retargeter_cfg( + self, agibot_embodiment, sim_device: str, enable_visualization: bool = False + ) -> RetargeterCfg | None: + return None diff --git a/isaaclab_arena/embodiments/agibot/agibot.py b/isaaclab_arena/embodiments/agibot/agibot.py index 2d33417f..3ea0107a 100644 --- a/isaaclab_arena/embodiments/agibot/agibot.py +++ b/isaaclab_arena/embodiments/agibot/agibot.py @@ -34,65 +34,44 @@ class AgibotEmbodiment(EmbodimentBase): default_arm_mode = ArmMode.LEFT def __init__( - self, - enable_cameras: bool = False, - initial_pose: Pose | None = None, - arm_mode: ArmMode | None = None, + self, enable_cameras: bool = False, initial_pose: Pose | None = None, arm_mode: ArmMode = ArmMode.LEFT ): super().__init__(enable_cameras, initial_pose) self.arm_mode = arm_mode or self.default_arm_mode - self.scene_config = AgibotSceneCfg() - self.scene_config.arm_mode = self.arm_mode - self.action_config = AgibotActionsCfg() - self.action_config.arm_mode = self.arm_mode + self.scene_config = AgibotLeftArmSceneCfg() if self.arm_mode == ArmMode.LEFT else AgibotRightArmSceneCfg() + self.action_config = AgibotLeftArmActionsCfg() if self.arm_mode == ArmMode.LEFT else AgibotRightArmActionsCfg() self.observation_config = AgibotObservationsCfg() self.mimic_env = AgibotMimicEnv @configclass class AgibotSceneCfg: - """Scene configuration for the Agibot A2D robot.""" + """Scene configuration for the Agibot.""" robot = AGIBOT_A2D_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - arm_mode: ArmMode = MISSING - ee_frame: FrameTransformerCfg = MISSING + +@configclass +class AgibotLeftArmSceneCfg(AgibotSceneCfg): + """Scene configuration for the Agibot left arm.""" + + ee_frame = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/base_link", + debug_vis=False, + target_frames=[ + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/gripper_center", + name="left_end_effector", + offset=OffsetCfg( + rot=(0.7071, 0.0, -0.7071, 0.0), + ), + ), + ], + ) + def __post_init__(self): - if self.arm_mode == ArmMode.LEFT: - self.ee_frame = FrameTransformerCfg( - prim_path="{ENV_REGEX_NS}/Robot/base_link", - debug_vis=False, - target_frames=[ - FrameTransformerCfg.FrameCfg( - prim_path="{ENV_REGEX_NS}/Robot/gripper_center", - name="left_end_effector", - offset=OffsetCfg( - pos=[0.0, 0.0, 0.0], - rot=[ - 0.7071, - 0.0, - -0.7071, - 0.0, - ], # rpy: [0, -90, 0] to make the gripper direction from +Z to +X - ), - ), - ], - ) - elif self.arm_mode == ArmMode.RIGHT: - self.ee_frame = FrameTransformerCfg( - prim_path="{ENV_REGEX_NS}/Robot/base_link", - debug_vis=False, - target_frames=[ - FrameTransformerCfg.FrameCfg( - prim_path="{ENV_REGEX_NS}/Robot/right_gripper_center", - name="right_end_effector", - ), - ], - ) - else: - raise ValueError(f"Arm mode: {self.arm_mode} is not supported yet.") # Add a marker to the end-effector frame marker_cfg = FRAME_MARKER_CFG.copy() marker_cfg.markers["frame"].scale = (0.1, 0.1, 0.1) @@ -101,52 +80,71 @@ def __post_init__(self): @configclass -class AgibotActionsCfg: - """Action configuration for the Agibot robot.""" - - arm_mode: ArmMode = MISSING - arm_action: RMPFlowActionCfg = MISSING - gripper_action: mdp.AbsBinaryJointPositionActionCfg = MISSING +class AgibotRightArmSceneCfg(AgibotSceneCfg): + """Scene configuration for the Agibot right arm.""" + + ee_frame = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/base_link", + debug_vis=False, + target_frames=[ + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/right_gripper_center", + name="right_end_effector", + ), + ], + ) def __post_init__(self): - if self.arm_mode == ArmMode.LEFT: - self.arm_action = RMPFlowActionCfg( - asset_name="robot", - joint_names=["left_arm_joint.*"], - body_name="gripper_center", - controller=AGIBOT_LEFT_ARM_RMPFLOW_CFG, - scale=1.0, - body_offset=RMPFlowActionCfg.OffsetCfg(pos=[0.0, 0.0, 0.0], rot=[0.7071, 0.0, -0.7071, 0.0]), - articulation_prim_expr="/World/envs/env_.*/Robot", - use_relative_mode=True, - ) + # Add a marker to the end-effector frame + marker_cfg = FRAME_MARKER_CFG.copy() + marker_cfg.markers["frame"].scale = (0.1, 0.1, 0.1) + marker_cfg.prim_path = "/Visuals/FrameTransformer" + self.ee_frame.visualizer_cfg = marker_cfg - self.gripper_action = mdp.AbsBinaryJointPositionActionCfg( - asset_name="robot", - threshold=0.5, - joint_names=["left_hand_joint1", "left_.*_Support_Joint"], - open_command_expr={"left_hand_joint1": 0.994, "left_.*_Support_Joint": 0.994}, - close_command_expr={"left_hand_joint1": 0.0, "left_.*_Support_Joint": 0.0}, - ) - elif self.arm_mode == ArmMode.RIGHT: - self.arm_action = RMPFlowActionCfg( - asset_name="robot", - joint_names=["right_arm_joint.*"], - body_name="right_gripper_center", - controller=AGIBOT_RIGHT_ARM_RMPFLOW_CFG, - scale=1.0, - body_offset=RMPFlowActionCfg.OffsetCfg(pos=[0.0, 0.0, 0.0]), - articulation_prim_expr="/World/envs/env_.*/Robot", - use_relative_mode=True, - ) - self.gripper_action = mdp.AbsBinaryJointPositionActionCfg( - asset_name="robot", - threshold=0.5, - joint_names=["right_hand_joint1", "right_.*_Support_Joint"], - open_command_expr={"right_hand_joint1": 0.994, "right_.*_Support_Joint": 0.994}, - close_command_expr={"right_hand_joint1": 0.0, "right_.*_Support_Joint": 0.0}, - ) +@configclass +class AgibotLeftArmActionsCfg: + """Action configuration for the Agibot left arm.""" + + arm_action = RMPFlowActionCfg( + asset_name="robot", + joint_names=["left_arm_joint.*"], + body_name="gripper_center", + controller=AGIBOT_LEFT_ARM_RMPFLOW_CFG, + scale=1.0, + body_offset=RMPFlowActionCfg.OffsetCfg(rot=[0.7071, 0.0, -0.7071, 0.0]), + articulation_prim_expr="/World/envs/env_.*/Robot", + use_relative_mode=True, + ) + + gripper_action = mdp.BinaryJointPositionActionCfg( + asset_name="robot", + joint_names=["left_hand_joint1", "left_.*_Support_Joint"], + open_command_expr={"left_hand_joint1": 0.994, "left_.*_Support_Joint": 0.994}, + close_command_expr={"left_hand_joint1": 0.0, "left_.*_Support_Joint": 0.0}, + ) + + +@configclass +class AgibotRightArmActionsCfg: + """Action configuration for the Agibot right arm.""" + + arm_action = RMPFlowActionCfg( + asset_name="robot", + joint_names=["right_arm_joint.*"], + body_name="right_gripper_center", + controller=AGIBOT_RIGHT_ARM_RMPFLOW_CFG, + scale=1.0, + articulation_prim_expr="/World/envs/env_.*/Robot", + use_relative_mode=True, + ) + + gripper_action = mdp.BinaryJointPositionActionCfg( + asset_name="robot", + joint_names=["right_hand_joint1", "right_.*_Support_Joint"], + open_command_expr={"right_hand_joint1": 0.994, "right_.*_Support_Joint": 0.994}, + close_command_expr={"right_hand_joint1": 0.0, "right_.*_Support_Joint": 0.0}, + ) @configclass diff --git a/isaaclab_arena/tasks/place_upright_task.py b/isaaclab_arena/tasks/place_upright_task.py new file mode 100644 index 00000000..c7dfb5e7 --- /dev/null +++ b/isaaclab_arena/tasks/place_upright_task.py @@ -0,0 +1,227 @@ +# Copyright (c) 2025, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +from dataclasses import MISSING + +import isaaclab.envs.mdp as mdp_isaac_lab +from isaaclab.envs.common import ViewerCfg +from isaaclab.envs.mimic_env_cfg import MimicEnvCfg, SubTaskConfig +from isaaclab.managers import EventTermCfg, SceneEntityCfg, TerminationTermCfg +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.utils import configclass + +from isaaclab_arena.affordances.placeable import Placeable +from isaaclab_arena.embodiments.common.arm_mode import ArmMode +from isaaclab_arena.metrics.metric_base import MetricBase +from isaaclab_arena.metrics.object_moved import ObjectMovedRateMetric +from isaaclab_arena.metrics.success_rate import SuccessRateMetric +from isaaclab_arena.tasks.task_base import TaskBase +from isaaclab_arena.terms.events import set_object_pose +from isaaclab_arena.utils.cameras import get_viewer_cfg_look_at_object + + +class PlaceUprightTask(TaskBase): + + def __init__( + self, + placeable_object: Placeable, + orientation_threshold: float | None = None, + episode_length_s: float | None = None, + task_description: str | None = None, + ): + super().__init__(episode_length_s=episode_length_s) + assert isinstance(placeable_object, Placeable), "Placeable object must be an instance of Placeable" + self.placeable_object = placeable_object + self.orientation_threshold = ( + orientation_threshold if orientation_threshold is not None else placeable_object.orientation_threshold + ) + self.scene_config = InteractiveSceneCfg(num_envs=1, env_spacing=3.0, replicate_physics=False) + self.events_cfg = PlaceUprightEventCfg(self.placeable_object) + self.termination_cfg = self.make_termination_cfg() + self.task_description = ( + f"Place the {placeable_object.name} upright" if task_description is None else task_description + ) + + def get_scene_cfg(self): + return self.scene_config + + def get_termination_cfg(self): + return self.termination_cfg + + def make_termination_cfg(self): + params = {} + if self.orientation_threshold is not None: + params["orientation_threshold"] = self.orientation_threshold + success = TerminationTermCfg( + func=self.placeable_object.is_placed_upright, + params=params, + ) + return TerminationsCfg(success=success) + + def get_events_cfg(self): + return self.events_cfg + + def get_mimic_env_cfg(self, arm_mode: ArmMode): + return PlaceUprightMimicEnvCfg( + arm_mode=arm_mode, + placeable_object_name=self.placeable_object.name, + ) + + def get_metrics(self) -> list[MetricBase]: + return [ + SuccessRateMetric(), + ObjectMovedRateMetric(self.placeable_object), + ] + + def get_viewer_cfg(self) -> ViewerCfg: + return get_viewer_cfg_look_at_object(lookat_object=self.placeable_object, offset=np.array([1.5, 1.5, 1.5])) + + +@configclass +class TerminationsCfg: + """Termination terms for the MDP.""" + + time_out: TerminationTermCfg = TerminationTermCfg(func=mdp_isaac_lab.time_out) + + # Dependent on the placeable object, so this is passed in from the task at + # construction time. + success: TerminationTermCfg = MISSING + + +@configclass +class PlaceUprightEventCfg: + """Configuration for Place Upright.""" + + reset_placeable_object_pose: EventTermCfg = MISSING + + def __init__(self, placeable_object: Placeable): + assert isinstance(placeable_object, Placeable), "Object pose must be an instance of Placeable" + initial_pose = placeable_object.get_initial_pose() + if initial_pose is not None: + self.reset_placeable_object_pose = EventTermCfg( + func=set_object_pose, + mode="reset", + params={ + "pose": initial_pose, + "asset_cfg": SceneEntityCfg(placeable_object.name), + }, + ) + else: + raise ValueError(f"Initial pose is not set for the placeable object {placeable_object.name}") + + +@configclass +class PlaceUprightMimicEnvCfg(MimicEnvCfg): + """ + Isaac Lab Mimic environment config class for Place Upright env. + """ + + arm_mode: ArmMode = ArmMode.SINGLE_ARM + + placeable_object_name: str = "placeable_object" + + def __post_init__(self): + # post init of parents + super().__post_init__() + + # Override the existing values + self.datagen_config.name = "demo_src_placeupright_isaac_lab_task_D0" + self.datagen_config.generation_guarantee = True + self.datagen_config.generation_keep_failed = False + self.datagen_config.generation_num_trials = 100 + self.datagen_config.generation_select_src_per_subtask = False + self.datagen_config.generation_select_src_per_arm = False + self.datagen_config.generation_relative = False + self.datagen_config.generation_joint_pos = False + self.datagen_config.generation_transform_first_robot_pose = False + self.datagen_config.generation_interpolate_from_last_target_pose = True + self.datagen_config.max_num_failures = 25 + self.datagen_config.seed = 1 + + # The following are the subtask configurations for the pick and place task. + subtask_configs = [] + subtask_configs.append( + SubTaskConfig( + # Each subtask involves manipulation with respect to a single object frame. + object_ref=self.placeable_object_name, + # This key corresponds to the binary indicator in "datagen_info" that signals + # when this subtask is finished (e.g., on a 0 to 1 edge). + subtask_term_signal="grasp", + # Specifies time offsets for data generation when splitting a trajectory into + # subtask segments. Random offsets are added to the termination boundary. + subtask_term_offset_range=(0, 0), + # Selection strategy for the source subtask segment during data generation + selection_strategy="nearest_neighbor_object", + # Optional parameters for the selection strategy function + selection_strategy_kwargs={"nn_k": 3}, + # Amount of action noise to apply during this subtask + action_noise=0.005, + # Number of interpolation steps to bridge to this subtask segment + num_interpolation_steps=5, + # Additional fixed steps for the robot to reach the necessary pose + num_fixed_steps=0, + # If True, apply action noise during the interpolation phase and execution + apply_noise_during_interpolation=False, + ) + ) + subtask_configs.append( + SubTaskConfig( + # Each subtask involves manipulation with respect to a single object frame. + # TODO(alexmillane, 2025.09.02): This is currently broken. FIX. + # We need a way to pass in a reference to an object that exists in the + # scene. + object_ref=self.placeable_object_name, + # End of final subtask does not need to be detected + subtask_term_signal=None, + # No time offsets for the final subtask + subtask_term_offset_range=(0, 0), + # Selection strategy for source subtask segment + selection_strategy="nearest_neighbor_object", + # Optional parameters for the selection strategy function + selection_strategy_kwargs={"nn_k": 3}, + # Amount of action noise to apply during this subtask + action_noise=0.005, + # Number of interpolation steps to bridge to this subtask segment + num_interpolation_steps=5, + # Additional fixed steps for the robot to reach the necessary pose + num_fixed_steps=0, + # If True, apply action noise during the interpolation phase and execution + apply_noise_during_interpolation=False, + ) + ) + if self.arm_mode == ArmMode.SINGLE_ARM: + self.subtask_configs["robot"] = subtask_configs + # We need to add the left and right subtasks for GR1. + elif self.arm_mode in [ArmMode.LEFT, ArmMode.RIGHT]: + self.subtask_configs[self.arm_mode.value] = subtask_configs + # EEF on opposite side (arm is static) + subtask_configs = [] + subtask_configs.append( + SubTaskConfig( + # Each subtask involves manipulation with respect to a single object frame. + object_ref=self.placeable_object_name, + # Corresponding key for the binary indicator in "datagen_info" for completion + subtask_term_signal=None, + # Time offsets for data generation when splitting a trajectory + subtask_term_offset_range=(0, 0), + # Selection strategy for source subtask segment + selection_strategy="nearest_neighbor_object", + # Optional parameters for the selection strategy function + selection_strategy_kwargs={"nn_k": 3}, + # Amount of action noise to apply during this subtask + action_noise=0.005, + # Number of interpolation steps to bridge to this subtask segment + num_interpolation_steps=0, + # Additional fixed steps for the robot to reach the necessary pose + num_fixed_steps=0, + # If True, apply action noise during the interpolation phase and execution + apply_noise_during_interpolation=False, + ) + ) + self.subtask_configs[self.arm_mode.get_other_arm().value] = subtask_configs + + else: + raise ValueError(f"Embodiment arm mode {self.arm_mode} not supported") diff --git a/isaaclab_arena/tests/test_object_with_cfg_addons.py b/isaaclab_arena/tests/test_object_with_cfg_addons.py new file mode 100644 index 00000000..3688d87a --- /dev/null +++ b/isaaclab_arena/tests/test_object_with_cfg_addons.py @@ -0,0 +1,52 @@ +# Copyright (c) 2025, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + + +from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function + +HEADLESS = False + + +def _test_object_with_cfg_addons(simulation_app): + + from isaaclab_arena.assets.object_base import ObjectType + from isaaclab_arena.assets.object_library import LibraryObject + from isaaclab_arena.utils.pose import Pose + + class ConeWithCfgAddons(LibraryObject): + """ + Cone with cfg addons. + """ + + name = "cone_with_cfg_addons" + tags = ["object"] + usd_path = "https://omniverse-content-production.s3-us-west-2.amazonaws.com/Assets/Isaac/4.5/Isaac/Props/Shapes/cone.usd" + default_prim_path = "{ENV_REGEX_NS}/target_cone_with_cfg_addons" + object_type = ObjectType.RIGID + spawn_cfg_addon = {"visible": False} # By default, the object is visible. + asset_cfg_addon = {"debug_vis": True} # By default, the object is not debug visualized. + + def __init__(self, prim_path: str = default_prim_path, initial_pose: Pose | None = None): + super().__init__(prim_path=prim_path, initial_pose=initial_pose) + + cone = ConeWithCfgAddons() + + # Check that the settings have been applied + assert cone.object_cfg.spawn.visible is False + assert cone.object_cfg.debug_vis is True + + return True + + +def test_object_with_cfg_addons(): + result = run_simulation_app_function( + _test_object_with_cfg_addons, + headless=HEADLESS, + ) + assert result, "Test failed" + + +if __name__ == "__main__": + test_object_with_cfg_addons() diff --git a/isaaclab_arena/tests/test_place_upright_task.py b/isaaclab_arena/tests/test_place_upright_task.py new file mode 100644 index 00000000..2f72cb4b --- /dev/null +++ b/isaaclab_arena/tests/test_place_upright_task.py @@ -0,0 +1,211 @@ +# Copyright (c) 2025, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import gymnasium as gym +import torch + +from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function + +NUM_STEPS = 5 +HEADLESS = True + + +def get_test_environment(remove_randomize_mug_positions_event: bool, num_envs: int): + """Returns a scene which we use for these tests.""" + from isaaclab_arena.assets.asset_registry import AssetRegistry + from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser + from isaaclab_arena.embodiments.agibot.agibot import AgibotEmbodiment + from isaaclab_arena.embodiments.common.arm_mode import ArmMode + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder + from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment + from isaaclab_arena.scene.scene import Scene + from isaaclab_arena.tasks.place_upright_task import PlaceUprightTask + from isaaclab_arena.utils.pose import Pose + + args_parser = get_isaaclab_arena_cli_parser() + args_cli = args_parser.parse_args(["--num_envs", str(num_envs)]) + + asset_registry = AssetRegistry() + background = asset_registry.get_asset_by_name("table")() + background.set_initial_pose(Pose(position_xyz=(0.50, 0.0, 0.625), rotation_wxyz=(0.7071, 0, 0, 0.7071))) + background.object_cfg.spawn.scale = (1.0, 1.0, 0.60) + # placeable object must have initial pose set + mug = asset_registry.get_asset_by_name("mug")( + initial_pose=Pose(position_xyz=(0.05, 0.0, 0.75), rotation_wxyz=(0.7071, 0.7071, 0.0, 0.0)) + ) + + light = asset_registry.get_asset_by_name("light")() + + scene = Scene(assets=[background, mug, light]) + + isaaclab_arena_environment = IsaacLabArenaEnvironment( + name="place_upright_mug", + embodiment=AgibotEmbodiment(arm_mode=ArmMode.LEFT), + scene=scene, + task=PlaceUprightTask(mug, mug.orientation_threshold), + ) + env_builder = ArenaEnvBuilder(isaaclab_arena_environment, args_cli) + name, cfg = env_builder.build_registered() + if remove_randomize_mug_positions_event: + cfg.events.reset_all = None + cfg.events.randomize_mug_positions = None + cfg.events.reset_placeable_object_pose = None + + env = gym.make(name, cfg=cfg).unwrapped + env.reset() + + return env, mug + + +def _test_place_upright_mug_single(simulation_app) -> bool: + from isaaclab.envs.manager_based_env import ManagerBasedEnv + + from isaaclab_arena.tests.utils.simulation import step_zeros_and_call + + env, placeable_obj = get_test_environment(remove_randomize_mug_positions_event=True, num_envs=1) + + def assert_upright(env: ManagerBasedEnv, terminated: torch.Tensor): + is_upright = placeable_obj.is_placed_upright(env) + assert is_upright.shape == torch.Size([1]), "Is upright shape is not correct" + assert is_upright.item(), "The mug is not upright when it should be" + if is_upright.item(): + print("Mug is placed upright") + # Check terminated. + assert terminated.shape == torch.Size([1]), "Terminated shape is not correct" + assert terminated.item(), "The task didn't terminate when it should have" + if terminated.item(): + print("Place upright mug task is completed") + + try: + print("Placing mug upright") + placeable_obj.place_upright(env, env_ids=None) + step_zeros_and_call(env, NUM_STEPS, assert_upright) + + except Exception as e: + print(f"Error: {e}") + return False + finally: + env.close() + + return True + + +def _test_place_upright_mug_multi(simulation_app) -> bool: + + from isaaclab_arena.tests.utils.simulation import step_zeros_and_call + + env, placeable_obj = get_test_environment(remove_randomize_mug_positions_event=True, num_envs=2) + + try: + + with torch.inference_mode(): + # place both mugs upright, upright percentage is a scalar + placeable_obj.place_upright(env, None, upright_percentage=1.0) + step_zeros_and_call(env, NUM_STEPS) + is_upright = placeable_obj.is_placed_upright(env) + print(f"expected: [True, True]: got: {is_upright}") + assert torch.all(is_upright == torch.tensor([True, True], device=env.device)) + + # reset place both mugs upright, upright percentage is a tensor + placeable_obj.place_upright(env, None, upright_percentage=torch.tensor([0.0, 0.0])) + step_zeros_and_call(env, NUM_STEPS) + is_upright = placeable_obj.is_placed_upright(env) + print(f"expected: [False, False]: got: {is_upright}") + assert torch.all(is_upright == torch.tensor([False, False], device=env.device)) + + # Place first mug upright, env_ids is a tensor + placeable_obj.place_upright(env, torch.tensor([0]), upright_percentage=1.0) + step_zeros_and_call(env, NUM_STEPS) + is_upright = placeable_obj.is_placed_upright(env) + print(f"expected: [True, False]: got: {is_upright}") + assert torch.all(is_upright == torch.tensor([True, False], device=env.device)) + + # Place second mug upright, env_ids and upright percentage are tensors + placeable_obj.place_upright(env, torch.tensor([1]), upright_percentage=torch.tensor([1.0])) + step_zeros_and_call(env, NUM_STEPS) + is_upright = placeable_obj.is_placed_upright(env) + print(f"expected: [True, True]: got: {is_upright}") + assert torch.all(is_upright == torch.tensor([True, True], device=env.device)) + + # Place second mug upright, env_ids and upright percentage are tensors + placeable_obj.place_upright(env, None, upright_percentage=torch.tensor([0.0, 1.0])) + step_zeros_and_call(env, NUM_STEPS) + is_upright = placeable_obj.is_placed_upright(env) + print(f"expected: [False, True]: got: {is_upright}") + assert torch.all(is_upright == torch.tensor([False, True], device=env.device)) + + except Exception as e: + print(f"Error: {e}") + return False + finally: + env.close() + + return True + + +def _test_place_upright_mug_condition(simulation_app) -> bool: + from isaaclab_arena.tests.utils.simulation import step_zeros_and_call + + env, placeable_obj = get_test_environment(remove_randomize_mug_positions_event=False, num_envs=2) + try: + with torch.inference_mode(): + # place both mugs upright + placeable_obj.place_upright(env, None, upright_percentage=1.0) + step_zeros_and_call(env, NUM_STEPS) + is_upright = placeable_obj.is_placed_upright(env) + print(f"expected: [False, False]: got: {is_upright}") + assert torch.all(is_upright == torch.tensor([False, False], device=env.device)) + + # reset place both mugs upright + placeable_obj.place_upright(env, None, upright_percentage=0.0) + step_zeros_and_call(env, NUM_STEPS) + is_upright = placeable_obj.is_placed_upright(env) + print(f"expected: [False, False]: got: {is_upright}") + assert torch.all(is_upright == torch.tensor([False, False], device=env.device)) + + # Place first mug upright + placeable_obj.place_upright(env, torch.tensor([0])) + step_zeros_and_call(env, NUM_STEPS) + is_upright = placeable_obj.is_placed_upright(env) + print(f"expected: [False, False]: got: {is_upright}") + assert torch.all(is_upright == torch.tensor([False, False], device=env.device)) + + except Exception as e: + print(f"Error: {e}") + return False + finally: + env.close() + + return True + + +def test_place_upright_mug_single(): + result = run_simulation_app_function( + _test_place_upright_mug_single, + headless=HEADLESS, + ) + assert result, f"Test {_test_place_upright_mug_single.__name__} failed" + + +def test_place_upright_mug_multi(): + result = run_simulation_app_function( + _test_place_upright_mug_multi, + headless=HEADLESS, + ) + assert result, f"Test {_test_place_upright_mug_multi.__name__} failed" + + +def test_place_upright_mug_condition(): + result = run_simulation_app_function( + _test_place_upright_mug_condition, + headless=HEADLESS, + ) + assert result, f"Test {_test_place_upright_mug_condition.__name__} failed" + + +if __name__ == "__main__": + test_place_upright_mug_single() + test_place_upright_mug_multi() + test_place_upright_mug_condition() diff --git a/isaaclab_arena_environments/cli.py b/isaaclab_arena_environments/cli.py index 65a2b2bf..006bd68d 100644 --- a/isaaclab_arena_environments/cli.py +++ b/isaaclab_arena_environments/cli.py @@ -16,6 +16,7 @@ from isaaclab_arena_environments.kitchen_pick_and_place_environment import KitchenPickAndPlaceEnvironment from isaaclab_arena_environments.lift_object_environment import LiftObjectEnvironment from isaaclab_arena_environments.press_button_environment import PressButtonEnvironment +from isaaclab_arena_environments.tabletop_place_upright_environment import TableTopPlaceUprightEnvironment # NOTE(alexmillane, 2025.09.04): There is an issue with type annotation in this file. # We cannot annotate types which require the simulation app to be started in order to @@ -32,6 +33,7 @@ GalileoG1LocomanipPickAndPlaceEnvironment.name: GalileoG1LocomanipPickAndPlaceEnvironment, PressButtonEnvironment.name: PressButtonEnvironment, LiftObjectEnvironment.name: LiftObjectEnvironment, + TableTopPlaceUprightEnvironment.name: TableTopPlaceUprightEnvironment, } diff --git a/isaaclab_arena_environments/tabletop_place_upright_environment.py b/isaaclab_arena_environments/tabletop_place_upright_environment.py new file mode 100644 index 00000000..9b698349 --- /dev/null +++ b/isaaclab_arena_environments/tabletop_place_upright_environment.py @@ -0,0 +1,103 @@ +# Copyright (c) 2025, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import argparse + +from isaaclab_arena_environments.example_environment_base import ExampleEnvironmentBase + + +class TableTopPlaceUprightEnvironment(ExampleEnvironmentBase): + """ + A place upright environment for the Seattle Lab table. + """ + + name = "tabletop_place_upright" + + def get_env(self, args_cli: argparse.Namespace): + import isaaclab.envs.mdp as mdp + from isaaclab.managers import EventTermCfg as EventTerm + from isaaclab.managers import SceneEntityCfg + from isaaclab.utils import configclass + from isaaclab_tasks.manager_based.manipulation.stack.mdp.franka_stack_events import randomize_object_pose + + from isaaclab_arena.embodiments.common.arm_mode import ArmMode + from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment + from isaaclab_arena.scene.scene import Scene + from isaaclab_arena.tasks.place_upright_task import PlaceUprightTask + from isaaclab_arena.utils.pose import Pose + + @configclass + class EventCfgPlaceUprightMug: + """Configuration for events.""" + + reset_all = EventTerm(func=mdp.reset_scene_to_default, mode="reset", params={"reset_joint_targets": True}) + + randomize_mug_positions = EventTerm( + func=randomize_object_pose, + mode="reset", + params={ + "pose_range": { + "x": (-0.05, 0.2), + "y": (-0.10, 0.10), + "z": (0.75, 0.75), + "roll": (-1.57, -1.57), + "yaw": (-0.57, 0.57), + }, + "asset_cfgs": [SceneEntityCfg("mug")], + }, + ) + + # Add the asset registry from the arena migration package + background = self.asset_registry.get_asset_by_name(args_cli.background)() + placeable_object = self.asset_registry.get_asset_by_name(args_cli.object)( + initial_pose=Pose(position_xyz=(0.05, 0.0, 0.75), rotation_wxyz=(0.0, 1.0, 0.0, 0.0)) + ) + if args_cli.embodiment == "agibot": + embodiment = self.asset_registry.get_asset_by_name(args_cli.embodiment)( + enable_cameras=args_cli.enable_cameras, arm_mode=ArmMode.LEFT + ) + else: + raise NotImplementedError( + f"Embodiment {args_cli.embodiment} not supported for tabletop place upright environment" + ) + + if args_cli.teleop_device is not None: + teleop_device = self.device_registry.get_device_by_name(args_cli.teleop_device)() + else: + teleop_device = None + + embodiment.set_initial_pose( + Pose( + position_xyz=(-0.60, 0.0, 0.0), + rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + ) + ) + background.set_initial_pose(Pose(position_xyz=(0.50, 0.0, 0.625), rotation_wxyz=(0.7071, 0, 0, 0.7071))) + background.object_cfg.spawn.scale = (1.0, 1.0, 0.60) + + ground_plane = self.asset_registry.get_asset_by_name("ground_plane")() + light = self.asset_registry.get_asset_by_name("light")() + + scene = Scene(assets=[background, placeable_object, ground_plane, light]) + + task = PlaceUprightTask(placeable_object) + task.events_cfg = EventCfgPlaceUprightMug() + + isaaclab_arena_environment = IsaacLabArenaEnvironment( + name=self.name, + embodiment=embodiment, + scene=scene, + task=task, + teleop_device=teleop_device, + ) + return isaaclab_arena_environment + + @staticmethod + def add_cli_args(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--object", type=str, default="mug") + parser.add_argument("--background", type=str, default="table") + parser.add_argument("--embodiment", type=str, default="agibot") + parser.add_argument("--enable_cameras", type=bool, default=False) + parser.add_argument("--teleop_device", type=str, default="keyboard")