code cleanup; runnning ik and control on separate threads

Author: zhigenzhao

scripts/hardware/teleop_a1x_hardware.py

@@ -0,0 +1,40 @@
+import threading
+
+import rospy
+
+from xrobotoolkit_teleop.hardware.galaxea_teleop_controller import GalaxeaA1XTeleopController
+from xrobotoolkit_teleop.utils.xr_client import XrClient
+
+
+def main():
+    xr_client = XrClient()
+    controller = GalaxeaA1XTeleopController(xr_client=xr_client)
+
+    stop_event = threading.Event()
+    rospy.on_shutdown(lambda: stop_event.set())
+    left_arm_thread = threading.Thread(
+        target=controller.run_control_thread,
+        args=(stop_event,),
+    )
+
+    ik_thread = threading.Thread(
+        target=controller.run_ik_thread,
+        args=(stop_event,),
+    )
+
+    left_arm_thread.start()
+    ik_thread.start()
+
+    try:
+        rospy.spin()
+    except KeyboardInterrupt:
+        print("Keyboard interrupt received, shutting down...")
+    finally:
+        stop_event.set()
+
+    left_arm_thread.join()
+    ik_thread.join()
+
+
+if __name__ == "__main__":
+    main()

xrobotoolkit_teleop/hardware/galaxea.py

@@ -19,23 +19,14 @@ def __init__(self, rate_hz=100):
         self.timestamp = 0.0
 
         # motor control parameters
-        self.q_des = None  # Desired joint positions
-        self.v_des = [0.0] * 6  # Desired joint velocities
-        self.kp = [2000, 2000, 1000, 100, 100, 100]
-        self.kd = [500.0, 500.0, 500, 100, 100, 100]
+        self.q_des = None
+        self.v_des = [0.0] * 6
+        self.kp = [2000, 2000, 1000, 200, 200, 200]
+        self.kd = [200.0, 200.0, 200, 100, 100, 100]
         self.t_ff = [0.0] * 6
-        # self.kd = [5000.0, 5000.0, 5000.0, 1000.0, 1000.0, 1000.0]
         self.arm_ctrl_msg = motor_control()
-        # self.joint_names = [
-        #     "arm_joint_1",
-        #     "arm_joint_2",
-        #     "arm_joint_3",
-        #     "arm_joint_4",
-        #     "arm_joint_5",
-        #     "arm_joint_6",
-        # ]
 
-        self.q_des_gripper = [-2]
+        self.q_des_gripper = [-2.1]
         self.v_des_gripper = [0.0]
         self.kp_gripper = [1]
         self.kd_gripper = [0.05]
@@ -46,15 +37,15 @@ def arm_state_callback(self, msg: JointState):
         """
         Callback function to handle joint state updates.
         """
-        self.qpos = msg.position[:6]  # Assuming first 6 are arm joints
-        self.qvel = msg.velocity[:6]  # Assuming first 6 are arm joints
+        self.qpos = msg.position[:6]
+        self.qvel = msg.velocity[:6]
         self.qpos_gripper = [msg.position[6]]
         self.qvel_gripper = [msg.velocity[6]]
         if self.q_des is None:
             self.q_des = self.qpos
         self.timestamp = msg.header.stamp.to_sec()
-        print(f"Received joint state: {self.qpos}, {self.qvel} at time {self.timestamp}")
-        print(f"Gripper state: {self.qpos_gripper}, {self.qvel_gripper}")
+        # print(f"Received joint state: {self.qpos}, {self.qvel} at time {self.timestamp}")
+        # print(f"Gripper state: {self.qpos_gripper}, {self.qvel_gripper}")
 
     def publish_arm_control(self):
         """
@@ -63,9 +54,9 @@ def publish_arm_control(self):
         if self.q_des is None:
             return
 
-        self.arm_ctrl_msg.header = Header()  # Create a new Header instance
+        self.arm_ctrl_msg.header = Header()
         self.arm_ctrl_msg.header.stamp = rospy.Time.now()
-        self.arm_ctrl_msg.header.frame_id = "base_link"  # Or your relevant frame
+        self.arm_ctrl_msg.header.frame_id = "base_link"
         self.arm_ctrl_msg.kp = self.kp
         self.arm_ctrl_msg.kd = self.kd
         self.arm_ctrl_msg.t_ff = self.t_ff
@@ -91,19 +82,5 @@ def run(self):
         Main loop to run the controller.
         """
         while not rospy.is_shutdown():
