diff --git a/embodichain/lab/sim/planners/curobo/curobo_planner.py b/embodichain/lab/sim/planners/curobo/curobo_planner.py index 6259134de..e914fee6e 100644 --- a/embodichain/lab/sim/planners/curobo/curobo_planner.py +++ b/embodichain/lab/sim/planners/curobo/curobo_planner.py @@ -54,6 +54,7 @@ PlanOptions, validate_plan_options, ) +from embodichain.lab.sim.planners.curobo.curobo_yaml import _named_rigid_objects from embodichain.lab.sim.planners.utils import MoveType, PlanResult, PlanState if TYPE_CHECKING: @@ -143,18 +144,6 @@ def __deepcopy__(self, memo: dict) -> "_RigidObjectRefMapping": # noqa: ARG002 return _RigidObjectRefMapping(self) -def _named_rigid_objects( - rigid_objects: list[RigidObject] | Mapping[str, RigidObject] | None, -) -> list[tuple[str, RigidObject]]: - """Return canonical cuRobo obstacle names paired with their live objects.""" - if isinstance(rigid_objects, Mapping): - return list(rigid_objects.items()) - return [ - (getattr(obj, "uid", None) or f"obstacle_{index}", obj) - for index, obj in enumerate(rigid_objects or ()) - ] - - @configclass class CuroboWorldCfg: """Static collision-world configuration for the cuRobo backend. diff --git a/embodichain/lab/sim/planners/curobo/curobo_yaml.py b/embodichain/lab/sim/planners/curobo/curobo_yaml.py index f470f8e4b..65015695e 100644 --- a/embodichain/lab/sim/planners/curobo/curobo_yaml.py +++ b/embodichain/lab/sim/planners/curobo/curobo_yaml.py @@ -42,6 +42,18 @@ __all__ = ["generate_curobo_robot_yaml", "generate_curobo_world_yaml"] +def _named_rigid_objects( + rigid_objects: Sequence[RigidObject] | Mapping[str, RigidObject] | None, +) -> list[tuple[str, RigidObject]]: + """Return canonical obstacle names paired with their live objects.""" + if isinstance(rigid_objects, Mapping): + return list(rigid_objects.items()) + return [ + (getattr(obj, "uid", None) or f"obstacle_{index}", obj) + for index, obj in enumerate(rigid_objects or ()) + ] + + def _parse_mimic_joint_names(urdf_path: str) -> set[str]: """Return the names of URDF joints that mimic another joint. @@ -598,13 +610,7 @@ def generate_curobo_world_yaml( import yaml registry_backed = isinstance(rigid_objects, Mapping) - if registry_backed: - named_rigid_objects = list(rigid_objects.items()) - else: - named_rigid_objects = [ - (getattr(obj, "uid", None) or f"obstacle_{idx}", obj) - for idx, obj in enumerate(rigid_objects) - ] + named_rigid_objects = _named_rigid_objects(rigid_objects) if not named_rigid_objects: raise ValueError("rigid_objects must contain at least one RigidObject.") diff --git a/scripts/tutorials/atomic_action/assemble.py b/scripts/tutorials/atomic_action/assemble.py index 99151be3d..e325f12bc 100644 --- a/scripts/tutorials/atomic_action/assemble.py +++ b/scripts/tutorials/atomic_action/assemble.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Demonstrate object assembly with a dual-arm UR5. +"""Demonstrate object assembly with a selectable dual-arm robot. The left arm picks up a soda can (object A) and places it directly above a cube (object B). The relative pose of the can with respect to the cube is declared on @@ -55,16 +55,16 @@ from embodichain.lab.sim.shapes import CubeCfg, MeshCfg from embodichain.utils import logger from scripts.tutorials.atomic_action.scenario_utils import ( - add_dual_ur5_robot, + add_dual_tutorial_robot, add_support_surface, - make_dual_ur5_solver_cfg, settle_object, ) from scripts.tutorials.atomic_action.tutorial_utils import ( + TutorialRobot, broadcast_pose_batch, clone_local_pose_from_first_env, create_antipodal_semantics, - create_toppra_motion_generator, + create_curobo_motion_generator, create_tutorial_argument_parser, create_tutorial_simulation, draw_axis_marker, @@ -72,6 +72,7 @@ prepare_tutorial_scene, publish_tutorial_scene, replay_trajectory, + run_tutorial, serve_tutorial_scene, ) @@ -87,7 +88,7 @@ # --- Adjustable scene placeholders ----------------------------------------- # Object A (soda can) is staged for the left arm to pick up; object B (cube) is -# the assembly base the can is placed onto. Tweak these to match the dual-UR5 +# the assembly base the can is placed onto. Tweak these to match the dual-arm # reach and the soda-can mesh geometry. OBJECT_A_XY = (0.0, 0.02) OBJECT_B_XY = (0.0, 0.20) @@ -132,22 +133,23 @@ def parse_arguments() -> argparse.Namespace: "headless_play", "visualize_axes", ), - default_device="cpu", default_renderer="hybrid", ) return parser.parse_args() -def create_dual_ur5_robot(sim: SimulationManager) -> Robot: - """Create a dual-UR5 robot with one PGI gripper on each arm.""" - return add_dual_ur5_robot( +def create_dual_robot( + sim: SimulationManager, + robot_type: TutorialRobot, +) -> Robot: + """Create the selected dual-arm robot with one PGI gripper per arm.""" + return add_dual_tutorial_robot( sim, - uid="DualUR5Assemble", - urdf_name="dual_ur5_assemble", - solver_cfg=make_dual_ur5_solver_cfg( - GRIPPER_TCP_Z, - ur_ik_nearest_weight=(1.0, 4.0, 1.0, 1.0, 1.0, 1.0), - ), + robot_type=robot_type, + uid=f"Dual{robot_type.title()}Assemble", + urdf_name=f"dual_{robot_type}_assemble", + tcp_z=GRIPPER_TCP_Z, + ur_ik_nearest_weight=(1.0, 4.0, 1.0, 1.0, 1.0, 1.0), hand_stiffness=1e2, hand_damping=1e1, hand_max_effort=1e3, @@ -257,7 +259,7 @@ def run_assemble_demo( n_sample=args.n_sample, force_reannotate=args.force_reannotate, ) - motion_gen = create_toppra_motion_generator(robot) + motion_gen = create_curobo_motion_generator(robot) left_open, left_close = get_hand_open_close_qpos( robot, hand_control_part="left_hand", close_qpos=HAND_CLOSE_QPOS ) @@ -326,14 +328,20 @@ def run_assemble_demo( "pick_up", GraspGoal(can_semantics), binding, - MotionPolicy(sample_count=PICKUP_SAMPLE_INTERVAL), + MotionPolicy( + strategy="motion_gen", + sample_count=PICKUP_SAMPLE_INTERVAL, + ), skill_options=pick_up_options, ), ActionInvocation( "place", AssembleGoal(affordance=assemble_affordance), binding, - MotionPolicy(sample_count=PLACE_SAMPLE_INTERVAL), + MotionPolicy( + strategy="motion_gen", + sample_count=PLACE_SAMPLE_INTERVAL, + ), skill_options=place_options, ), ) @@ -374,13 +382,10 @@ def main() -> None: arena_space=3.0, light_pos=(0.0, -0.4, 3.0), ) - robot = create_dual_ur5_robot(sim) - try: - run_assemble_demo(args, sim, robot) - serve_tutorial_scene(sim, args) - finally: - sim.destroy() + robot = create_dual_robot(sim, args.robot) + run_assemble_demo(args, sim, robot) + serve_tutorial_scene(sim, args) if __name__ == "__main__": - main() + run_tutorial(main) diff --git a/scripts/tutorials/atomic_action/coordinated_pickment.py b/scripts/tutorials/atomic_action/coordinated_pickment.py index b953ee466..8d2e15a50 100644 --- a/scripts/tutorials/atomic_action/coordinated_pickment.py +++ b/scripts/tutorials/atomic_action/coordinated_pickment.py @@ -16,8 +16,8 @@ """Demonstrate dual-arm coordinated pickment with selectable object meshes. -The two UR5 arms pinch opposite sides of one object, lift it together, and move -the object to an object-centric target pose while both grippers stay closed. +The two selected arms pinch opposite sides of one object, lift it together, and +move the object to an object-centric target pose while both grippers stay closed. """ from __future__ import annotations @@ -54,17 +54,17 @@ from embodichain.utils import logger from embodichain.utils.math import matrix_from_euler from scripts.tutorials.atomic_action.scenario_utils import ( - add_dual_ur5_robot, + add_dual_tutorial_robot, add_support_surface, compute_world_bounds, get_local_vertices, log_action_plan, - make_dual_ur5_solver_cfg, resolve_cached_data_path, rotate_pose_about_world_z, settle_object, ) from scripts.tutorials.atomic_action.tutorial_utils import ( + TutorialRobot, broadcast_pose_batch, clone_local_pose_from_first_env, create_antipodal_semantics, @@ -189,15 +189,18 @@ def parse_arguments() -> argparse.Namespace: return parser.parse_args() -def create_dual_ur5_robot(sim: SimulationManager) -> Robot: - """Create a dual-UR5 robot with one PGI gripper on each arm.""" - return add_dual_ur5_robot( +def create_dual_robot( + sim: SimulationManager, + robot_type: TutorialRobot, +) -> Robot: + """Create the selected dual-arm robot with one PGI gripper per arm.""" + return add_dual_tutorial_robot( sim, - uid="DualUR5CoordinatedPickment", - urdf_name="dual_ur5_coordinated_pickment", - arm_urdf_path=resolve_cached_data_path("UniversalRobots/UR5/UR5.urdf"), - gripper_urdf_path=resolve_cached_data_path("DH_PGI_140_80/DH_PGI_140_80.urdf"), - solver_cfg=make_dual_ur5_solver_cfg(GRIPPER_TCP_Z, solver="pytorch"), + robot_type=robot_type, + uid=f"Dual{robot_type.title()}CoordinatedPickment", + urdf_name=f"dual_{robot_type}_coordinated_pickment", + tcp_z=GRIPPER_TCP_Z, + solver="pytorch", ) @@ -278,12 +281,14 @@ def compute_left_to_right_arm_direction( Returns: A normalized ``(3,)`` direction vector. """ - left_base = robot.get_link_pose( - link_name="left_base_link", env_ids=[0], to_matrix=True - )[0, :3, 3] - right_base = robot.get_link_pose( - link_name="right_base_link", env_ids=[0], to_matrix=True - )[0, :3, 3] + left_root = robot.cfg.solver_cfg["left_arm"].root_link_name + right_root = robot.cfg.solver_cfg["right_arm"].root_link_name + left_base = robot.get_link_pose(link_name=left_root, env_ids=[0], to_matrix=True)[ + 0, :3, 3 + ] + right_base = robot.get_link_pose(link_name=right_root, env_ids=[0], to_matrix=True)[ + 0, :3, 3 + ] direction = (right_base - left_base).to(device=device, dtype=torch.float32) return direction / direction.norm().clamp_min(1e-6) @@ -446,7 +451,10 @@ def run_coordinated_pickment_demo( manipulators={"left": "left_arm", "right": "right_arm"}, end_effectors={"left": "left_hand", "right": "right_hand"}, ), - MotionPolicy(sample_count=PICKMENT_SAMPLE_INTERVAL), + MotionPolicy( + strategy="motion_gen", + sample_count=PICKMENT_SAMPLE_INTERVAL, + ), skill_options=pickment_options, ), ) @@ -506,7 +514,7 @@ def main() -> None: arena_space=3.0, light_pos=(0.0, -0.4, 3.0), ) - robot = create_dual_ur5_robot(sim) + robot = create_dual_robot(sim, args.robot) run_coordinated_pickment_demo(args, sim, robot) diff --git a/scripts/tutorials/atomic_action/coordinated_placement.py b/scripts/tutorials/atomic_action/coordinated_placement.py index 0fe36f519..c3c9244d4 100644 --- a/scripts/tutorials/atomic_action/coordinated_placement.py +++ b/scripts/tutorials/atomic_action/coordinated_placement.py @@ -16,8 +16,8 @@ """Demonstrate dual-arm coordinated placement with bread and pan meshes. -The left UR5 picks up bread. The right UR5 picks up a pan and moves it to the -lower alignment pose. The left UR5 places the bread above the pan and releases +The left arm picks up bread. The right arm picks up a pan and moves it to the +lower alignment pose. The left arm places the bread above the pan and releases it while the right hand keeps holding the pan. """ @@ -59,14 +59,13 @@ from embodichain.lab.sim.shapes import MeshCfg from embodichain.utils import logger from scripts.tutorials.atomic_action.scenario_utils import ( - add_dual_ur5_robot, + add_dual_tutorial_robot, compute_local_bounds, compute_world_bounds, create_manual_object_semantics, get_local_vertices, invert_pose, log_action_plan, - make_dual_ur5_solver_cfg, normalize_vector, resolve_cached_data_path, rotate_pose_about_world_z, @@ -74,6 +73,7 @@ transform_points, ) from scripts.tutorials.atomic_action.tutorial_utils import ( + TutorialRobot, broadcast_pose_batch, clone_local_pose_from_first_env, create_toppra_motion_generator, @@ -212,21 +212,18 @@ def parse_arguments() -> argparse.Namespace: return parser.parse_args() -def create_dual_ur5_robot(sim: SimulationManager) -> Robot: - """Create a dual-UR5 robot with one PGI gripper on each arm.""" - return add_dual_ur5_robot( +def create_dual_robot( + sim: SimulationManager, + robot_type: TutorialRobot, +) -> Robot: + """Create the selected dual-arm robot with one PGI gripper per arm.""" + return add_dual_tutorial_robot( sim, - uid="DualUR5CoordinatedPlacement", - urdf_name="dual_ur5_coordinated_placement", - arm_urdf_path=resolve_cached_data_path("UniversalRobots/UR5/UR5.urdf"), - gripper_urdf_path=resolve_cached_data_path("DH_PGI_140_80/DH_PGI_140_80.urdf"), - solver_cfg=make_dual_ur5_solver_cfg( - GRIPPER_TCP_Z, - clear_urdf_path=True, - ), + robot_type=robot_type, + uid=f"Dual{robot_type.title()}CoordinatedPlacement", + urdf_name=f"dual_{robot_type}_coordinated_placement", + tcp_z=GRIPPER_TCP_Z, init_pos=ROBOT_INIT_POS, - joint_name_case="lower", - set_urdf_name_case=False, ) @@ -631,7 +628,10 @@ def run_coordinated_placement_demo( manipulators={"primary": "left_arm"}, end_effectors={"primary": "left_hand"}, ), - motion_policy=MotionPolicy(sample_count=PICK_SAMPLE_INTERVAL), + motion_policy=MotionPolicy( + strategy="motion_gen", + sample_count=PICK_SAMPLE_INTERVAL, + ), skill_options=left_pick_options, ), ActionInvocation( @@ -644,7 +644,10 @@ def run_coordinated_placement_demo( manipulators={"primary": "right_arm"}, end_effectors={"primary": "right_hand"}, ), - motion_policy=MotionPolicy(sample_count=PAN_PICK_SAMPLE_INTERVAL), + motion_policy=MotionPolicy( + strategy="motion_gen", + sample_count=PAN_PICK_SAMPLE_INTERVAL, + ), skill_options=right_pick_options, ), ) @@ -799,7 +802,10 @@ def log_trajectory_execution(step_idx: int, total_steps: int) -> None: "support": "right_hand", }, ), - motion_policy=MotionPolicy(sample_count=COORDINATED_SAMPLE_INTERVAL), + motion_policy=MotionPolicy( + strategy="motion_gen", + sample_count=COORDINATED_SAMPLE_INTERVAL, + ), skill_options=coordinated_options, ), ), @@ -867,7 +873,7 @@ def main() -> None: arena_space=3.0, light_pos=(0.0, -0.4, 3.0), ) - robot = create_dual_ur5_robot(sim) + robot = create_dual_robot(sim, args.robot) run_coordinated_placement_demo(args, sim, robot) diff --git a/scripts/tutorials/atomic_action/dynamic_obstacle_recovery.py b/scripts/tutorials/atomic_action/dynamic_obstacle_recovery.py index fc679530d..1f0ceb1e6 100644 --- a/scripts/tutorials/atomic_action/dynamic_obstacle_recovery.py +++ b/scripts/tutorials/atomic_action/dynamic_obstacle_recovery.py @@ -30,7 +30,6 @@ import torch -from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim import SimulationManager, VisualMaterialCfg from embodichain.lab.sim.atomic_actions import ( ActionBinding, @@ -56,11 +55,12 @@ CuroboPlannerCfg, CuroboWorldCfg, ) -from embodichain.lab.sim.robots import FrankaPandaCfg from embodichain.lab.sim.shapes import CubeCfg from embodichain.lab.visualization import SceneOverlays, TrajectoryOverlay from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( + add_tutorial_robot, + create_tutorial_argument_parser, create_tutorial_simulation, draw_axis_marker, prepare_tutorial_scene, @@ -71,7 +71,6 @@ stop_auto_play_recording, ) -ROBOT_UID = "dynamic_scene_franka" OBSTACLE_UID = "dynamic_obstacle" CONTROL_PART = "arm" SAMPLE_COUNT = 80 @@ -373,11 +372,9 @@ def _publish_path_overlays( def parse_arguments() -> argparse.Namespace: """Parse command-line arguments for the dynamic-obstacle tutorial.""" - parser = argparse.ArgumentParser( - description="Demonstrate collision-world revision recovery with cuRobo." + parser = create_tutorial_argument_parser( + "Demonstrate collision-world revision recovery with cuRobo." ) - add_env_launcher_args_to_parser(parser) - parser.add_argument("--auto_play", action="store_true") parser.add_argument( "--no_obstacle_motion", action="store_true", @@ -390,9 +387,7 @@ def main() -> None: """Move an obstacle during execution and replan from the latest snapshot.""" args = parse_arguments() sim = create_tutorial_simulation(args) - robot = sim.add_robot( - cfg=FrankaPandaCfg.from_dict({"uid": ROBOT_UID, "robot_type": "panda"}) - ) + robot = add_tutorial_robot(sim, args.robot) obstacle = sim.add_rigid_object( cfg=RigidObjectCfg( uid=OBSTACLE_UID, @@ -417,7 +412,7 @@ def main() -> None: motion_gen = MotionGenerator( MotionGenCfg( planner_cfg=CuroboPlannerCfg( - robot_uid=ROBOT_UID, + robot_uid=robot.uid, # The coarse default voxel fit under-covers the hand and # fingertips. A denser fit plus modest padding matches the # physical gripper without making the arm path infeasible. diff --git a/scripts/tutorials/atomic_action/hand_over.py b/scripts/tutorials/atomic_action/hand_over.py index bec85110e..5c75bbb73 100644 --- a/scripts/tutorials/atomic_action/hand_over.py +++ b/scripts/tutorials/atomic_action/hand_over.py @@ -50,12 +50,12 @@ from embodichain.lab.sim.shapes import MeshCfg from embodichain.utils import logger from scripts.tutorials.atomic_action.scenario_utils import ( - add_dual_ur5_robot, + add_dual_tutorial_robot, add_support_surface, - make_dual_ur5_solver_cfg, settle_object, ) from scripts.tutorials.atomic_action.tutorial_utils import ( + TutorialRobot, create_antipodal_semantics, create_toppra_motion_generator, create_tutorial_argument_parser, @@ -65,6 +65,7 @@ prepare_tutorial_scene, publish_tutorial_scene, replay_trajectory, + run_tutorial, serve_tutorial_scene, ) @@ -81,7 +82,7 @@ # --- Adjustable scene placeholders ----------------------------------------- # The object starts on the left side, is handed over at a lifted middle pose, # and is delivered to the right side. Tweak these to match the mesh geometry -# and the dual-UR5 reach. +# and the selected dual-arm robot's reach. OBJECT_INIT_XY = (0.0, 0.02) MIDDLE_OBJECT_XYZ = (0.0, 0.02, 0.82) MIDDLE_OBJECT_YAW_DEG = 0.0 @@ -119,16 +120,18 @@ def parse_arguments() -> argparse.Namespace: return parser.parse_args() -def create_dual_ur5_robot(sim: SimulationManager) -> Robot: - """Create a dual-UR5 robot with one PGI gripper on each arm.""" - return add_dual_ur5_robot( +def create_dual_robot( + sim: SimulationManager, + robot_type: TutorialRobot, +) -> Robot: + """Create the selected dual-arm robot with one PGI gripper per arm.""" + return add_dual_tutorial_robot( sim, - uid="DualUR5HandOver", - urdf_name="dual_ur5_hand_over", - solver_cfg=make_dual_ur5_solver_cfg( - GRIPPER_TCP_Z, - ur_ik_nearest_weight=(1.0, 4.0, 1.0, 1.0, 1.0, 1.0), - ), + robot_type=robot_type, + uid=f"Dual{robot_type.title()}HandOver", + urdf_name=f"dual_{robot_type}_hand_over", + tcp_z=GRIPPER_TCP_Z, + ur_ik_nearest_weight=(1.0, 4.0, 1.0, 1.0, 1.0, 1.0), hand_stiffness=1e2, hand_damping=1e1, hand_max_effort=1e3, @@ -267,7 +270,10 @@ def run_handover_demo( manipulators={"primary": "left_arm"}, end_effectors={"primary": "left_hand"}, ), - MotionPolicy(sample_count=PICKUP_SAMPLE_INTERVAL), + MotionPolicy( + strategy="motion_gen", + sample_count=PICKUP_SAMPLE_INTERVAL, + ), skill_options=pick_up_options, ), ActionInvocation( @@ -283,7 +289,10 @@ def run_handover_demo( "destination": "right_hand", }, ), - MotionPolicy(sample_count=HANDOVER_SAMPLE_INTERVAL), + MotionPolicy( + strategy="motion_gen", + sample_count=HANDOVER_SAMPLE_INTERVAL, + ), skill_options=handover_options, ), ) @@ -324,13 +333,10 @@ def main() -> None: arena_space=3.0, light_pos=(0.0, -0.4, 3.0), ) - robot = create_dual_ur5_robot(sim) - try: - run_handover_demo(args, sim, robot) - serve_tutorial_scene(sim, args) - finally: - sim.destroy() + robot = create_dual_robot(sim, args.robot) + run_handover_demo(args, sim, robot) + serve_tutorial_scene(sim, args) if __name__ == "__main__": - main() + run_tutorial(main) diff --git a/scripts/tutorials/atomic_action/move_end_effector.py b/scripts/tutorials/atomic_action/move_end_effector.py index 6e777a599..6beb89f8e 100644 --- a/scripts/tutorials/atomic_action/move_end_effector.py +++ b/scripts/tutorials/atomic_action/move_end_effector.py @@ -37,14 +37,13 @@ ) from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( - add_ur5_gripper_robot, + add_tutorial_robot, broadcast_pose_batch, broadcast_waypoint_pose_batch, - create_toppra_motion_generator, + create_curobo_motion_generator, create_tutorial_argument_parser, create_tutorial_simulation, draw_axis_marker, - make_top_down_eef_pose, prepare_tutorial_scene, replay_trajectory, run_tutorial, @@ -67,18 +66,19 @@ def main() -> None: """Move the robot end effector through two pose waypoints.""" args = parse_arguments() sim = create_tutorial_simulation(args) - robot = add_ur5_gripper_robot(sim) - motion_gen = create_toppra_motion_generator(robot) + robot = add_tutorial_robot(sim, args.robot) + motion_gen = create_curobo_motion_generator(robot) engine = AtomicActionEngine(motion_generator=motion_gen) - poses = torch.stack( - [ - make_top_down_eef_pose( - torch.tensor([0.30, -0.20, 0.36], device=sim.device) - ), - make_top_down_eef_pose(torch.tensor([0.45, 0.10, 0.30], device=sim.device)), - ] + start_pose = robot.compute_fk( + robot.get_qpos(name="arm"), name="arm", to_matrix=True + )[0] + poses = start_pose.unsqueeze(0).repeat(2, 1, 1) + poses[:, :3, 3] += torch.tensor( + [[-0.08, -0.08, 0.08], [0.04, 0.12, 0.04]], + dtype=poses.dtype, + device=poses.device, ) num_envs = robot.get_qpos().shape[0] if not args.no_vis_eef_axis: @@ -100,7 +100,10 @@ def main() -> None: broadcast_waypoint_pose_batch(poses, num_envs) ), binding=ActionBinding(manipulators={"primary": "arm"}), - motion_policy=MotionPolicy(sample_count=MOVE_SAMPLE_INTERVAL), + motion_policy=MotionPolicy( + strategy="motion_gen", + sample_count=MOVE_SAMPLE_INTERVAL, + ), ), ) ) diff --git a/scripts/tutorials/atomic_action/move_held_object.py b/scripts/tutorials/atomic_action/move_held_object.py index 6b0e1bf24..7c5f85f6d 100644 --- a/scripts/tutorials/atomic_action/move_held_object.py +++ b/scripts/tutorials/atomic_action/move_held_object.py @@ -45,11 +45,11 @@ from embodichain.lab.sim.shapes import MeshCfg from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( - add_ur5_gripper_robot, + add_tutorial_robot, broadcast_pose_batch, clone_local_pose_from_first_env, create_antipodal_semantics, - create_toppra_motion_generator, + create_curobo_motion_generator, create_tutorial_argument_parser, create_tutorial_simulation, draw_axis_marker, @@ -117,9 +117,9 @@ def main() -> None: """Plan MoveEndEffector -> PickUp -> MoveHeldObject.""" args = parse_arguments() sim = create_tutorial_simulation(args) - robot = add_ur5_gripper_robot(sim) + robot = add_tutorial_robot(sim, args.robot) obj = create_pick_object(sim) - motion_gen = create_toppra_motion_generator(robot) + motion_gen = create_curobo_motion_generator(robot) hand_open, hand_close = get_hand_open_close_qpos(robot) engine = AtomicActionEngine( @@ -158,13 +158,19 @@ def main() -> None: "move_end_effector", EndEffectorPoseGoal(move_target), binding, - MotionPolicy(sample_count=MOVE_SAMPLE_INTERVAL), + MotionPolicy( + strategy="motion_gen", + sample_count=MOVE_SAMPLE_INTERVAL, + ), ), ActionInvocation( "pick_up", GraspGoal(semantics), binding, - MotionPolicy(sample_count=PICK_SAMPLE_INTERVAL), + MotionPolicy( + strategy="motion_gen", + sample_count=PICK_SAMPLE_INTERVAL, + ), skill_options=PickUpOptions( pre_grasp_distance=0.15, lift_height=0.16, @@ -175,7 +181,10 @@ def main() -> None: "move_held_object", HeldObjectPoseGoal(object_target), binding, - MotionPolicy(sample_count=MOVE_HELD_OBJECT_SAMPLE_INTERVAL), + MotionPolicy( + strategy="motion_gen", + sample_count=MOVE_HELD_OBJECT_SAMPLE_INTERVAL, + ), ), ) ) diff --git a/scripts/tutorials/atomic_action/move_joints.py b/scripts/tutorials/atomic_action/move_joints.py index 7fbee6948..2c3ef3de1 100644 --- a/scripts/tutorials/atomic_action/move_joints.py +++ b/scripts/tutorials/atomic_action/move_joints.py @@ -38,8 +38,8 @@ ) from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( - add_ur5_gripper_robot, - create_toppra_motion_generator, + add_tutorial_robot, + create_curobo_motion_generator, create_tutorial_argument_parser, create_tutorial_simulation, draw_axis_marker, @@ -65,17 +65,22 @@ def main() -> None: """Move the robot arm through a named target and two explicit waypoints.""" args = parse_arguments() sim = create_tutorial_simulation(args) - robot = add_ur5_gripper_robot(sim) - motion_gen = create_toppra_motion_generator(robot) - - ready, mid, home = ( - torch.tensor(qpos, dtype=torch.float32, device=sim.device) - for qpos in ( - [0.35, -1.20, 1.30, -1.65, -1.57, 0.20], - [0.15, -1.40, 1.45, -1.60, -1.57, 0.10], - [0.0, -1.57, 1.57, -1.57, -1.57, 0.0], + robot = add_tutorial_robot(sim, args.robot) + motion_gen = create_curobo_motion_generator(robot) + + home = robot.get_qpos(name="arm")[0].clone() + limits = robot.get_qpos_limits(name="arm")[0] + + def offset_from_home(offsets: tuple[float, ...]) -> torch.Tensor: + target = home.clone() + count = min(target.numel(), len(offsets)) + target[:count] += torch.tensor( + offsets[:count], dtype=target.dtype, device=target.device ) - ) + return torch.minimum(torch.maximum(target, limits[:, 0]), limits[:, 1]) + + ready = offset_from_home((0.35, 0.37, -0.27, -0.08, 0.0, 0.20)) + mid = offset_from_home((0.15, 0.17, -0.12, -0.03, 0.0, 0.10)) engine = AtomicActionEngine( motion_generator=motion_gen, control_profiles={ @@ -96,7 +101,10 @@ def main() -> None: torch.stack([mid, home]).unsqueeze(0).repeat(robot.get_qpos().shape[0], 1, 1) ) binding = ActionBinding(manipulators={"primary": "arm"}) - policy = MotionPolicy(sample_count=MOVE_JOINTS_SAMPLE_INTERVAL) + policy = MotionPolicy( + strategy="motion_gen", + sample_count=MOVE_JOINTS_SAMPLE_INTERVAL, + ) compiled = engine.compile( ( ActionInvocation( diff --git a/scripts/tutorials/atomic_action/moving_target_recovery.py b/scripts/tutorials/atomic_action/moving_target_recovery.py index a34efec07..0a4546b9e 100644 --- a/scripts/tutorials/atomic_action/moving_target_recovery.py +++ b/scripts/tutorials/atomic_action/moving_target_recovery.py @@ -59,8 +59,8 @@ from embodichain.lab.sim.shapes import CubeCfg from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( - add_ur5_gripper_robot, - create_toppra_motion_generator, + add_tutorial_robot, + create_curobo_motion_generator, create_tutorial_argument_parser, create_tutorial_simulation, draw_axis_marker, @@ -243,7 +243,7 @@ def main() -> None: """Replan a late-bound PickUp request and lift the relocated cube.""" args = parse_arguments() sim = create_tutorial_simulation(args) - robot = add_ur5_gripper_robot(sim) + robot = add_tutorial_robot(sim, args.robot) target = _create_moving_target(sim) sim.update(step=10) target_scene = _MovingTargetScene(target, MOVED_TARGET_POSITION) @@ -252,7 +252,7 @@ def main() -> None: robot, scene_supplier=target_scene.snapshot, ) - motion_gen = create_toppra_motion_generator(robot) + motion_gen = create_curobo_motion_generator(robot) hand_open, hand_close = get_hand_open_close_qpos(robot) initialize_pre_pick_robot_pose(robot, target, hand_open) if args.no_target_motion: @@ -301,6 +301,7 @@ def main() -> None: ), binding=binding, motion_policy=MotionPolicy( + strategy="motion_gen", sample_count=PICK_SAMPLE_COUNT, control_dt=2.0 * sim_runtime.physics_dt, ), diff --git a/scripts/tutorials/atomic_action/pickup.py b/scripts/tutorials/atomic_action/pickup.py index a9c361671..4751dc4f5 100644 --- a/scripts/tutorials/atomic_action/pickup.py +++ b/scripts/tutorials/atomic_action/pickup.py @@ -42,10 +42,10 @@ from embodichain.lab.sim.shapes import CubeCfg from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( - add_ur5_gripper_robot, + add_tutorial_robot, clone_local_pose_from_first_env, create_antipodal_semantics, - create_toppra_motion_generator, + create_curobo_motion_generator, create_tutorial_argument_parser, create_tutorial_simulation, draw_axis_marker, @@ -126,11 +126,11 @@ def main() -> None: """Plan and replay a sampled antipodal PickUp trajectory.""" args = parse_arguments() sim = create_tutorial_simulation(args) - robot = add_ur5_gripper_robot(sim) + robot = add_tutorial_robot(sim, args.robot) obj = create_pick_object(sim) hand_open, hand_close = get_hand_open_close_qpos(robot) initialize_pre_pick_robot_pose(robot, obj, hand_open) - motion_gen = create_toppra_motion_generator(robot) + motion_gen = create_curobo_motion_generator(robot) engine = AtomicActionEngine( motion_generator=motion_gen, @@ -162,7 +162,10 @@ def main() -> None: manipulators={"primary": "arm"}, end_effectors={"primary": "hand"}, ), - motion_policy=MotionPolicy(sample_count=PICK_SAMPLE_INTERVAL), + motion_policy=MotionPolicy( + strategy="motion_gen", + sample_count=PICK_SAMPLE_INTERVAL, + ), skill_options=PickUpOptions( approach_direction=resolve_approach_direction(args, sim.device), pre_grasp_distance=0.15, diff --git a/scripts/tutorials/atomic_action/place.py b/scripts/tutorials/atomic_action/place.py index ae30526ad..7465d5575 100644 --- a/scripts/tutorials/atomic_action/place.py +++ b/scripts/tutorials/atomic_action/place.py @@ -44,12 +44,12 @@ from embodichain.lab.sim.shapes import CubeCfg from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( - add_ur5_gripper_robot, + add_tutorial_robot, broadcast_pose_batch, broadcast_waypoint_pose_batch, clone_local_pose_from_first_env, create_antipodal_semantics, - create_toppra_motion_generator, + create_curobo_motion_generator, create_tutorial_argument_parser, create_tutorial_simulation, draw_axis_marker, @@ -124,9 +124,9 @@ def main() -> None: """Plan and replay PickUp followed by a multi-waypoint Place.""" args = parse_arguments() sim = create_tutorial_simulation(args) - robot = add_ur5_gripper_robot(sim) + robot = add_tutorial_robot(sim, args.robot) obj = create_pick_object(sim) - motion_gen = create_toppra_motion_generator(robot) + motion_gen = create_curobo_motion_generator(robot) hand_open, hand_close = get_hand_open_close_qpos(robot) initialize_pre_pick_robot_pose(robot, obj, hand_open) @@ -166,7 +166,10 @@ def main() -> None: "pick_up", GraspGoal(semantics), binding, - MotionPolicy(sample_count=PICK_SAMPLE_INTERVAL), + MotionPolicy( + strategy="motion_gen", + sample_count=PICK_SAMPLE_INTERVAL, + ), skill_options=PickUpOptions( pre_grasp_distance=0.15, lift_height=0.16, @@ -181,7 +184,10 @@ def main() -> None: ) ), binding, - MotionPolicy(sample_count=PLACE_SAMPLE_INTERVAL), + MotionPolicy( + strategy="motion_gen", + sample_count=PLACE_SAMPLE_INTERVAL, + ), skill_options=PlaceOptions( lift_height=PLACE_LIFT_HEIGHT, hand_interp_steps=HAND_INTERP_STEPS, diff --git a/scripts/tutorials/atomic_action/press.py b/scripts/tutorials/atomic_action/press.py index a2e066a9b..ca30a1927 100644 --- a/scripts/tutorials/atomic_action/press.py +++ b/scripts/tutorials/atomic_action/press.py @@ -47,9 +47,9 @@ from embodichain.lab.sim.shapes import CubeCfg from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( - add_ur5_gripper_robot, + add_tutorial_robot, broadcast_pose_batch, - create_toppra_motion_generator, + create_curobo_motion_generator, create_tutorial_argument_parser, create_tutorial_simulation, draw_axis_marker, @@ -150,7 +150,7 @@ def main() -> None: """Plan, verify, and replay MoveEndEffector followed by Press.""" args = parse_arguments() sim = create_tutorial_simulation(args) - robot = add_ur5_gripper_robot(sim) + robot = add_tutorial_robot(sim, args.robot) block = create_wooden_block(sim, [*args.block_pos, 0.5 * BLOCK_SIZE[2]]) if sim.device.type == "cuda": sim.init_gpu_physics() @@ -158,7 +158,7 @@ def main() -> None: sim.update(step=5) block.clear_dynamics() - motion_gen = create_toppra_motion_generator(robot) + motion_gen = create_curobo_motion_generator(robot) hand_open, hand_close = get_hand_open_close_qpos(robot) engine = AtomicActionEngine( motion_generator=motion_gen, @@ -195,13 +195,19 @@ def main() -> None: "move_end_effector", EndEffectorPoseGoal(move_target), binding, - MotionPolicy(sample_count=MOVE_SAMPLE_INTERVAL), + MotionPolicy( + strategy="motion_gen", + sample_count=MOVE_SAMPLE_INTERVAL, + ), ), ActionInvocation( "press", PressGoal(press_target), binding, - MotionPolicy(sample_count=PRESS_SAMPLE_INTERVAL), + MotionPolicy( + strategy="motion_gen", + sample_count=PRESS_SAMPLE_INTERVAL, + ), skill_options=PressOptions( hand_interp_steps=HAND_INTERP_STEPS, ), diff --git a/scripts/tutorials/atomic_action/scenario_utils.py b/scripts/tutorials/atomic_action/scenario_utils.py index b51aefb91..4fb0f8fda 100644 --- a/scripts/tutorials/atomic_action/scenario_utils.py +++ b/scripts/tutorials/atomic_action/scenario_utils.py @@ -39,8 +39,14 @@ ) from embodichain.lab.sim.objects import RigidObject, Robot from embodichain.lab.sim.shapes import CubeCfg +from embodichain.lab.sim.robots import build_dual_arm_cfg from embodichain.lab.sim.solvers import PytorchSolverCfg, SolverCfg, URSolverCfg from embodichain.utils import logger +from scripts.tutorials.atomic_action.tutorial_utils import ( + GRIPPER_HAND_JOINT_PATTERN, + TutorialRobot, + create_tutorial_robot_cfg, +) ARM_URDF_PATH = "UniversalRobots/UR5/UR5.urdf" GRIPPER_URDF_PATH = "DH_PGI_140_80/DH_PGI_140_80.urdf" @@ -210,6 +216,190 @@ def add_dual_ur5_robot( return sim.add_robot(cfg=cfg) +def create_dual_tutorial_robot_cfg( + *, + robot_type: TutorialRobot, + uid: str, + urdf_name: str, + tcp_z: float, + solver: Literal["ur", "pytorch"] = "ur", + ur_ik_nearest_weight: Sequence[float] | None = None, + pytorch_num_samples: int = 30, + init_pos: Sequence[float] = DUAL_UR5_INIT_POS, + init_rot: Sequence[float] = DUAL_UR5_INIT_ROT, + left_arm_home: Sequence[float] | None = None, + right_arm_home: Sequence[float] | None = None, + hand_stiffness: float = 1e3, + hand_damping: float = 1e2, + hand_max_effort: float = 1e4, +) -> RobotCfg: + """Build a dual tutorial robot from the selected arm and shared PGI hand. + + Franka always uses its PyTorch kinematics solver; ``solver="ur"`` selects + the analytical solver only when ``robot_type="ur5"``. The mounting layout, + control-part names, gripper component, and downstream action bindings stay + identical across both robot choices. + + Args: + robot_type: Arm family to mount on both sides. + uid: Simulation robot identifier. + urdf_name: Cache name for the assembled dual-arm URDF. + tcp_z: PGI tool-center-point offset along local Z. + solver: Preferred UR5 solver implementation. + ur_ik_nearest_weight: Optional nearest-solution weights for UR5 IK. + pytorch_num_samples: Number of PyTorch IK seed samples. + init_pos: Root position of the assembled robot. + init_rot: Root xyz Euler rotation in degrees. + left_arm_home: Optional left-arm initial configuration. + right_arm_home: Optional right-arm initial configuration. + hand_stiffness: PGI joint drive stiffness. + hand_damping: PGI joint drive damping. + hand_max_effort: PGI joint maximum effort. + + Returns: + A dual-arm robot configuration with two PGI grippers. + """ + base_cfg = create_tutorial_robot_cfg(robot_type) + tcp = [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, tcp_z], + [0.0, 0.0, 0.0, 1.0], + ] + base_solver = base_cfg.solver_cfg["arm"] + if robot_type == "ur5" and solver == "ur": + base_solver.tcp = tcp + base_solver.ik_nearest_weight = ur_ik_nearest_weight + else: + base_solver = PytorchSolverCfg( + end_link_name=base_solver.end_link_name, + root_link_name=base_solver.root_link_name, + tcp=tcp, + num_samples=pytorch_num_samples, + ) + base_cfg.solver_cfg["arm"] = base_solver + + for property_name, value in ( + ("stiffness", hand_stiffness), + ("damping", hand_damping), + ("max_effort", hand_max_effort), + ): + getattr(base_cfg.drive_pros, property_name)[GRIPPER_HAND_JOINT_PATTERN] = value + + arm_facing_rotation = make_yaw_transform( + (0.0, 0.0, 0.0), + math.radians(float(base_cfg.init_rot[2])), + ) + mounts = { + "left": make_yaw_transform((-0.3, -1.45, 0.4), np.pi / 2) @ arm_facing_rotation, + "right": make_yaw_transform((0.3, -1.45, 0.4), np.pi / 2) @ arm_facing_rotation, + } + cfg = build_dual_arm_cfg(base_cfg, mounts) + + # ``build_dual_arm_cfg`` duplicates the arm component and all control + # parts. The tutorial robots intentionally keep the PGI as a separate URDF + # component, so mount one copy on each assembled arm as well. + cfg.urdf_cfg.fname = urdf_name + hand_component = base_cfg.urdf_cfg.components["hand"] + for side in ("left", "right"): + cfg.urdf_cfg.add_component( + f"{side}_hand", + hand_component["urdf_path"], + hand_component["transform"], + **hand_component.get("params", {}), + ) + + arm_dof = len(base_cfg.control_parts["arm"]) + base_arm_home = list(base_cfg.init_qpos[:arm_dof]) + base_hand_home = list(base_cfg.init_qpos[arm_dof:]) + if left_arm_home is None: + left_arm_home = base_arm_home + if right_arm_home is None: + right_arm_home = base_arm_home + if len(left_arm_home) != arm_dof or len(right_arm_home) != arm_dof: + raise ValueError( + f"Dual {robot_type} arm homes must each contain {arm_dof} joints." + ) + + cfg.uid = uid + cfg.init_pos = list(init_pos) + cfg.init_rot = list(init_rot) + # DexSim traverses the two arm branches breadth-first, so their active + # joints appear left/right interleaved even though the URDF components are + # emitted one after the other. Match that runtime order before appending + # the two gripper components. + cfg.init_qpos = ( + [ + qpos + for joint_pair in zip(left_arm_home, right_arm_home, strict=True) + for qpos in joint_pair + ] + + base_hand_home + + base_hand_home + ) + return cfg + + +def add_dual_tutorial_robot( + sim: SimulationManager, + *, + robot_type: TutorialRobot, + uid: str, + urdf_name: str, + tcp_z: float, + solver: Literal["ur", "pytorch"] = "ur", + ur_ik_nearest_weight: Sequence[float] | None = None, + pytorch_num_samples: int = 30, + init_pos: Sequence[float] = DUAL_UR5_INIT_POS, + init_rot: Sequence[float] = DUAL_UR5_INIT_ROT, + left_arm_home: Sequence[float] | None = None, + right_arm_home: Sequence[float] | None = None, + hand_stiffness: float = 1e3, + hand_damping: float = 1e2, + hand_max_effort: float = 1e4, +) -> Robot: + """Add a dual UR5 or Franka tutorial robot to a simulation. + + Args: + sim: Simulation manager that owns the robot. + robot_type: Arm family to mount on both sides. + uid: Simulation robot identifier. + urdf_name: Cache name for the assembled dual-arm URDF. + tcp_z: PGI tool-center-point offset along local Z. + solver: Preferred UR5 solver implementation. + ur_ik_nearest_weight: Optional nearest-solution weights for UR5 IK. + pytorch_num_samples: Number of PyTorch IK seed samples. + init_pos: Root position of the assembled robot. + init_rot: Root xyz Euler rotation in degrees. + left_arm_home: Optional left-arm initial configuration. + right_arm_home: Optional right-arm initial configuration. + hand_stiffness: PGI joint drive stiffness. + hand_damping: PGI joint drive damping. + hand_max_effort: PGI joint maximum effort. + + Returns: + The added dual-arm robot instance. + """ + return sim.add_robot( + cfg=create_dual_tutorial_robot_cfg( + robot_type=robot_type, + uid=uid, + urdf_name=urdf_name, + tcp_z=tcp_z, + solver=solver, + ur_ik_nearest_weight=ur_ik_nearest_weight, + pytorch_num_samples=pytorch_num_samples, + init_pos=init_pos, + init_rot=init_rot, + left_arm_home=left_arm_home, + right_arm_home=right_arm_home, + hand_stiffness=hand_stiffness, + hand_damping=hand_damping, + hand_max_effort=hand_max_effort, + ) + ) + + def add_support_surface( sim: SimulationManager, *, @@ -333,11 +523,13 @@ def log_action_plan( "GRIPPER_URDF_PATH", "LEFT_ARM_HOME", "RIGHT_ARM_HOME", + "add_dual_tutorial_robot", "add_dual_ur5_robot", "add_support_surface", "compute_local_bounds", "compute_world_bounds", "create_manual_object_semantics", + "create_dual_tutorial_robot_cfg", "get_local_vertices", "invert_pose", "log_action_plan", diff --git a/scripts/tutorials/atomic_action/tutorial_utils.py b/scripts/tutorials/atomic_action/tutorial_utils.py index cf8dc0731..b1ddd5a2b 100644 --- a/scripts/tutorials/atomic_action/tutorial_utils.py +++ b/scripts/tutorials/atomic_action/tutorial_utils.py @@ -36,8 +36,13 @@ ) from embodichain.lab.sim.cfg import LightCfg, MarkerCfg, RenderCfg, RobotCfg from embodichain.lab.sim.objects import RigidObject, Robot -from embodichain.lab.sim.planners import MotionGenCfg, MotionGenerator, ToppraPlannerCfg -from embodichain.lab.sim.robots import URRobotCfg +from embodichain.lab.sim.planners import ( + CuroboPlannerCfg, + MotionGenCfg, + MotionGenerator, + ToppraPlannerCfg, +) +from embodichain.lab.sim.robots import FrankaPandaCfg, URRobotCfg from embodichain.lab.sim.solvers import URSolverCfg from embodichain.toolkits.graspkit.pg_grasp.antipodal_generator import ( AntipodalSamplerCfg, @@ -72,6 +77,13 @@ GRIPPER_Y_THICKNESS = 0.040 DEFAULT_GRIPPER_CLOSE_QPOS = 0.024 DEFAULT_TUTORIAL_LIGHT_POS = (1.0, 0.0, 3.0) +_FRANKA_TUTORIAL_BASE_ROTATION = (0.0, 0.0, 180.0) +_GRIPPER_TCP = ( + (1.0, 0.0, 0.0, 0.0), + (0.0, 1.0, 0.0, 0.0), + (0.0, 0.0, 1.0, GRIPPER_TCP_Z), + (0.0, 0.0, 0.0, 1.0), +) TOP_DOWN_EEF_ROTATION = ( (-0.0539, -0.9985, -0.0022), (-0.9977, 0.0540, -0.0401), @@ -85,6 +97,8 @@ "headless_play", "visualize_axes", ] +TutorialRobot = Literal["ur5", "franka"] +TUTORIAL_ROBOTS: tuple[TutorialRobot, ...] = ("ur5", "franka") def create_tutorial_argument_parser( @@ -110,6 +124,12 @@ def create_tutorial_argument_parser( action="store_true", help="Run the demo without waiting for keyboard input.", ) + parser.add_argument( + "--robot", + choices=TUTORIAL_ROBOTS, + default="ur5", + help="Robot construction to use (default: ur5).", + ) if "debug_state" in features: parser.add_argument( "--debug_state", @@ -241,6 +261,58 @@ def add_ur5_gripper_robot( ) +def add_franka_panda_robot( + sim: SimulationManager, + init_pos: Sequence[float] = (0.0, 0.0, 0.0), + init_qpos: Sequence[float] | None = None, +) -> Robot: + """Add a Franka arm with the standard PGI tutorial gripper. + + Args: + sim: Simulation manager that owns the robot. + init_pos: Root position of the robot in its arena. + init_qpos: Optional full robot joint configuration. + + Returns: + The added robot instance. + """ + return sim.add_robot( + cfg=create_franka_panda_robot_cfg( + init_pos=init_pos, + init_qpos=init_qpos, + ) + ) + + +def add_tutorial_robot( + sim: SimulationManager, + robot_type: TutorialRobot, + init_pos: Sequence[float] = (0.0, 0.0, 0.0), + init_qpos: Sequence[float] | None = None, +) -> Robot: + """Add a selected tutorial robot with the shared PGI gripper. + + Args: + sim: Simulation manager that owns the robot. + robot_type: Tutorial robot family to construct. + init_pos: Root position of the robot in its arena. + init_qpos: Optional full robot joint configuration. + + Returns: + The added robot instance. + + Raises: + ValueError: If ``robot_type`` is not supported. + """ + return sim.add_robot( + cfg=create_tutorial_robot_cfg( + robot_type, + init_pos=init_pos, + init_qpos=init_qpos, + ) + ) + + def create_toppra_motion_generator(robot: Robot) -> MotionGenerator: """Create the standard TOPPRA motion generator for a tutorial robot. @@ -255,13 +327,27 @@ def create_toppra_motion_generator(robot: Robot) -> MotionGenerator: ) +def create_curobo_motion_generator(robot: Robot) -> MotionGenerator: + """Create a cuRobo-backed motion generator for a tutorial robot. + + Args: + robot: Robot whose trajectories will be planned. + + Returns: + The configured motion generator with an empty external collision world. + """ + return MotionGenerator( + cfg=MotionGenCfg(planner_cfg=CuroboPlannerCfg(robot_uid=robot.uid)) + ) + + def get_hand_open_close_qpos( robot: Robot, *, hand_control_part: str = "hand", close_qpos: float = DEFAULT_GRIPPER_CLOSE_QPOS, ) -> tuple[torch.Tensor, torch.Tensor]: - """Return the open limit and a safe closed position for a gripper. + """Return the open limit and a safe closed position for a PGI gripper. Args: robot: Robot containing the gripper control part. @@ -389,7 +475,7 @@ def initialize_pre_pick_robot_pose( *, height: float = 0.36, ) -> None: - """Set a UR5 at a deterministic open-gripper pose above an object. + """Set a tutorial robot at a deterministic open-gripper pose above an object. Args: robot: Robot to initialize. @@ -809,22 +895,108 @@ def create_ur5_gripper_robot_cfg( GRIPPER_HAND_JOINT_PATTERN: 1e4, }, }, - "solver_cfg": { - "arm": { - "tcp": [ - [1.0, 0.0, 0.0, 0.0], - [0.0, 1.0, 0.0, 0.0], - [0.0, 0.0, 1.0, GRIPPER_TCP_Z], - [0.0, 0.0, 0.0, 1.0], - ] - } - }, + "solver_cfg": {"arm": {"tcp": _GRIPPER_TCP}}, "init_qpos": qpos, "init_pos": init_pos, } ) +def create_franka_panda_robot_cfg( + init_pos: Sequence[float] = (0.0, 0.0, 0.0), + init_qpos: Sequence[float] | None = None, +) -> RobotCfg: + """Build a Franka arm + PGI gripper configuration for the tutorials. + + The shared scenes place manipulation targets on the robot's negative-X + side, so the base is rotated 180 degrees. The arm-only Franka URDF is + assembled with the same DH_PGI_140_80 component, hand control part, drive + properties, open state, and TCP offset used by the UR5 tutorial robot. + + Args: + init_pos: Initial root position of the robot in the arena. + init_qpos: Optional full robot joint configuration. + + Returns: + A fully populated :class:`~embodichain.lab.sim.cfg.RobotCfg`. + """ + overrides = { + "robot_type": "panda", + "uid": "FrankaPanda", + "init_pos": init_pos, + "init_rot": _FRANKA_TUTORIAL_BASE_ROTATION, + "urdf_cfg": { + "components": [ + { + "component_type": "arm", + "urdf_path": "Franka/Panda/Panda.urdf", + }, + { + "component_type": "hand", + "urdf_path": GRIPPER_URDF_PATH, + }, + ], + }, + "control_parts": {"hand": [GRIPPER_HAND_JOINT_PATTERN]}, + "drive_pros": { + "stiffness": {GRIPPER_HAND_JOINT_PATTERN: 1e3}, + "damping": {GRIPPER_HAND_JOINT_PATTERN: 1e2}, + "max_effort": {GRIPPER_HAND_JOINT_PATTERN: 1e4}, + }, + "solver_cfg": { + "arm": { + "end_link_name": "fr3_link8", + "tcp": _GRIPPER_TCP, + } + }, + } + if init_qpos is not None: + overrides["init_qpos"] = list(init_qpos) + cfg = FrankaPandaCfg.from_dict(overrides) + if init_qpos is None: + cfg.init_qpos[-2:] = [0.0, 0.0] + for drive_values in ( + cfg.drive_pros.stiffness, + cfg.drive_pros.damping, + cfg.drive_pros.max_effort, + ): + drive_values.pop("fr3_finger_joint[1-2]", None) + return cfg + + +def create_tutorial_robot_cfg( + robot_type: TutorialRobot, + init_pos: Sequence[float] = (0.0, 0.0, 0.0), + init_qpos: Sequence[float] | None = None, +) -> RobotCfg: + """Build a selected tutorial arm with the common PGI gripper contract. + + Args: + robot_type: Tutorial robot family to construct. + init_pos: Initial root position of the robot in its arena. + init_qpos: Optional full robot joint configuration. + + Returns: + A UR5 or Franka robot configuration exposing ``arm`` and ``hand``. + + Raises: + ValueError: If ``robot_type`` is not supported. + """ + if robot_type == "ur5": + return create_ur5_gripper_robot_cfg( + init_pos=init_pos, + init_qpos=init_qpos, + ) + if robot_type == "franka": + return create_franka_panda_robot_cfg( + init_pos=init_pos, + init_qpos=init_qpos, + ) + raise ValueError( + f"Unsupported tutorial robot {robot_type!r}; expected one of {TUTORIAL_ROBOTS}." + ) + + __all__ = [ "DEFAULT_AUTO_PLAY_LOOK_AT", "DEFAULT_AXIS_LEN", @@ -836,13 +1008,20 @@ def create_ur5_gripper_robot_cfg( "GRIPPER_URDF_PATH", "TOP_DOWN_EEF_ROTATION", "TutorialCliFeature", + "TutorialRobot", + "TUTORIAL_ROBOTS", + "add_franka_panda_robot", + "add_tutorial_robot", "add_ur5_gripper_robot", "broadcast_pose_batch", "broadcast_waypoint_pose_batch", "clone_local_pose_from_first_env", "create_antipodal_semantics", + "create_curobo_motion_generator", + "create_franka_panda_robot_cfg", "create_toppra_motion_generator", "create_tutorial_argument_parser", + "create_tutorial_robot_cfg", "create_tutorial_simulation", "create_ur5_gripper_robot_cfg", "format_tensor", diff --git a/tests/sim/atomic_actions/test_trajectory_ops.py b/tests/sim/atomic_actions/test_trajectory_ops.py index 47e1b3cf6..a23384834 100644 --- a/tests/sim/atomic_actions/test_trajectory_ops.py +++ b/tests/sim/atomic_actions/test_trajectory_ops.py @@ -93,7 +93,7 @@ def unexpected_current_device(): with pytest.raises(ValueError, match="CUDA device requested"): normalize_success_mask( True, - n_envs=2, + num_envs=2, device="cuda", name="IK success", ) diff --git a/tests/sim/atomic_actions/test_tutorial_utils.py b/tests/sim/atomic_actions/test_tutorial_utils.py index 07d101939..def9e17a3 100644 --- a/tests/sim/atomic_actions/test_tutorial_utils.py +++ b/tests/sim/atomic_actions/test_tutorial_utils.py @@ -18,7 +18,9 @@ from __future__ import annotations +import importlib from argparse import Namespace +from types import SimpleNamespace from unittest.mock import MagicMock, call, patch import pytest @@ -31,11 +33,24 @@ _maximum_path_deviation, _minimum_cuboid_clearance, ) +from scripts.tutorials.atomic_action.coordinated_pickment import ( + compute_left_to_right_arm_direction, +) +from scripts.tutorials.atomic_action.scenario_utils import ( + create_dual_tutorial_robot_cfg, +) from scripts.tutorials.atomic_action.tutorial_utils import ( + TUTORIAL_ROBOTS, broadcast_pose_batch, broadcast_waypoint_pose_batch, clone_local_pose_from_first_env, create_antipodal_semantics, + create_curobo_motion_generator, + create_franka_panda_robot_cfg, + create_tutorial_argument_parser, + create_tutorial_robot_cfg, + create_ur5_gripper_robot_cfg, + get_hand_open_close_qpos, replay_trajectory, should_open_tutorial_window, should_wait_for_tutorial_input, @@ -46,6 +61,30 @@ Y_OFFSET = 0.18 EXPECTED_STEP_COUNT = 3 CUBOID_SIZE = (0.2, 0.2, 0.2) +FRANKA_TUTORIAL_BASE_ROTATION = (0.0, 0.0, 180.0) +DUAL_FRANKA_MOUNT_X_AXIS = torch.tensor([0.0, -1.0, 0.0]) +PGI_TUTORIAL_TCP = torch.tensor( + [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.15], + [0.0, 0.0, 0.0, 1.0], + ] +) +ATOMIC_ACTION_TUTORIAL_MODULES = ( + "assemble", + "coordinated_pickment", + "coordinated_placement", + "dynamic_obstacle_recovery", + "hand_over", + "move_end_effector", + "move_held_object", + "move_joints", + "moving_target_recovery", + "pickup", + "place", + "press", +) def _run_obstacle_animation(*, pace_wall_time: bool) -> tuple[MagicMock, MagicMock]: @@ -199,6 +238,173 @@ def test_create_antipodal_semantics_keeps_mesh_data_on_affordance() -> None: assert semantics.affordance.generator_cfg.antipodal_sampler_cfg.n_sample == 64 +def test_franka_tutorial_config_uses_ur5_gripper_component() -> None: + ur5_cfg = create_ur5_gripper_robot_cfg() + franka_cfg = create_franka_panda_robot_cfg() + + assert franka_cfg.urdf_cfg.components["hand"]["urdf_path"] == ( + ur5_cfg.urdf_cfg.components["hand"]["urdf_path"] + ) + assert franka_cfg.urdf_cfg.components["arm"]["urdf_path"].endswith( + "/Franka/Panda/Panda.urdf" + ) + assert franka_cfg.init_qpos[-2:] == [0.0, 0.0] + assert franka_cfg.init_rot == FRANKA_TUTORIAL_BASE_ROTATION + for property_name in ("stiffness", "damping", "max_effort"): + ur5_values = getattr(ur5_cfg.drive_pros, property_name) + franka_values = getattr(franka_cfg.drive_pros, property_name) + assert franka_values["gripper_finger1_joint_1"] == ( + ur5_values["gripper_finger1_joint_1"] + ) + assert "fr3_finger_joint[1-2]" not in franka_values + + +def test_ur5_and_franka_configs_share_place_binding_contract() -> None: + configs = ( + create_ur5_gripper_robot_cfg(), + create_franka_panda_robot_cfg(), + ) + + assert all(set(cfg.control_parts) == {"arm", "hand"} for cfg in configs) + assert all(set(cfg.solver_cfg) == {"arm"} for cfg in configs) + assert configs[0].control_parts["hand"] == configs[1].control_parts["hand"] + assert all( + torch.allclose(torch.as_tensor(cfg.solver_cfg["arm"].tcp), PGI_TUTORIAL_TCP) + for cfg in configs + ) + + +@pytest.mark.parametrize( + ("robot_type", "arm_dof", "solver_name"), + (("ur5", 6, "URSolverCfg"), ("franka", 7, "PytorchSolverCfg")), +) +def test_dual_tutorial_configs_share_pgi_binding_contract( + robot_type: str, + arm_dof: int, + solver_name: str, +) -> None: + single_cfg = create_tutorial_robot_cfg(robot_type) + dual_cfg = create_dual_tutorial_robot_cfg( + robot_type=robot_type, + uid=f"test_{robot_type}", + urdf_name=f"test_dual_{robot_type}", + tcp_z=0.121, + ) + expected_tcp = PGI_TUTORIAL_TCP.clone() + expected_tcp[2, 3] = 0.121 + + assert tuple(dual_cfg.urdf_cfg.components) == ( + "left_arm", + "right_arm", + "left_hand", + "right_hand", + ) + expected_arm_home = list(single_cfg.init_qpos[:arm_dof]) + assert dual_cfg.init_qpos[: 2 * arm_dof : 2] == expected_arm_home + assert dual_cfg.init_qpos[1 : 2 * arm_dof : 2] == expected_arm_home + for side in ("left", "right"): + assert len(dual_cfg.control_parts[f"{side}_arm"]) == arm_dof + assert dual_cfg.control_parts[f"{side}_hand"] == [ + f"{side}_gripper_finger1_joint_1" + ] + assert dual_cfg.urdf_cfg.components[f"{side}_hand"]["urdf_path"] == ( + single_cfg.urdf_cfg.components["hand"]["urdf_path"] + ) + solver = dual_cfg.solver_cfg[f"{side}_arm"] + assert type(solver).__name__ == solver_name + assert torch.allclose(torch.as_tensor(solver.tcp), expected_tcp) + + +def test_dual_franka_mount_preserves_single_arm_facing_direction() -> None: + cfg = create_dual_tutorial_robot_cfg( + robot_type="franka", + uid="test_franka_orientation", + urdf_name="test_dual_franka_orientation", + tcp_z=0.121, + ) + + for side in ("left", "right"): + mount = torch.as_tensor( + cfg.urdf_cfg.components[f"{side}_arm"]["transform"], + dtype=torch.float32, + ) + assert torch.allclose( + mount[:3, 0], + DUAL_FRANKA_MOUNT_X_AXIS, + atol=1e-6, + ) + + +def test_hand_commands_use_pgi_open_limit() -> None: + robot = MagicMock() + robot.device = torch.device("cpu") + robot.get_qpos_limits.return_value = torch.tensor([[[0.0, 0.04]]]) + + hand_open, hand_close = get_hand_open_close_qpos(robot) + + assert torch.allclose(hand_open, torch.tensor([0.0])) + assert torch.allclose(hand_close, torch.tensor([0.024])) + + +def test_curobo_motion_generator_factory_selects_curobo_backend() -> None: + robot = MagicMock(uid="tutorial_robot") + + with patch( + "scripts.tutorials.atomic_action.tutorial_utils.MotionGenerator" + ) as motion_generator_cls: + result = create_curobo_motion_generator(robot) + + cfg = motion_generator_cls.call_args.kwargs["cfg"] + assert result is motion_generator_cls.return_value + assert cfg.planner_cfg.planner_type == "curobo" + assert cfg.planner_cfg.robot_uid == "tutorial_robot" + + +def test_shared_robot_selection_keeps_ur5_default_and_accepts_franka() -> None: + parser = create_tutorial_argument_parser("test parser") + default_args = parser.parse_args([]) + franka_args = parser.parse_args(["--robot", "franka"]) + + assert TUTORIAL_ROBOTS == ("ur5", "franka") + assert default_args.robot == "ur5" + assert franka_args.robot == "franka" + + +def test_arm_direction_uses_selected_robot_solver_roots() -> None: + robot = MagicMock() + robot.cfg.solver_cfg = { + "left_arm": SimpleNamespace(root_link_name="left_franka_root"), + "right_arm": SimpleNamespace(root_link_name="right_franka_root"), + } + left_pose = torch.eye(4).unsqueeze(0) + right_pose = torch.eye(4).unsqueeze(0) + right_pose[0, 1, 3] = 2.0 + robot.get_link_pose.side_effect = (left_pose, right_pose) + + direction = compute_left_to_right_arm_direction(robot, "cpu") + + assert torch.allclose(direction, torch.tensor([0.0, 1.0, 0.0])) + assert robot.get_link_pose.call_args_list == [ + call(link_name="left_franka_root", env_ids=[0], to_matrix=True), + call(link_name="right_franka_root", env_ids=[0], to_matrix=True), + ] + + +@pytest.mark.parametrize("module_name", ATOMIC_ACTION_TUTORIAL_MODULES) +def test_all_atomic_action_tutorials_accept_both_robot_choices( + module_name: str, +) -> None: + module = importlib.import_module(f"scripts.tutorials.atomic_action.{module_name}") + + with patch("sys.argv", [f"{module_name}.py"]): + default_args = module.parse_arguments() + with patch("sys.argv", [f"{module_name}.py", "--robot", "franka"]): + franka_args = module.parse_arguments() + + assert default_args.robot == "ur5" + assert franka_args.robot == "franka" + + def test_replay_timed_trajectory_uses_arrival_intervals() -> None: sim = MagicMock() sim.sim_config.physics_dt = 0.1