diff --git a/isaaclab_arena/assets/object_library.py b/isaaclab_arena/assets/object_library.py index 511ace198..61c7b651b 100644 --- a/isaaclab_arena/assets/object_library.py +++ b/isaaclab_arena/assets/object_library.py @@ -617,6 +617,10 @@ class RedCube(LibraryObject): name = "red_cube" tags = ["object"] + # TODO(lanceli, 2026.02.04): There is a known bug where rigid body attributes can only bind to the root layer. + # As a workaround, the original assets from ISAAC_NUCLEUS_DIR have been adjusted and uploaded to ISAAC_NUCLEUS_STAGING_DIR. + # Once this bug is resolved, the original assets can be used instead. + # usd_path =f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/red_block.usd" # not support, rigid body attribute need to be bind to root xform. usd_path = f"{ISAACLAB_STAGING_NUCLEUS_DIR}/Arena/assets/object_library/isaac_blocks/red_block_root_rigid.usd" @@ -637,6 +641,10 @@ class GreenCube(LibraryObject): name = "green_cube" tags = ["object"] + # TODO(lanceli, 2026.02.04): There is a known bug where rigid body attributes can only bind to the root layer. + # As a workaround, the original assets from ISAAC_NUCLEUS_DIR have been adjusted and uploaded to ISAAC_NUCLEUS_STAGING_DIR. + # Once this bug is resolved, the original assets can be used instead. + # usd_path = f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/green_block.usd" # not support, rigid body attribute need to be bind to root xform. usd_path = f"{ISAACLAB_STAGING_NUCLEUS_DIR}/Arena/assets/object_library/isaac_blocks/green_block_root_rigid.usd" object_type = ObjectType.RIGID diff --git a/isaaclab_arena/examples/compile_env_notebook.py b/isaaclab_arena/examples/compile_env_notebook.py index e85d93381..1e25c4c66 100644 --- a/isaaclab_arena/examples/compile_env_notebook.py +++ b/isaaclab_arena/examples/compile_env_notebook.py @@ -24,7 +24,7 @@ from isaaclab_arena.tasks.pick_and_place_task import PickAndPlaceTask from isaaclab_arena.tasks.sequential_task_base import SequentialTaskBase from isaaclab_arena.tasks.task_base import TaskBase -from isaaclab_arena.utils.pose import Pose, PoseRange +from isaaclab_arena.utils.pose import Pose asset_registry = AssetRegistry() @@ -71,48 +71,72 @@ ) ) - RANDOMIZATION_HALF_RANGE_X_M = 0.04 RANDOMIZATION_HALF_RANGE_Y_M = 0.01 RANDOMIZATION_HALF_RANGE_Z_M = 0.0 -z_position = { - "sweet_potato": 1.0, - "jug": 1.1, -}[pickup_object_name] -yaw = { - "sweet_potato": 0.0, - "jug": -90.0, -}[pickup_object_name] -pickup_object.set_initial_pose( - # Bench (no randomization) - # Pose(position_xyz=(3.922, -0.565, 1.019), rotation_wxyz=(0.7071068, 0.0, 0.0, 0.7071068)) - # Bench (with randomization) - PoseRange( - position_xyz_min=( - 4.1 - RANDOMIZATION_HALF_RANGE_X_M, - -0.6 - RANDOMIZATION_HALF_RANGE_Y_M, - z_position - RANDOMIZATION_HALF_RANGE_Z_M, - ), - position_xyz_max=( - 4.1 + RANDOMIZATION_HALF_RANGE_X_M, - -0.6 + RANDOMIZATION_HALF_RANGE_Y_M, - z_position + RANDOMIZATION_HALF_RANGE_Z_M, - ), - # position_xyz_max=( - # 3.922 + RANDOMIZATION_HALF_RANGE_X_M, - # -0.565 + RANDOMIZATION_HALF_RANGE_Y_M, - # 1.019 + RANDOMIZATION_HALF_RANGE_Z_M - # ), - rpy_min=(0.0, 0.0, yaw), - rpy_max=(0.0, 0.0, yaw), + +from isaaclab_arena.relations.relations import AtPosition, IsAnchor, On, RandomAroundSolution + +kitchen_counter_top = ObjectReference( + name="kitchen_counter_top", + prim_path="{ENV_REGEX_NS}/lightwheel_robocasa_kitchen/counter_right_main_group/top_geometry", + parent_asset=kitchen_background, +) +kitchen_counter_top.add_relation(IsAnchor()) +pickup_object.add_relation(On(kitchen_counter_top)) +pickup_object.add_relation(AtPosition(x=4.1, y=-0.6)) +pickup_object.add_relation( + RandomAroundSolution( + x_half_m=RANDOMIZATION_HALF_RANGE_X_M, + y_half_m=RANDOMIZATION_HALF_RANGE_Y_M, + z_half_m=RANDOMIZATION_HALF_RANGE_Z_M, + yaw_base_rad=-90.0, ) - # Above shelf - # Pose( - # position_xyz=(4.625, -0.395, 1.224), - # rotation_wxyz=(0.7071068, 0.0, 0.0, 0.7071068) - # ) ) +# Add reference to the scene. +# scene = Scene(assets=[kitchen_background, kitchen_counter_top, refrigerator, refrigerator_shelf, pickup_object, light]) + + +# z_position = { +# "sweet_potato": 1.0, +# "jug": 1.1, +# }[pickup_object_name] +# yaw = { +# "sweet_potato": 0.0, +# "jug": -90.0, +# }[pickup_object_name] + +# pickup_object.set_initial_pose( +# # Bench (no randomization) +# # Pose(position_xyz=(3.922, -0.565, 1.019), rotation_wxyz=(0.7071068, 0.0, 0.0, 0.7071068)) +# # Bench (with randomization) +# PoseRange( +# position_xyz_min=( +# 4.1 - RANDOMIZATION_HALF_RANGE_X_M, +# -0.6 - RANDOMIZATION_HALF_RANGE_Y_M, +# z_position - RANDOMIZATION_HALF_RANGE_Z_M, +# ), +# position_xyz_max=( +# 4.1 + RANDOMIZATION_HALF_RANGE_X_M, +# -0.6 + RANDOMIZATION_HALF_RANGE_Y_M, +# z_position + RANDOMIZATION_HALF_RANGE_Z_M, +# ), +# # position_xyz_max=( +# # 3.922 + RANDOMIZATION_HALF_RANGE_X_M, +# # -0.565 + RANDOMIZATION_HALF_RANGE_Y_M, +# # 1.019 + RANDOMIZATION_HALF_RANGE_Z_M +# # ), +# rpy_min=(0.0, 0.0, yaw), +# rpy_max=(0.0, 0.0, yaw), +# ) +# # Above shelf +# # Pose( +# # position_xyz=(4.625, -0.395, 1.224), +# # rotation_wxyz=(0.7071068, 0.0, 0.0, 0.7071068) +# # ) +# ) + class PutAndCloseDoorTask(SequentialTaskBase): @@ -138,7 +162,7 @@ def get_mimic_env_cfg(self, arm_mode): task = PutAndCloseDoorTask(subtasks=[pick_and_place_task, close_door_task]) -scene = Scene(assets=[kitchen_background, refrigerator, refrigerator_shelf, pickup_object, light]) +scene = Scene(assets=[kitchen_background, kitchen_counter_top, refrigerator, refrigerator_shelf, pickup_object, light]) isaaclab_arena_environment = IsaacLabArenaEnvironment( name="reference_object_test", embodiment=embodiment, diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py index 8c35278ef..5ca7d7117 100644 --- a/isaaclab_arena/relations/object_placer.py +++ b/isaaclab_arena/relations/object_placer.py @@ -11,7 +11,7 @@ from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams from isaaclab_arena.relations.placement_result import PlacementResult from isaaclab_arena.relations.relation_solver import RelationSolver -from isaaclab_arena.relations.relations import RandomAroundSolution, get_anchor_objects +from isaaclab_arena.relations.relations import RandomAroundSolution, RotateAroundSolution, get_anchor_objects from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox, get_random_pose_within_bounding_box from isaaclab_arena.utils.pose import Pose @@ -208,18 +208,24 @@ def _apply_positions( ) -> None: """Apply solved positions to objects (skipping anchors). - If an object has a RandomAroundSolution marker, a PoseRange is created - centered on the solved position. Otherwise, a fixed Pose is set. + If RandomAroundSolution marker is present, sets a PoseRange (for reset-time randomization). + Rotation is taken from RotateAroundSolution marker if present, otherwise keep the identity rotation. """ for obj, pos in positions.items(): if obj in anchor_objects: continue random_marker = self._get_random_around_solution(obj) + rotate_marker = self._get_rotate_around_solution(obj) + rotation_wxyz = rotate_marker.get_rotation_wxyz() if rotate_marker else (1.0, 0.0, 0.0, 0.0) + if random_marker is not None: - obj.set_initial_pose(random_marker.to_pose_range(pos)) + # We need to set a PoseRange for the randomization to be picked up on reset. + # Set a PoseRange with the explicit rotation from RotateAroundSolution if present + obj.set_initial_pose(random_marker.to_pose_range_centered_at(pos, rotation_wxyz=rotation_wxyz)) else: - obj.set_initial_pose(Pose(position_xyz=pos, rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + # Without randomization, we can set a fixed Pose. + obj.set_initial_pose(Pose(position_xyz=pos, rotation_wxyz=rotation_wxyz)) def _get_random_around_solution(self, obj: Object | ObjectReference) -> RandomAroundSolution | None: """Get RandomAroundSolution marker from object if present. @@ -235,6 +241,20 @@ def _get_random_around_solution(self, obj: Object | ObjectReference) -> RandomAr return rel return None + def _get_rotate_around_solution(self, obj: Object | ObjectReference) -> RotateAroundSolution | None: + """Get RotateAroundSolution marker from object if present. + + Args: + obj: Object to check for the marker. + + Returns: + The RotateAroundSolution marker if found, None otherwise. + """ + for rel in obj.get_relations(): + if isinstance(rel, RotateAroundSolution): + return rel + return None + @property def last_loss_history(self) -> list[float]: """Loss values from the most recent place() call.""" diff --git a/isaaclab_arena/relations/relations.py b/isaaclab_arena/relations/relations.py index 3cb5baf80..820a26b32 100644 --- a/isaaclab_arena/relations/relations.py +++ b/isaaclab_arena/relations/relations.py @@ -5,9 +5,12 @@ from __future__ import annotations +import torch from enum import Enum from typing import TYPE_CHECKING +from isaaclab.utils.math import euler_xyz_from_quat + from isaaclab_arena.utils.pose import PoseRange if TYPE_CHECKING: @@ -151,6 +154,9 @@ def __init__( roll_half_rad: float = 0.0, pitch_half_rad: float = 0.0, yaw_half_rad: float = 0.0, + roll_base_rad: float = 0.0, + pitch_base_rad: float = 0.0, + yaw_base_rad: float = 0.0, ): """ Args: @@ -160,6 +166,9 @@ def __init__( roll_half_rad: Half-extent for roll (radians). Rotation will be randomized ±roll_half_rad. pitch_half_rad: Half-extent for pitch (radians). Rotation will be randomized ±pitch_half_rad. yaw_half_rad: Half-extent for yaw (radians). Rotation will be randomized ±yaw_half_rad. + roll_base_rad: Base roll angle (radians). Center of the roll randomization range. + pitch_base_rad: Base pitch angle (radians). Center of the pitch randomization range. + yaw_base_rad: Base yaw angle (radians). Center of the yaw randomization range. """ self.x_half_m = x_half_m self.y_half_m = y_half_m @@ -167,16 +176,32 @@ def __init__( self.roll_half_rad = roll_half_rad self.pitch_half_rad = pitch_half_rad self.yaw_half_rad = yaw_half_rad + self.roll_base_rad = roll_base_rad + self.pitch_base_rad = pitch_base_rad + self.yaw_base_rad = yaw_base_rad - def to_pose_range(self, position: tuple[float, float, float]) -> PoseRange: - """Create a PoseRange centered on the given position. + def to_pose_range_centered_at( + self, + position: tuple[float, float, float], + rotation_wxyz: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0), + ) -> PoseRange: + """Create a PoseRange centered on the given position and rotation. Args: position: Center position (x, y, z) for the range. + rotation_wxyz: Center rotation as quaternion (w, x, y, z) for the range. + Defaults to identity quaternion. Returns: - PoseRange spanning ± half-extents around the position. + PoseRange spanning ± half-extents around the position and rotation. """ + # Convert quaternion to euler angles (roll, pitch, yaw) + quat_tensor = torch.tensor([rotation_wxyz]) + roll, pitch, yaw = euler_xyz_from_quat(quat_tensor) + center_roll = float(roll[0]) + center_pitch = float(pitch[0]) + center_yaw = float(yaw[0]) + return PoseRange( position_xyz_min=( position[0] - self.x_half_m, @@ -189,18 +214,65 @@ def to_pose_range(self, position: tuple[float, float, float]) -> PoseRange: position[2] + self.z_half_m, ), rpy_min=( - -self.roll_half_rad, - -self.pitch_half_rad, - -self.yaw_half_rad, + center_roll - self.roll_half_rad, + center_pitch - self.pitch_half_rad, + center_yaw - self.yaw_half_rad, ), rpy_max=( - self.roll_half_rad, - self.pitch_half_rad, - self.yaw_half_rad, + center_roll + self.roll_half_rad, + center_pitch + self.pitch_half_rad, + center_yaw + self.yaw_half_rad, ), ) +class RotateAroundSolution(RelationBase): + """Marker specifying an explicit rotation to apply on top of the solver solution. + + When ObjectPlacer applies positions, objects with this marker will have the + specified rotation applied on top of the solved position to create a fixed Pose. + + Note: This is NOT a spatial relation - the RelationSolver ignores it. It only + affects how ObjectPlacer applies the solved position to the object. + + Usage: + import math + box.add_relation(On(desk)) + box.add_relation(RotateAroundSolution(yaw_rad=math.pi / 4)) + # -> ObjectPlacer sets a Pose with solved position and 45° yaw rotation + """ + + def __init__( + self, + roll_rad: float = 0.0, + pitch_rad: float = 0.0, + yaw_rad: float = 0.0, + ): + """ + Args: + roll_rad: Roll rotation in radians. + pitch_rad: Pitch rotation in radians. + yaw_rad: Yaw rotation in radians. + """ + self.roll_rad = roll_rad + self.pitch_rad = pitch_rad + self.yaw_rad = yaw_rad + + def get_rotation_wxyz(self) -> tuple[float, float, float, float]: + """Get the rotation as a quaternion (w, x, y, z). + + Returns: + Quaternion rotation converted from roll/pitch/yaw. + """ + from isaaclab.utils.math import quat_from_euler_xyz + + roll = torch.tensor(self.roll_rad) + pitch = torch.tensor(self.pitch_rad) + yaw = torch.tensor(self.yaw_rad) + quat = quat_from_euler_xyz(roll, pitch, yaw) + return tuple(quat.tolist()) + + class AtPosition(RelationBase): """Constrains object to specific world coordinates. diff --git a/isaaclab_arena/tasks/sorting_task.py b/isaaclab_arena/tasks/sorting_task.py index 3ffb758fe..427e842ed 100644 --- a/isaaclab_arena/tasks/sorting_task.py +++ b/isaaclab_arena/tasks/sorting_task.py @@ -18,6 +18,7 @@ from isaaclab_arena.tasks.task_base import TaskBase from isaaclab_arena.tasks.terminations import objects_on_destinations, root_height_below_minimum_multi_objects from isaaclab_arena.utils.cameras import get_viewer_cfg_look_at_object +from isaaclab_arena.utils.configclass import make_configclass class SortMultiObjectTask(TaskBase): @@ -30,34 +31,35 @@ def __init__( episode_length_s: float | None = None, ): super().__init__(episode_length_s=episode_length_s) + assert len(pick_up_object_list) == len(destination_location_list) + self.pick_up_object_list = pick_up_object_list self.destination_location_list = destination_location_list self.background_scene = background_scene - assert len(pick_up_object_list) == len(destination_location_list) - - pick_up_object_contact_sensor_list = [] - for pick_up_object, destination_location in zip(pick_up_object_list, destination_location_list): - pick_up_object_contact_sensor_list.append( - pick_up_object.get_contact_sensor_cfg(contact_against_prim_paths=[destination_location.get_prim_path()]) + self.pick_up_object_contact_sensor_list = [] + self.contact_sensor_name_list = [] + for pick_up_object, destination in zip(pick_up_object_list, destination_location_list): + self.pick_up_object_contact_sensor_list.append( + pick_up_object.get_contact_sensor_cfg(contact_against_prim_paths=[destination.get_prim_path()]) ) - self.pick_up_object_contact_sensor_list = pick_up_object_contact_sensor_list - self.contact_sensor_name_list = [ - f"contact_sensor_{i}" for i in range(len(self.pick_up_object_contact_sensor_list)) - ] + self.contact_sensor_name_list.append(f"contact_sensor_{pick_up_object.name}") self.events_cfg = None self.scene_config = self.make_scene_cfg() self.termination_cfg = self.make_termination_cfg() def make_scene_cfg(self): - self.scene_config = SceneCfg() - for name, pick_up_object_contact_sensor in zip( + # Support variable number of contact sensors. + fields: list[tuple[str, type, ContactSensorCfg]] = [] + for contact_sensor_name, contact_sensor_cfg in zip( self.contact_sensor_name_list, self.pick_up_object_contact_sensor_list ): - setattr(self.scene_config, name, pick_up_object_contact_sensor) - return self.scene_config + fields.append((contact_sensor_name, type(contact_sensor_cfg), contact_sensor_cfg)) + SceneCfg = make_configclass("SceneCfg", fields) + scene_cfg = SceneCfg() + return scene_cfg def get_scene_cfg(self): return self.scene_config @@ -109,17 +111,6 @@ def get_viewer_cfg(self) -> ViewerCfg: ) -@configclass -class SceneCfg: - """ - Scene configuration for the pick and place task. - Note: only support <4 objects. Need to figure out a more flexible method, like __post_init__() - """ - - contact_sensor_0: ContactSensorCfg = MISSING - contact_sensor_1: ContactSensorCfg = MISSING - - @configclass class TerminationsCfg: """Termination terms for the MDP.""" diff --git a/isaaclab_arena_environments/sorting_environment.py b/isaaclab_arena_environments/sorting_environment.py index b351c36cc..dc15c1759 100644 --- a/isaaclab_arena_environments/sorting_environment.py +++ b/isaaclab_arena_environments/sorting_environment.py @@ -22,7 +22,9 @@ def get_env(self, args_cli: argparse.Namespace): from isaaclab_arena.tasks.sorting_task import SortMultiObjectTask from isaaclab_arena.utils.pose import Pose - assert len(args_cli.destination) == len(args_cli.object) + assert ( + len(args_cli.destinations) == len(args_cli.objects) == 2 + ), "Only 2 objects and 2 destinations are supported in this environment." # Add the asset registry from the arena migration package light = self.asset_registry.get_asset_by_name("light")() @@ -62,32 +64,32 @@ def get_env(self, args_cli: argparse.Namespace): else: teleop_device = None - destination_location1 = self.asset_registry.get_asset_by_name(args_cli.destination[0])() - destination_location1.set_initial_pose( + destination_location_1 = self.asset_registry.get_asset_by_name(args_cli.destinations[0])() + destination_location_1.set_initial_pose( Pose( position_xyz=(0.0, 0.1, 0.1), rotation_wxyz=(1.0, 0.0, 0.0, 0.0), ) ) - destination_location2 = self.asset_registry.get_asset_by_name(args_cli.destination[1])() - destination_location2.set_initial_pose( + destination_location_2 = self.asset_registry.get_asset_by_name(args_cli.destinations[1])() + destination_location_2.set_initial_pose( Pose( position_xyz=(0.0, -0.1, 0.1), rotation_wxyz=(1.0, 0.0, 0.0, 0.0), ) ) - pick_up_object1 = self.asset_registry.get_asset_by_name(args_cli.object[0])() - pick_up_object1.set_initial_pose( + pick_up_object_1 = self.asset_registry.get_asset_by_name(args_cli.objects[0])() + pick_up_object_1.set_initial_pose( Pose( position_xyz=(0.0, 0.3, 0.1), rotation_wxyz=(1.0, 0.0, 0.0, 0.0), ) ) - pick_up_object2 = self.asset_registry.get_asset_by_name(args_cli.object[1])() - pick_up_object2.set_initial_pose( + pick_up_object_2 = self.asset_registry.get_asset_by_name(args_cli.objects[1])() + pick_up_object_2.set_initial_pose( Pose( position_xyz=(0.0, -0.3, 0.1), rotation_wxyz=(1.0, 0.0, 0.0, 0.0), @@ -95,11 +97,20 @@ def get_env(self, args_cli: argparse.Namespace): ) scene = Scene( - assets=[background, light, pick_up_object1, pick_up_object2, destination_location1, destination_location2] + assets=[ + background, + light, + pick_up_object_1, + pick_up_object_2, + destination_location_1, + destination_location_2, + ] ) task = SortMultiObjectTask( - [pick_up_object1, pick_up_object2], [destination_location1, destination_location2], background + pick_up_object_list=[pick_up_object_1, pick_up_object_2], + destination_location_list=[destination_location_1, destination_location_2], + background_scene=background, ) # add custom force threshold for success termination @@ -117,16 +128,16 @@ def get_env(self, args_cli: argparse.Namespace): @staticmethod def add_cli_args(parser: argparse.ArgumentParser) -> None: parser.add_argument( - "--object", + "--objects", nargs="*", default=["red_cube", "green_cube"], - help="object list (example: --object red_cube green_cube)", + help="object list (example: --objects red_cube green_cube)", ) parser.add_argument( - "--destination", + "--destinations", nargs="*", default=["red_container", "green_container"], - help="destination list (example: --destination red_container green_container)", + help="destination list (example: --destinations red_container green_container)", ) parser.add_argument("--background", type=str, default="table") parser.add_argument("--embodiment", type=str, default="franka")