Fix H2 Head Motor Control in DDS Example

example/h2/high_level/h2_arm_sdk_dds_example.py

@@ -6,12 +6,17 @@
 from unitree_sdk2py.core.channel import ChannelFactoryInitialize, ChannelPublisher, ChannelSubscriber
 from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_
 from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowCmd_, LowState_
+from unitree_sdk2py.h2.loco.h2_loco_client import LocoClient
 from unitree_sdk2py.utils.crc import CRC
 from unitree_sdk2py.utils.thread import RecurrentThread
 
 
 SMALL_SHOULDER_ROLL = np.deg2rad(20.0)
 SMALL_ELBOW = np.deg2rad(30.0)
+SMALL_HEAD_YAW = np.deg2rad(20.0)
+SMALL_HEAD_PITCH = np.deg2rad(15.0)
+
+ARM_SDK_SUPPORTED_FSM_IDS = {4, 703}
 
 
 class H2JointIndex:
@@ -59,8 +64,8 @@ class H2JointIndex:
     RightWristYaw = 28
 
     # Head
-    HeadYaw = 29
-    HeadPitch = 30
+    HeadPitch = 29
+    HeadYaw = 30
 
     kNotUsedJoint = 31  # NOTE: Weight
 
@@ -72,6 +77,8 @@ def __init__(self):
         self.duration_ = 3.0
         self.kp = 80.0
         self.kd = 1.5
+        self.head_kp = 30.0
+        self.head_kd = 1.0
         self.low_cmd = unitree_hg_msg_dds__LowCmd_()
         self.low_state = None
         self.first_update_low_state = False
@@ -81,9 +88,10 @@ def __init__(self):
         self.target_pos = [
             0.0, SMALL_SHOULDER_ROLL, 0.0, SMALL_ELBOW, 0.0, 0.0, 0.0,
             0.0, -SMALL_SHOULDER_ROLL, 0.0, SMALL_ELBOW, 0.0, 0.0, 0.0,
+            SMALL_HEAD_PITCH, SMALL_HEAD_YAW,
         ]
 
