Add gizmo helper tool for robot control (#132)

Author: yuecideng

embodichain/lab/sim/sim_manager.py

@@ -1130,7 +1130,7 @@ def get_robot_uid_list(self) -> List[str]:
 
     def enable_gizmo(
         self, uid: str, control_part: str | None = None, gizmo_cfg: object = None
-    ) -> None:
+    ) -> Gizmo:
         """Enable gizmo control for any simulation object (Robot, RigidObject, Camera, etc.).
 
         Args:
@@ -1190,6 +1190,8 @@ def enable_gizmo(
                 f"Failed to create gizmo for {object_type} '{uid}' with control_part '{control_part}': {e}"
             )
 
+        return gizmo
+
     def disable_gizmo(self, uid: str, control_part: str | None = None) -> None:
         """Disable and remove gizmo for a robot.
 

embodichain/lab/sim/utility/gizmo_utils.py

@@ -21,8 +21,12 @@
 """
 
 from typing import Callable
+from typing import TYPE_CHECKING
 from dexsim.types import TransformMask
 
+if TYPE_CHECKING:
+    from embodichain.lab.sim.objects import Robot
+
 
 def create_gizmo_callback() -> Callable:
     """Create a standard gizmo transform callback function.
@@ -44,3 +48,176 @@ def gizmo_transform_callback(node, translation, rotation, flag):
                 node.set_rotation_rpy(rotation)
 
     return gizmo_transform_callback
