g1-teleop/server/joint_test.py

#!/usr/bin/env python3
"""
Systematic joint-by-joint testing for sign calibration.

Animates ONE joint at a time with a slow sine wave so the user can see
exactly what each joint does, then match the motion with their Quest 3.

Usage:
    python joint_test.py <joint_name> [--amplitude 0.5] [--period 4.0]

Joint names:
    waist_yaw       (joint 12) - Z axis
    waist_roll      (joint 13) - X axis
    waist_pitch     (joint 14) - Y axis
    l_shoulder_pitch (joint 15) - Y axis
    l_shoulder_roll  (joint 16) - X axis
    l_shoulder_yaw   (joint 17) - Z axis
    l_elbow         (joint 18) - Y axis
    l_wrist_roll    (joint 19) - X axis
    l_wrist_pitch   (joint 20) - Y axis
    l_wrist_yaw     (joint 21) - Z axis
    r_shoulder_pitch (joint 22) - Y axis
    r_shoulder_roll  (joint 23) - X axis
    r_shoulder_yaw   (joint 24) - Z axis
    r_elbow         (joint 25) - Y axis
    r_wrist_roll    (joint 26) - X axis
    r_wrist_pitch   (joint 27) - Y axis
    r_wrist_yaw     (joint 28) - Z axis

    Also: "list" to show all joints, "all_shoulders" to cycle through shoulder joints
"""

import sys
import os
import time
import argparse
import signal
import numpy as np

NUM_JOINTS = 29

JOINTS = {
    'waist_yaw':        12,
    'waist_roll':       13,
    'waist_pitch':      14,
    'l_shoulder_pitch': 15,
    'l_shoulder_roll':  16,
    'l_shoulder_yaw':   17,
    'l_elbow':          18,
    'l_wrist_roll':     19,
    'l_wrist_pitch':    20,
    'l_wrist_yaw':      21,
    'r_shoulder_pitch': 22,
    'r_shoulder_roll':  23,
    'r_shoulder_yaw':   24,
    'r_elbow':          25,
    'r_wrist_roll':     26,
    'r_wrist_pitch':    27,
    'r_wrist_yaw':      28,
}

# Joint descriptions - what POSITIVE values should do
DESCRIPTIONS = {
    'waist_yaw':        'Torso rotates LEFT (counterclockwise from above)',
    'waist_roll':       'Torso leans RIGHT',
    'waist_pitch':      'Torso leans FORWARD',
    'l_shoulder_pitch': 'Left arm swings FORWARD',
    'l_shoulder_roll':  'Left arm raises OUT to the LEFT side',
    'l_shoulder_yaw':   'Left upper arm rotates INWARD',
    'l_elbow':          'Left elbow BENDS (forearm toward shoulder)',
    'l_wrist_roll':     'Left wrist rotates (palm faces down→up)',
    'l_wrist_pitch':    'Left wrist bends DOWN (flexion)',
    'l_wrist_yaw':      'Left wrist bends toward THUMB side',
    'r_shoulder_pitch': 'Right arm swings FORWARD',
    'r_shoulder_roll':  'Right arm raises OUT to the RIGHT side (NEGATIVE values)',
    'r_shoulder_yaw':   'Right upper arm rotates INWARD',
    'r_elbow':          'Right elbow BENDS (forearm toward shoulder)',
    'r_wrist_roll':     'Right wrist rotates',
    'r_wrist_pitch':    'Right wrist bends DOWN (flexion)',
    'r_wrist_yaw':      'Right wrist bends toward THUMB side',
}

# Default amplitudes (stay within safe range)
DEFAULT_AMPLITUDES = {
    'waist_yaw':        0.4,
    'waist_roll':       0.3,
    'waist_pitch':      0.3,
    'l_shoulder_pitch': 0.8,
    'l_shoulder_roll':  0.8,
    'l_shoulder_yaw':   0.5,
    'l_elbow':          0.8,
    'l_wrist_roll':     0.5,
    'l_wrist_pitch':    0.5,
    'l_wrist_yaw':      0.5,
    'r_shoulder_pitch': 0.8,
    'r_shoulder_roll':  0.8,
    'r_shoulder_yaw':   0.5,
    'r_elbow':          0.8,
    'r_wrist_roll':     0.5,
    'r_wrist_pitch':    0.5,
    'r_wrist_yaw':      0.5,
}

# Default Kp/Kd gains
KP = [
    60, 60, 60, 100, 40, 40,
    60, 60, 60, 100, 40, 40,
    60, 40, 40,
    40, 40, 40, 40, 40, 40, 40,
    40, 40, 40, 40, 40, 40, 40,
]
KD = [
    1, 1, 1, 2, 1, 1,
    1, 1, 1, 2, 1, 1,
    1, 1, 1,
    1, 1, 1, 1, 1, 1, 1,
    1, 1, 1, 1, 1, 1, 1,
]


def init_dds(domain_id=1):
    from unitree_sdk2py.core.channel import ChannelPublisher, ChannelFactoryInitialize
    from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowCmd_
    from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_
    from unitree_sdk2py.utils.crc import CRC

    ChannelFactoryInitialize(domain_id)
    pub = ChannelPublisher("rt/lowcmd", LowCmd_)
    pub.Init()
    cmd = unitree_hg_msg_dds__LowCmd_()
    crc = CRC()
    return pub, cmd, crc