-        self.arm_joints = [
+        self.upper_body_joints = [
             H2JointIndex.LeftShoulderPitch, H2JointIndex.LeftShoulderRoll,
             H2JointIndex.LeftShoulderYaw, H2JointIndex.LeftElbow,
             H2JointIndex.LeftWristRoll, H2JointIndex.LeftWristPitch,
@@ -92,8 +100,11 @@ def __init__(self):
             H2JointIndex.RightShoulderYaw, H2JointIndex.RightElbow,
             H2JointIndex.RightWristRoll, H2JointIndex.RightWristPitch,
             H2JointIndex.RightWristYaw,
+            H2JointIndex.HeadPitch, H2JointIndex.HeadYaw,
         ]
 
+        self.head_joints = {H2JointIndex.HeadYaw, H2JointIndex.HeadPitch}
+
     def Init(self):
         # create publisher
         self.arm_sdk_publisher = ChannelPublisher("rt/arm_sdk", LowCmd_)
@@ -118,52 +129,48 @@ def LowStateHandler(self, msg: LowState_):
         if not self.first_update_low_state:
             self.first_update_low_state = True
 
+    def SetJointCommand(self, joint, q):
+        command = self.low_cmd.motor_cmd[joint]
+        command.tau = 0.0
+        command.q = q
+        command.dq = 0.0
+        if joint in self.head_joints:
+            command.kp = self.head_kp
+            command.kd = self.head_kd
+        else:
+            command.kp = self.kp
+            command.kd = self.kd
+
     def LowCmdWrite(self):
         if self.low_state is None:
             return
 
         self.time_ += self.control_dt_
 
         if self.time_ < self.duration_:
-            print("[Stage 1]: set arms to zero posture.")
-            # [Stage 1]: set arms to zero posture
+            print("[Stage 1]: set arms and head to zero posture.")
             self.low_cmd.motor_cmd[H2JointIndex.kNotUsedJoint].q = 1.0
 
-            for i, joint in enumerate(self.arm_joints):
+            for joint in self.upper_body_joints:
                 ratio = np.clip(self.time_ / self.duration_, 0.0, 1.0)
-                self.low_cmd.motor_cmd[joint].tau = 0.0
-                self.low_cmd.motor_cmd[joint].q = (
-                    (1.0 - ratio) * self.low_state.motor_state[joint].q
-                )
-                self.low_cmd.motor_cmd[joint].dq = 0.0
-                self.low_cmd.motor_cmd[joint].kp = self.kp
-                self.low_cmd.motor_cmd[joint].kd = self.kd
+                q = (1.0 - ratio) * self.low_state.motor_state[joint].q
+                self.SetJointCommand(joint, q)
 
         elif self.time_ < self.duration_ * 3:
-            print("[Stage 2]: lift arms up.")
-            # [Stage 2]: lift arms up
-            for i, joint in enumerate(self.arm_joints):
+            print("[Stage 2]: lift arms and move head.")
+            for i, joint in enumerate(self.upper_body_joints):
                 ratio = np.clip((self.time_ - self.duration_) / (self.duration_ * 2), 0.0, 1.0)
-                self.low_cmd.motor_cmd[joint].tau = 0.0
-                self.low_cmd.motor_cmd[joint].q = (
+                q = (
                     ratio * self.target_pos[i] + (1.0 - ratio) * self.low_state.motor_state[joint].q
                 )
-                self.low_cmd.motor_cmd[joint].dq = 0.0
-                self.low_cmd.motor_cmd[joint].kp = self.kp
-                self.low_cmd.motor_cmd[joint].kd = self.kd
+                self.SetJointCommand(joint, q)
 
         elif self.time_ < self.duration_ * 6:
-            print("[Stage 3]: set arms back to zero posture.")
-            # [Stage 3]: set arms back to zero posture
-            for i, joint in enumerate(self.arm_joints):
+            print("[Stage 3]: set arms and head back to zero posture.")
+            for joint in self.upper_body_joints:
                 ratio = np.clip((self.time_ - self.duration_ * 3) / (self.duration_ * 3), 0.0, 1.0)
-                self.low_cmd.motor_cmd[joint].tau = 0.0
-                self.low_cmd.motor_cmd[joint].q = (
-                    (1.0 - ratio) * self.low_state.motor_state[joint].q
-                )
-                self.low_cmd.motor_cmd[joint].dq = 0.0
-                self.low_cmd.motor_cmd[joint].kp = self.kp
-                self.low_cmd.motor_cmd[joint].kd = self.kd
+                q = (1.0 - ratio) * self.low_state.motor_state[joint].q
+                self.SetJointCommand(joint, q)
 
         elif self.time_ < self.duration_ * 7:
             print("[Stage 4]: release arm_sdk.")
@@ -187,12 +194,43 @@ def LowCmdWrite(self):
     else:
         ChannelFactoryInitialize(0)
 
-    custom = Custom()
-    custom.Init()
-    custom.Start()
+    loco_client = LocoClient()
+    loco_client.SetTimeout(5.0)
+    loco_client.Init()
+
+    ret, fsm_id = loco_client.GetFsmId()
+    if ret != 0:
+        print(f"Failed to get H2 FSM ID, error code: {ret}")
+        sys.exit(1)
+    if fsm_id not in ARM_SDK_SUPPORTED_FSM_IDS:
+        print(
+            f"Current FSM {fsm_id} does not accept rt/arm_sdk commands. "
+            f"Switch to one of {sorted(ARM_SDK_SUPPORTED_FSM_IDS)} first."
+        )
+        sys.exit(1)
+
+    ret = loco_client.EnableArmSDK()
+    if ret != 0:
+        print(f"Failed to enable the external Arm SDK, error code: {ret}")
+        sys.exit(1)
 
-    while True:
-        time.sleep(1)
-        if custom.done:
-            print("Done!")
-            sys.exit(-1)
+    exit_code = 0
+    try:
+        custom = Custom()
+        custom.Init()
+        custom.Start()
+
+        while True:
+            time.sleep(1)
+            if custom.done:
+                print("Done!")
+                break
+    except KeyboardInterrupt:
+        print("Interrupted; releasing Arm SDK control.")
+    finally:
+        ret = loco_client.DisableArmSDK()
+        if ret != 0:
+            print(f"Failed to disable the external Arm SDK, error code: {ret}")
+            exit_code = 1
+
+    sys.exit(exit_code)

unitree_sdk2py/h2/loco/h2_loco_client.py

@@ -170,10 +170,10 @@ def ShakeHand(self, stage: int = -1):
             return self.SetTaskId(3 if self.first_shake_hand_stage_ else 2)
 
     def EnableArmSDK(self):
-        self.SetArmSdkStatus(True)
+        return self.SetArmSdkStatus(True)
 
     def DisableArmSDK(self):
-        self.SetArmSdkStatus(False)
+        return self.SetArmSdkStatus(False)
 
     def GetFsmId(self):
         p = {}
@@ -246,4 +246,3 @@ def GetAvailableFsmIds(self):
             return code, None, None
         js = json.loads(data)
         return code, js.get("ids", []), js.get("names", [])
-