+
+
+def run_gizmo_robot_control_loop(
+    robot: "Robot", control_part: str = "arm", end_link_name: str | None = None
+):
+    """Run a control loop for testing gizmo controls on a robot.
+
+    This function implements a control loop that allows users to manipulate a robot
+    using gizmo controls with keyboard input for additional commands.
+
+    Args:
+        robot (Robot): The robot to control with the gizmo.
+        control_part (str, optional): The part of the robot to control. Defaults to "arm".
+        end_link_name (str | None, optional): The name of the end link for FK calculations. Defaults to None.
+
+    Keyboard Controls:
+        Q/ESC: Exit the control loop
+        P: Print current robot state (joint positions, end-effector pose)
+        G: Toggle gizmo visibility
+        R: Reset robot to initial pose
+        I: Print control information
+    """
+    import select
+    import sys
+    import tty
+    import termios
+    import time
+    import numpy as np
+
+    np.set_printoptions(precision=5, suppress=True)
+
+    from embodichain.lab.sim import SimulationManager
+    from embodichain.lab.sim.objects import Robot
+    from embodichain.lab.sim.solvers import PinkSolverCfg
+
+    from embodichain.utils.logger import log_info, log_warning, log_error
+
+    sim = SimulationManager.get_instance()
+
+    # Enter auto-update mode.
+    sim.set_manual_update(False)
+
+    # Replace robot's default solver with PinkSolver for gizmo control.
+    robot_solver = robot.get_solver(name=control_part)
+    control_part_link_names = robot.get_control_part_link_names(name=control_part)
+    end_link_name = (
+        control_part_link_names[-1] if end_link_name is None else end_link_name
+    )
+    pink_solver_cfg = PinkSolverCfg(
+        urdf_path=robot.cfg.fpath,
+        end_link_name=end_link_name,
+        root_link_name=robot_solver.root_link_name,
+        pos_eps=1e-2,
+        rot_eps=5e-2,
+        max_iterations=300,
+        dt=0.1,
+    )
+    robot.init_solver(cfg={control_part: pink_solver_cfg})
+
+    # Enable gizmo for the robot
+    gizmo = sim.enable_gizmo(uid=robot.uid, control_part=control_part)
+
+    # Store initial robot configuration
+    initial_qpos = robot.get_qpos(name=control_part)
+
+    gizmo_visible = True
+
+    log_info("\n=== Gizmo Robot Control ===")
+    log_info("Gizmo Controls:")
+    log_info("  Use the 3D gizmo to drag and manipulate the robot")
+    log_info("\nKeyboard Controls:")
+    log_info("  Q/ESC: Exit control loop")
+    log_info("  P: Print current robot state")
+    log_info("  G: Toggle gizmo visibility")
+    log_info("  R: Reset robot to initial pose")
+    log_info("  I: Print this information again")
+
+    # Save terminal settings
+    old_settings = termios.tcgetattr(sys.stdin)
+    tty.setcbreak(sys.stdin.fileno())
+
+    def get_key():
+        """Non-blocking keyboard input."""
+        if select.select([sys.stdin], [], [], 0)[0]:
+            return sys.stdin.read(1)
+        return None
+
+    try:
+        while True:
+            time.sleep(0.033)  # ~30Hz
+            sim.update_gizmos()
+
+            # Check for keyboard input
+            key = get_key()
+
+            if key:
+                # Exit controls
+                if key in ["q", "Q", "\x1b"]:  # Q or ESC
+                    log_info("Exiting gizmo control loop...")
+                    sim.disable_gizmo(uid=robot.uid, control_part=control_part)
+                    if robot_solver:
+                        robot.init_solver(
+                            cfg={control_part: robot_solver.cfg}
+                        )  # Restore original solver
+                    break
+
+                # Print robot state
+                elif key in ["p", "P"]:
+                    current_qpos = robot.get_qpos(name=control_part)
+                    eef_pose = robot.compute_fk(name=control_part, qpos=current_qpos)
+                    log_info(f"\n=== Robot State ===")
+                    log_info(f"Control part: {control_part}")
+                    log_info(f"Joint positions: {current_qpos.squeeze().tolist()}")
+                    log_info(f"End-effector pose:\n{eef_pose.squeeze().numpy()}")
+
+                    if eef_pose is None:
+                        log_info(
+                            "End-effector pose unavailable: compute_fk returned None "
+                            f"for control part '{control_part}'."
+                        )
+                    else:
+                        eef_pose_np = eef_pose.detach().cpu().numpy().squeeze()
+                        log_info(f"End-effector pose:\n{eef_pose_np}")
+                elif key in ["g", "G"]:
+                    if gizmo_visible:
+                        sim.set_gizmo_visibility(
+                            uid=robot.uid, control_part=control_part, visible=False
+                        )
+                        log_info("Gizmo hidden")
+                        gizmo_visible = False
+                    else:
+                        sim.set_gizmo_visibility(
+                            uid=robot.uid, control_part=control_part, visible=True
+                        )
+                        log_info("Gizmo shown")
+                        gizmo_visible = True
+
+                # Reset to initial pose
+                elif key in ["r", "R"]:
+                    # TODO: Workaround for reset. Gizmo pose should be fixed in the future.
+                    sim.disable_gizmo(uid=robot.uid, control_part=control_part)
+                    robot.clear_dynamics()
+                    robot.set_qpos(qpos=initial_qpos, name=control_part, target=False)
+                    sim.enable_gizmo(uid=robot.uid, control_part=control_part)
+                    log_info("Robot reset to initial pose")
+
+                # Print info
+                elif key in ["i", "I"]:
+                    log_info("\n=== Gizmo Robot Control ===")
+                    log_info("Gizmo Controls:")
+                    log_info("  Use the 3D gizmo to drag and manipulate the robot")
+                    log_info("\nKeyboard Controls:")
+                    log_info("  Q/ESC: Exit control loop")
+                    log_info("  P: Print current robot state")
+                    log_info("  G: Toggle gizmo visibility")
+                    log_info("  R: Reset robot to initial pose")
+                    log_info("  I: Print this information again")
+
+    except KeyboardInterrupt:
+        sim.disable_gizmo(uid=robot.uid, control_part=control_part)
+        if robot_solver:
+            robot.init_solver(
+                cfg={control_part: robot_solver.cfg}
+            )  # Restore original solver
+        log_info("\nControl loop interrupted by user (Ctrl+C)")
+
+    finally:
+        try:
+            # Restore terminal settings
+            termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_settings)
+        except:
+            pass
+        log_info("Gizmo control loop terminated")