diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 1c97789a..35ae05b5 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/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. diff --git a/embodichain/lab/sim/utility/gizmo_utils.py b/embodichain/lab/sim/utility/gizmo_utils.py index 676a93fa..b2a2fafc 100644 --- a/embodichain/lab/sim/utility/gizmo_utils.py +++ b/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")