Support Franka in atomic-action tutorials

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.

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.")
 

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,23 +55,24 @@
 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,
     get_hand_open_close_qpos,
     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)

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)