def publish(pub, cmd, crc, angles):
    cmd.mode_pr = 0
    cmd.mode_machine = 0
    for i in range(NUM_JOINTS):
        cmd.motor_cmd[i].mode = 1
        cmd.motor_cmd[i].q = float(angles[i])
        cmd.motor_cmd[i].dq = 0.0
        cmd.motor_cmd[i].tau = 0.0
        cmd.motor_cmd[i].kp = float(KP[i])
        cmd.motor_cmd[i].kd = float(KD[i])
    # Zero legs
    for i in range(12):
        cmd.motor_cmd[i].mode = 0
        cmd.motor_cmd[i].kp = 0.0
        cmd.motor_cmd[i].kd = 0.0
    cmd.crc = crc.Crc(cmd)
    pub.Write(cmd)


def main():
    parser = argparse.ArgumentParser(description="Joint-by-joint animation tester")
    parser.add_argument("joint", help="Joint name (or 'list' to show all)")
    parser.add_argument("--amplitude", "-a", type=float, default=None,
                        help="Sine wave amplitude in radians (default: joint-specific)")
    parser.add_argument("--period", "-p", type=float, default=4.0,
                        help="Sine wave period in seconds (default: 4.0)")
    parser.add_argument("--hz", type=float, default=50.0)
    parser.add_argument("--offset", "-o", type=float, default=0.0,
                        help="Static offset added to sine wave")
    parser.add_argument("--dds-domain", type=int, default=1)
    parser.add_argument("--pause-at-peak", choices=["pos", "neg"], default=None,
                        help="Pause at peak each cycle: 'pos' or 'neg'")
    parser.add_argument("--pause-duration", type=float, default=3.0,
                        help="How long to pause in seconds (default: 3.0)")
    args = parser.parse_args()

    if args.joint == 'list':
        print("\nAvailable joints:\n")
        for name, idx in sorted(JOINTS.items(), key=lambda x: x[1]):
            desc = DESCRIPTIONS.get(name, '')
            amp = DEFAULT_AMPLITUDES.get(name, 0.5)
            print(f"  {name:20s} (joint {idx:2d})  +val: {desc}")
        print()
        return

    if args.joint not in JOINTS:
        print(f"Unknown joint: {args.joint}")
        print(f"Available: {', '.join(sorted(JOINTS.keys()))}")
        return

    joint_idx = JOINTS[args.joint]
    amplitude = args.amplitude if args.amplitude is not None else DEFAULT_AMPLITUDES.get(args.joint, 0.5)
    desc = DESCRIPTIONS.get(args.joint, '')

    print(f"\n{'='*60}")
    print(f"Testing: {args.joint} (joint index {joint_idx})")
    print(f"Positive direction: {desc}")
    print(f"Animation: sine wave, amplitude={amplitude:.2f} rad, period={args.period:.1f}s")
    print(f"Offset: {args.offset:.2f} rad")
    print(f"{'='*60}\n")

    pub, cmd, crc = init_dds(args.dds_domain)

    running = True
    def handler(sig, frame):
        nonlocal running
        running = False
    signal.signal(signal.SIGINT, handler)
    signal.signal(signal.SIGTERM, handler)

    interval = 1.0 / args.hz
    t0 = time.time()
    step = 0

    print("Animating... (Ctrl+C to stop)\n")

    last_pause_time = 0
    while running:
        t_start = time.perf_counter()
        elapsed = time.time() - t0
        angles = np.zeros(NUM_JOINTS)

        value = args.offset + amplitude * np.sin(2 * np.pi * elapsed / args.period)
        angles[joint_idx] = value

        publish(pub, cmd, crc, angles)

        step += 1
        if step % int(args.hz * 1) == 0:
            print(f"  {args.joint} = {value:+.3f} rad ({np.degrees(value):+.1f} deg)")

        # Pause at peak (with cooldown to avoid re-triggering)
        if args.pause_at_peak and (time.time() - last_pause_time) > args.period * 0.5:
            if args.pause_at_peak == "pos":
                peak_val = args.offset + amplitude
            else:
                peak_val = args.offset - amplitude
            if abs(value - peak_val) < 0.05:
                print(f"\n  >>> PAUSED at {args.pause_at_peak} peak: {value:+.3f} rad ({np.degrees(value):+.1f} deg) <<<")
                pause_end = time.time() + args.pause_duration
                while time.time() < pause_end and running:
                    publish(pub, cmd, crc, angles)
                    time.sleep(interval)
                last_pause_time = time.time()
                print(f"  >>> Resuming...\n")

        sleep_time = interval - (time.perf_counter() - t_start)
        if sleep_time > 0:
            time.sleep(sleep_time)

    # Return to zero
    print("\nReturning to zero...")
    angles = np.zeros(NUM_JOINTS)
    for _ in range(25):
        publish(pub, cmd, crc, angles)
        time.sleep(0.02)
    print("Done.")


if __name__ == "__main__":
    main()