-            # Example: Publish current joint positions and velocities
             self.publish_motor_control()
-            # Sleep to maintain the loop rate
             self.rate.sleep()
-
-
-def main():
-    controller = A1XController()
-    try:
-        controller.run()
-    except rospy.ROSInterruptException:
-        pass
-
-
-if __name__ == "__main__":
-    main()

xrobotoolkit_teleop/hardware/galaxea_teleop_controller.py

@@ -1,4 +1,5 @@
 import os
+import threading
 import time
 import webbrowser
 
@@ -18,28 +19,28 @@
 from xrobotoolkit_teleop.utils.path_utils import ASSET_PATH
 from xrobotoolkit_teleop.utils.xr_client import XrClient
 
-# Default paths and configurations for a single left arm
+# Default paths and configurations for a single right arm
 DEFAULT_SINGLE_A1X_URDF_PATH = os.path.join(ASSET_PATH, "galaxea/A1X/a1x.urdf")
 DEFAULT_SCALE_FACTOR = 1.0
 CONTROLLER_DEADZONE = 0.1
 
-# Default end-effector configuration for a single left arm without a gripper
+# Default end-effector configuration for a single right arm without a gripper
 DEFAULT_END_EFFECTOR_CONFIG = {
-    "left_arm": {
+    "right_arm": {
         "link_name": "gripper_link",  # URDF link name for the single arm
-        "pose_source": "left_controller",
-        "control_trigger": "left_grip",
+        "pose_source": "right_controller",
+        "control_trigger": "right_grip",
         "gripper_config": {
-            "joint_name": "left_gripper_finger_joint1",
-            "gripper_trigger": "left_trigger",
+            "joint_name": "right_gripper_finger_joint1",
+            "gripper_trigger": "right_trigger",
             "open_pos": -2.0,
             "close_pos": 0.0,
         },
     },
 }
 
 
-class GalaxeaTeleopControllerSingleArm:
+class GalaxeaA1XTeleopController:
     def __init__(
         self,
         xr_client: XrClient,
@@ -50,31 +51,22 @@ def __init__(
         visualize_placo: bool = True,
         ros_rate_hz: int = 100,
     ):
-        # rospy.init_node("galaxea_teleop_controller_single", anonymous=True)
 
         self.xr_client = xr_client
         self.robot_urdf_path = robot_urdf_path
         self.R_headset_world = R_headset_world
         self.scale_factor = scale_factor
         self.visualize_placo = visualize_placo
         self.end_effector_config = end_effector_config
-        # self.rate = rospy.Rate(ros_rate_hz)
 
-        # Define joint names for the single left arm (no gripper)
-        # left_joint_names = [f"arm_joint_{i}" for i in range(1, 7)]
+        self.right_controller = A1XController()
 
-        # Instantiate hardware controller for the left arm
-        # Assuming the single arm uses the "left" prefix for its topics
-        self.left_controller = A1XController()
-
-        # Placo Setup
         self.placo_robot = placo.RobotWrapper(self.robot_urdf_path)
         self.solver = placo.KinematicsSolver(self.placo_robot)
         self.solver.dt = 1.0 / ros_rate_hz
         self.solver.mask_fbase(True)
         self.solver.add_kinetic_energy_regularization_task(1e-6)
 
-        # Setup Placo tasks for the single arm
         self.effector_task = {}
         self.init_ee_xyz = {}
         self.init_ee_quat = {}
@@ -97,14 +89,12 @@ def __init__(
                 gripper_config = config["gripper_config"]
                 self.gripper_pos_target[gripper_config["joint_name"]] = gripper_config["open_pos"]
 
-        # Wait for the first joint state message to initialize Placo state
         print("Waiting for initial joint state from the Galaxea arm...")
-        while not rospy.is_shutdown() and self.left_controller.timestamp == 0:
+        while not rospy.is_shutdown() and self.right_controller.timestamp == 0:
             rospy.sleep(0.1)
         print("Initial joint state received.")
 
-        # Set initial Placo state from hardware. Assumes single arm URDF has 7 DoF base + 6 DoF arm
-        self.placo_robot.state.q[7:13] = self.left_controller.qpos
+        self.placo_robot.state.q[7:13] = self.right_controller.qpos
 
         if self.visualize_placo:
             self.placo_robot.update_kinematics()
@@ -118,7 +108,7 @@ def __init__(
 
     def _process_xr_pose(self, xr_pose, arm_name: str):
         controller_xyz = np.array([xr_pose[0], xr_pose[1], xr_pose[2]])
-        controller_quat = np.array([xr_pose[6], xr_pose[3], xr_pose[4], xr_pose[5]])  # w, x, y, z
+        controller_quat = np.array([xr_pose[6], xr_pose[3], xr_pose[4], xr_pose[5]])
 
         controller_xyz = self.R_headset_world @ controller_xyz
         R_transform = np.eye(4)
@@ -138,12 +128,10 @@ def _process_xr_pose(self, xr_pose, arm_name: str):
             delta_rot = quat_diff_as_angle_axis(self.init_controller_quat[arm_name], controller_quat)
         return delta_xyz, delta_rot
 
-    def update_ik_and_publish(self):
-        # Update placo state with the latest from hardware
-        self.placo_robot.state.q[7:13] = self.left_controller.qpos
+    def update_ik(self):
+        self.placo_robot.state.q[7:13] = self.right_controller.qpos
         self.placo_robot.update_kinematics()
 
-        # Process XR inputs and update IK task (only one loop for the single arm)
         for arm_name, config in self.end_effector_config.items():
             xr_grip_val = self.xr_client.get_key_value_by_name(config["control_trigger"])
             active = xr_grip_val > (1.0 - CONTROLLER_DEADZONE)
@@ -165,11 +153,9 @@ def update_ik_and_publish(self):
                 target_pose[:3, 3] = target_xyz
                 self.effector_task[arm_name].T_world_frame = target_pose
 
-                # Solve IK
                 try:
                     self.solver.solve(True)
-                    # Update target joint positions from IK solution
-                    self.left_controller.q_des = self.placo_robot.state.q[7:13].copy()
+                    self.right_controller.q_des = self.placo_robot.state.q[7:13].copy()
 
                     if self.visualize_placo:
                         self.placo_vis.display(self.placo_robot.state.q)
@@ -185,8 +171,8 @@ def update_ik_and_publish(self):
                     self.init_ee_xyz[arm_name] = None
                     self.init_controller_xyz[arm_name] = None
                     print(f"{arm_name} deactivated.")
-                # Keep task target at current robot pose when not active
                 self.effector_task[arm_name].T_world_frame = self.placo_robot.get_T_world_frame(config["link_name"])
+
             # Update gripper position based on XR input
             if "gripper_config" in config:
                 gripper_config = config["gripper_config"]
@@ -199,28 +185,23 @@ def update_ik_and_publish(self):
                 joint_name = gripper_config["joint_name"]
                 self.gripper_pos_target[joint_name] = [gripper_pos]
 
-        # if abs(self.gripper_pos_target["left_gripper_finger_joint1"][0] - self.left_controller.qpos_gripper[0]) > 0.01:
-        self.left_controller.q_des_gripper = self.gripper_pos_target["left_gripper_finger_joint1"]
-        self.left_controller.publish_gripper_control()
-        # Publish commands to hardware
-        self.left_controller.publish_arm_control()
+        self.right_controller.q_des_gripper = self.gripper_pos_target[
+            self.end_effector_config["right_arm"]["gripper_config"]["joint_name"]
+        ]
+
+    def run_ik_thread(self, stop_event: threading.Event):
+        print("Starting Galaxea A1X single arm teleop controller IK loop...")
+        while not (rospy.is_shutdown() or stop_event.is_set()):
+            start = time.time()
+            self.update_ik()
+            elapsed = time.time() - start
+            time.sleep(max(0, (self.solver.dt - elapsed)))
+        print("Galaxea teleop controller IK loop shutting down.")
 
-    def run(self):
+    def run_control_thread(self, stop_event: threading.Event):
         print("Starting Galaxea A1X single arm teleop controller loop...")
-        while not rospy.is_shutdown():
-            self.update_ik_and_publish()
-            self.left_controller.rate.sleep()
+        while not (rospy.is_shutdown() or stop_event.is_set()):
+            self.right_controller.publish_arm_control()
+            self.right_controller.publish_gripper_control()
+            self.right_controller.rate.sleep()
         print("Galaxea teleop controller shutting down.")
-
-
-if __name__ == "__main__":
-    # This is an example of how to run the controller
-    # You would need to have your XR client and ROS environment set up
-    try:
-        xr_client = XrClient()
-        controller = GalaxeaTeleopControllerSingleArm(xr_client=xr_client)
-        controller.run()
-    except rospy.ROSInterruptException:
-        print("ROS node interrupted.")
-    except Exception as e:
-        print(f"An error occurred: {e}")