Add EE keyboard teleop/sim for AX arm with robust IK

- teleoperate_ee_keyboard.py: velocity IK (dual-quaternion resolved-rate
  ported from realtime_servoing, plus a position-Jacobian solver) integrated
  to position commands; local/global frame toggle where global uses the
  position Jacobian with damped least-squares; constant Cartesian-speed
  scaling with a joint-step cap; ramped homing to the reference pose.
- simulate_ee_keyboard.py: hardware-free 3D sim reusing the same IK.
- urdf_mapping.py: widen shoulder_pan (+/-90) and elbow_flex (0-180) limits.
- ax_arm.py: guard send_action against non-.pos / non-string keys.

lerobot_robot_ax_arm/examples/teleoperate_ee_keyboard.py

@@ -1,15 +1,25 @@
 #!/usr/bin/env python
 
-"""Keyboard end-effector teleoperation for the 4-DoF AX arm (position-only IK on the URDF).
+"""Keyboard end-effector teleoperation for the 4-DoF AX arm (velocity IK, position output).
 
-The arm has 3 revolute joints (base yaw + shoulder/elbow pitch) giving exactly 3 task-space DoF,
-all spent on reaching a 3D *position* (orientation is not controllable). Each frame we:
+Adapted from a resolved-rate (twist) servoing controller: instead of commanding joint
+velocities, the per-tick joint velocity is integrated into a joint *position* target and sent as
+``Goal_Position`` (the AX arm runs in position mode).
 
-  1. read the raw motor ticks and map them to URDF joint degrees (see ``urdf_mapping``),
-  2. forward-kinematics to the current end-effector pose,
-  3. offset the target position by the pressed keys,
-  4. solve position-only IK (``orientation_weight=0.0``) for new URDF joint degrees,
-  5. map back to ticks and command the servos.
+Each frame we read the raw motor ticks, map them to URDF joint angles, solve for a joint velocity
+that produces the requested Cartesian motion, scale it to a fixed end-effector Cartesian step
+``CART_STEP_M`` per tick (capped in joint space near singularities), and command ``Goal_Position``.
+
+Two IK solvers are selectable at runtime (``--solver`` / ``m`` key):
+  - "dq"  : dual-quaternion resolved-rate (matches the full pose velocity via ``scipy.least_squares``;
+            faithful to the source controller, but rotation/translation coupling on a 3-DoF arm
+            causes axis leakage),
+  - "pos" : position-only Jacobian ``dEE_pos/dq`` solved with least-squares (crisp axis-aligned
+            Cartesian motion).
+
+The motion frame is also selectable (``--frame`` / ``t`` key): "local" moves along the
+end-effector's own axes, "global" along the fixed world axes. Global always uses the position
+Jacobian (the dq solver's held-orientation constraint leaks axes on this underactuated arm).
 
 The tick<->URDF mapping is established once by ``lerobot-calibrate`` (reference pose + travel
 limits), so no separate alignment step is needed here.
@@ -19,6 +29,8 @@ Controls (letter keys; hold to keep moving via terminal key-repeat):
     - a / d : +Y / -Y   (left / right)
     - r / f : +Z / -Z   (up / down)
     - o / c : open / close gripper
+    - t     : toggle global <-> local frame
+    - m     : toggle dq <-> pos solver
     - ESC / q : stop
 
 Run:
@@ -29,30 +41,175 @@ import argparse
 import time
 from importlib.resources import files
 
+import casadi as cs
 import numpy as np
+import scipy as sp
+from urdf2casadi import urdfparser as u2c
+from urdf2casadi.geometry import dual_quaternion, quaternion
 
-from lerobot.model.kinematics import RobotKinematics
 from lerobot.utils.keyboard_input import create_key_listener
 from lerobot.utils.robot_utils import precise_sleep
 
 from lerobot_robot_ax_arm import AXArm, AXArmConfig
 from lerobot_robot_ax_arm.urdf_mapping import (
     ARM_JOINTS,
-    URDF_JOINT_NAMES,
+    REFERENCE_URDF_DEG,
     ticks_to_urdf_vector,
     urdf_vector_to_ticks,
 )
 
 FPS = 30
-LINEAR_STEP_M = 0.005  # EE position change per pressed frame
+HOME_TIME_S = 2.0  # duration of the ramped move to the reference pose at startup
+CART_STEP_M = 0.008  # default end-effector Cartesian motion per tick, meters (override with --speed)
+MAX_JOINT_STEP_RAD = 0.15  # safety cap on joint motion per tick (keeps motion bounded near singularities)
+DLS_LAMBDA = 0.02  # damping factor for the position IK, well-behaved near singularities
 GRIP_STEP_TICK = 15  # gripper ticks per press
+FIT_THRESHOLD = 0.1  # only fit dual-quaternion derivative components above this magnitude
+
+ROOT_LINK = "base_link"
+TIP_LINK = "gripper_link"
+CHAIN_LINKS = ("robot_link_1", "robot_link_2", "robot_link_3", "gripper_link")
+
+
+def _skew(x: np.ndarray) -> np.ndarray:
+    return np.array([[0, -x[2], x[1]], [x[2], 0, -x[0]], [-x[1], x[0], 0]])
+
+
+def _build_kinematics(urdf_path: str) -> dict:
+    """FK + Jacobians (dual-quaternion and position-only) and joint limits from the URDF."""
+    parser = u2c.URDFparser()
+    parser.from_file(urdf_path)
+    fk = parser.get_forward_kinematics(ROOT_LINK, TIP_LINK)
+    q_sym = fk["q"]
+    fk_dq = fk["dual_quaternion_fk"]
+    fk_T = fk["T_fk"]
+    return {
+        "fk_dq": fk_dq,
+        "fk_T": fk_T,
+        "dq_jac": cs.Function("dq_jac", [q_sym], [cs.jacobian(fk_dq(q_sym), q_sym)]),
+        "pos_jac": cs.Function("pos_jac", [q_sym], [cs.jacobian(fk_T(q_sym)[:3, 3], q_sym)]),
+        "lower": np.array(fk["lower"], dtype=float).flatten(),
+        "upper": np.array(fk["upper"], dtype=float).flatten(),
+    }
+
+
+def build_link_fks(urdf_path: str):
+    """Casadi T_fk functions base->each link in CHAIN_LINKS, for drawing the arm."""
+    fks = []
+    for tip in CHAIN_LINKS:
+        parser = u2c.URDFparser()
+        parser.from_file(urdf_path)
+        fk = parser.get_forward_kinematics(ROOT_LINK, tip)
+        fks.append((fk["T_fk"], fk["q"].shape[0]))
+    return fks
+
+
+def chain_points(link_fks, q_rad: np.ndarray) -> np.ndarray:
+    """3D positions of the base and each link origin along the kinematic chain."""
+    pts = [np.zeros(3)]
+    for T, n in link_fks:
+        pts.append(np.array(T(q_rad[:n]))[:3, 3])
+    return np.array(pts)
+
+
+def open_live_view(link_fks, reach: float = 0.42):
+    """Open a non-blocking 3D plot; returns an update(q_rad, title) callback (False once closed)."""
+    import os
+    import tempfile
+
+    os.environ.setdefault("MPLCONFIGDIR", tempfile.gettempdir())
+    import matplotlib.pyplot as plt
+
+    fig = plt.figure(figsize=(7, 6))
+    ax = fig.add_subplot(projection="3d")
+    (chain_line,) = ax.plot([], [], [], "-o", lw=3, color="tab:blue")
+    (ee_pt,) = ax.plot([], [], [], "o", ms=10, color="tab:red")
+    ax.set_xlim(-reach, reach); ax.set_ylim(-reach, reach); ax.set_zlim(-0.05, reach)
+    ax.set_xlabel("X"); ax.set_ylabel("Y"); ax.set_zlabel("Z")
+    plt.ion(); plt.show(block=False)
+
+    def update(q_rad: np.ndarray, title: str) -> bool:
+        if not plt.fignum_exists(fig.number):
+            return False
+        pts = chain_points(link_fks, q_rad)
+        chain_line.set_data(pts[:, 0], pts[:, 1]); chain_line.set_3d_properties(pts[:, 2])
+        ee_pt.set_data(pts[-1:, 0], pts[-1:, 1]); ee_pt.set_3d_properties(pts[-1:, 2])
+        ax.set_title(title)
+        fig.canvas.draw_idle(); fig.canvas.flush_events()
+        return True
+
+    return update
+
+
+def _world_twist_to_dq_dot(pose: np.ndarray, twist_world: np.ndarray) -> np.ndarray:
+    """World-frame twist [v, w] -> dual-quaternion derivative x_dot at the current pose."""
+    T = dual_quaternion.to_numpy_transformation_matrix(pose)
+    adj = np.zeros((6, 6))
+    adj[:3, :3] = T[:3, :3]
+    adj[3:, 3:] = T[:3, :3]
+    adj[:3, 3:] = _skew(T[:3, -1]) @ T[:3, :3]
+    twist = adj @ twist_world
+
+    v = np.append(twist[:3], 0)
+    w = np.append(twist[3:], 0)
+    primal = pose[:4]
+    dual = pose[4:]
+    primal_conj = np.append(-primal[:3], primal[3])
+    p = 2 * quaternion.numpy_product(dual, primal_conj)
+
+    xi = np.zeros(8)
+    xi[:4] = np.append(0, twist[3:])
+    xi[4:] = v + (quaternion.numpy_product(p, w) - quaternion.numpy_product(w, p)) / 2
+    return 0.5 * dual_quaternion.numpy_product(xi, pose)
+
+
+def _fitness(q_dot, jacobian, x_dot):
+    index = np.where(np.abs(x_dot) > FIT_THRESHOLD)
+    delta = (np.dot(jacobian, q_dot) - x_dot)[index]
+    return np.dot(delta.T, delta)
+
+
+def _fitness_jacobian(q_dot, jacobian, x_dot):
+    return 2 * np.dot(jacobian.T, np.dot(jacobian, q_dot) - x_dot)
+
+
+def _joint_velocity(kin: dict, q_rad: np.ndarray, cmd: np.ndarray, frame: str, solver: str) -> np.ndarray:
+    """Joint velocity producing the requested unit Cartesian motion ``cmd`` (x, y, z)."""
+    rot = np.array(kin["fk_T"](q_rad))[:3, :3]
+    # World-frame ("global") translation on this 3-DoF (no-wrist) arm is only reliable via the position
+    # Jacobian: the dq resolved-rate tries to hold EE orientation fixed, which an underactuated arm
+    # cannot do, leaking the motion onto other axes as the arm reorients. So global always uses pos.
+    if solver == "pos" or frame == "global":
+        v_world = cmd.astype(float) if frame == "global" else rot @ cmd
+        J = np.array(kin["pos_jac"](q_rad))
+        # Damped least squares: bounded, well-conditioned joint velocity even near singularities.
+        return J.T @ np.linalg.solve(J @ J.T + DLS_LAMBDA**2 * np.eye(3), v_world)
+
+    # Dual-quaternion resolved-rate (local frame only): move along the end-effector's own axes.
+    lin = cmd.astype(float)
+    pose = np.array(kin["fk_dq"](q_rad)).flatten()
+    x_dot = _world_twist_to_dq_dot(pose, np.concatenate([lin, np.zeros(3)]))
+    jacobian = np.array(kin["dq_jac"](q_rad))
+    sol = sp.optimize.least_squares(
+        _fitness, np.zeros(len(ARM_JOINTS)), args=(jacobian, x_dot), xtol=1e-4, jac=_fitness_jacobian,
+    )
+    return sol.x if sol.success else np.zeros(len(ARM_JOINTS))
 
 
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--port", required=True, help="Serial port of the AX arm")
     parser.add_argument("--id", default="my_ax_arm", help="Robot id used for calibration files")
+    parser.add_argument("--speed", type=float, default=CART_STEP_M,
+                        help="End-effector Cartesian motion per tick in meters (higher = faster)")
+    parser.add_argument("--frame", choices=("local", "global"), default="local",
+                        help="Motion frame: 'local' = end-effector axes, 'global' = world axes")
+    parser.add_argument("--solver", choices=("dq", "pos"), default="dq",
+                        help="IK solver: 'dq' = dual-quaternion resolved-rate, 'pos' = position Jacobian")
+    parser.add_argument("--view", action="store_true",
+                        help="Show a live 3D plot of the real arm (digital twin) alongside teleop")
     args = parser.parse_args()
+    cart_step = args.speed
 
     robot = AXArm(AXArmConfig(port=args.port, id=args.id, use_degrees=True))
     robot.connect(calibrate=False)
@@ -60,13 +217,29 @@ def main():
         raise RuntimeError(f"No calibration found for id '{args.id}'. Run lerobot-calibrate first.")
 
     urdf_path = str(files("lerobot_robot_ax_arm") / "urdf" / "ax_arm.urdf")
-    kin = RobotKinematics(urdf_path, target_frame_name="gripper_link", joint_names=URDF_JOINT_NAMES)
+    kin = _build_kinematics(urdf_path)
+    q_lower, q_upper = kin["lower"], kin["upper"]
+    view_update = open_live_view(build_link_fks(urdf_path)) if args.view else None
 
     grip_calib = robot.calibration["gripper"]
     grip_tick = int(robot.bus.read("Present_Position", "gripper", normalize=False))
 
+    # Slowly ramp to the reference ("zero") pose (0, 45, 90) before starting teleop.
+    home_deg = np.array([REFERENCE_URDF_DEG[j] for j in ARM_JOINTS])
+    home_ticks = urdf_vector_to_ticks(home_deg, robot.calibration)
+    start_ticks = {j: float(robot.bus.read("Present_Position", j, normalize=False)) for j in ARM_JOINTS}
+    print("Homing to zero pose (0, 45, 90)...")
+    steps = max(1, int(HOME_TIME_S * FPS))
+    for i in range(1, steps + 1):
+        alpha = i / steps
+        for j in ARM_JOINTS:
+            c = robot.calibration[j]
+            tick = (1 - alpha) * start_ticks[j] + alpha * home_ticks[j]
+            robot.bus.write("Goal_Position", j, int(np.clip(tick, c.range_min, c.range_max)), normalize=False)
+        precise_sleep(1.0 / FPS)
+
     pending = {"x": 0.0, "y": 0.0, "z": 0.0, "g": 0.0}
-    state = {"quit": False}
+    state = {"quit": False, "frame": args.frame, "solver": args.solver}
     keymap = {"w": ("x", 1), "s": ("x", -1), "a": ("y", 1), "d": ("y", -1), "r": ("z", 1), "f": ("z", -1),
               "o": ("g", 1), "c": ("g", -1)}
 
@@ -74,29 +247,47 @@ def main():
         k = name.lower()
         if k in ("esc", "q"):
             state["quit"] = True
+        elif k == "t":
+            state["frame"] = "global" if state["frame"] == "local" else "local"
+        elif k == "m":
+            state["solver"] = "pos" if state["solver"] == "dq" else "dq"
         elif k in keymap:
             axis, direction = keymap[k]
             pending[axis] += direction
 
-    listener = create_key_listener(on_key, controls_help="w/s a/d r/f = XYZ, o/c = gripper, esc = stop")
+    listener = create_key_listener(
+        on_key, controls_help="w/s a/d r/f = XYZ, o/c = gripper, t = frame, m = solver, esc = stop"
+    )
     if listener is None:
         raise RuntimeError("Needs an interactive terminal with a usable key listener.")
 
-    print("Keyboard EE teleop. w/s=X a/d=Y r/f=Z o/c=gripper, ESC=stop.")
+    print("Keyboard EE teleop (velocity IK -> position). w/s=X a/d=Y r/f=Z o/c=gripper, t=frame, m=solver, ESC=stop.")
     try:
         while not state["quit"]:
             t0 = time.perf_counter()
 
             ticks = {j: float(robot.bus.read("Present_Position", j, normalize=False)) for j in ARM_JOINTS}
             q_deg = ticks_to_urdf_vector(ticks, robot.calibration)
+            q_rad = np.deg2rad(q_deg)
 
-            pose = kin.forward_kinematics(q_deg)
-            dx, dy, dz = (np.sign(pending[a]) for a in ("x", "y", "z"))
+            lin = np.array([np.sign(pending["x"]), np.sign(pending["y"]), np.sign(pending["z"])])
+            cmd_disp = lin.copy()
             pending["x"] = pending["y"] = pending["z"] = 0.0
-            pose[:3, 3] += LINEAR_STEP_M * np.array([dx, dy, dz])
 
-            q_target = kin.inverse_kinematics(q_deg, pose, orientation_weight=0.0)
-            target_ticks = urdf_vector_to_ticks(q_target, robot.calibration)
+            q_target_rad = q_rad
+            if np.any(lin):
+                q_dot = _joint_velocity(kin, q_rad, lin.astype(float), state["frame"], state["solver"])
+                # Scale for a consistent end-effector Cartesian speed (uniform across axes/poses),
+                # then cap the joint step so motion stays bounded near singularities.
+                ee_speed = float(np.linalg.norm(np.array(kin["pos_jac"](q_rad)) @ q_dot))
+                if ee_speed > 1e-6:
+                    q_step = q_dot * (cart_step / ee_speed)
+                    step_norm = np.linalg.norm(q_step)
+                    if step_norm > MAX_JOINT_STEP_RAD:
+                        q_step *= MAX_JOINT_STEP_RAD / step_norm
+                    q_target_rad = np.clip(q_rad + q_step, q_lower, q_upper)
+
+            target_ticks = urdf_vector_to_ticks(np.rad2deg(q_target_rad), robot.calibration)
             for j in ARM_JOINTS:
                 c = robot.calibration[j]
                 tick = int(np.clip(target_ticks[j], c.range_min, c.range_max))
@@ -108,10 +299,15 @@ def main():
                 grip_tick = int(np.clip(grip_tick + g * GRIP_STEP_TICK, grip_calib.range_min, grip_calib.range_max))
                 robot.bus.write("Goal_Position", "gripper", grip_tick, normalize=False)
 
-            urdf_str = " ".join(f"{j}={v:+6.1f}" for j, v in zip(ARM_JOINTS, q_target))
-            print(f"cmd[x={dx:+.0f} y={dy:+.0f} z={dz:+.0f} g={g:+.0f}] -> urdf[{urdf_str}] gripper={grip_tick}",
+            urdf_str = " ".join(f"{j}={v:+6.1f}" for j, v in zip(ARM_JOINTS, np.rad2deg(q_target_rad)))
+            print(f"[{state['frame']:>6}|{state['solver']}] cmd[x={cmd_disp[0]:+.0f} y={cmd_disp[1]:+.0f}"
+                  f" z={cmd_disp[2]:+.0f} g={g:+.0f}] -> urdf[{urdf_str}] gripper={grip_tick}",
                   end="\r", flush=True)
 
+            if view_update is not None:
+                # Draw the arm from the measured joint angles (real state, not the command).
+                view_update(q_rad, f"REAL arm | frame={state['frame']} solver={state['solver']}")
+
             precise_sleep(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
     except KeyboardInterrupt:
         pass