#!/usr/bin/env python """Hardware-free 3D simulation of the AX-arm keyboard EE teleop. Reuses the real IK (``_build_kinematics`` / ``_joint_velocity`` from ``teleoperate_ee_keyboard``) and a synthetic calibration, driving a simulated (ideal) servo bus instead of a real one. Lets you sanity-check the solver/frame behaviour and joint limits in a matplotlib window. Controls (focus the plot window): - w / s : +X / -X - a / d : +Y / -Y - r / f : +Z / -Z - o / c : open / close gripper - t : toggle global <-> local frame - m : toggle dq <-> pos solver - q / esc : quit Run: python examples/simulate_ee_keyboard.py [--solver pos] [--frame global] [--speed 0.06] """ import argparse import importlib.util import os import tempfile from pathlib import Path os.environ.setdefault("MPLCONFIGDIR", tempfile.gettempdir()) import numpy as np import matplotlib.pyplot as plt from matplotlib.animation import FuncAnimation from lerobot.motors import MotorCalibration from lerobot_robot_ax_arm.urdf_mapping import ( ARM_JOINTS, REFERENCE_URDF_DEG, SCALE, URDF_LIMITS_DEG, ticks_to_urdf_vector, urdf_vector_to_ticks, ) # Reuse the real teleop IK helpers without duplicating them. _EE = importlib.util.spec_from_file_location( "_ee_teleop", str(Path(__file__).with_name("teleoperate_ee_keyboard.py")) ) ee = importlib.util.module_from_spec(_EE) _EE.loader.exec_module(ee) TICK_REF = 512 # tick chosen to sit at each joint's URDF reference angle GRIP_RANGE = (350, 600) def _synthetic_calibration() -> dict[str, MotorCalibration]: """Calibration consistent with urdf_mapping: reference tick + travel limits per joint.""" calib: dict[str, MotorCalibration] = {} for i, j in enumerate(ARM_JOINTS): lo_deg, hi_deg = URDF_LIMITS_DEG[j] ticks = sorted(int(round(TICK_REF + (d - REFERENCE_URDF_DEG[j]) / SCALE)) for d in (lo_deg, hi_deg)) calib[j] = MotorCalibration(id=i + 1, drive_mode=0, homing_offset=TICK_REF, range_min=ticks[0], range_max=ticks[1]) calib["gripper"] = MotorCalibration(id=4, drive_mode=0, homing_offset=0, range_min=GRIP_RANGE[0], range_max=GRIP_RANGE[1]) return calib class _SimBus: """Ideal servo bus: Present_Position instantly follows the last commanded Goal_Position.""" def __init__(self, ticks: dict[str, float]): self.ticks = ticks def read(self, _reg, motor, normalize=False): return self.ticks[motor] def write(self, _reg, motor, value, normalize=False): self.ticks[motor] = float(value) def main(): parser = argparse.ArgumentParser() parser.add_argument("--speed", type=float, default=ee.CART_STEP_M) parser.add_argument("--frame", choices=("local", "global"), default="local") parser.add_argument("--solver", choices=("dq", "pos"), default="dq") args = parser.parse_args() from importlib.resources import files urdf_path = str(files("lerobot_robot_ax_arm") / "urdf" / "ax_arm.urdf") kin = ee._build_kinematics(urdf_path) link_fks = ee.build_link_fks(urdf_path) calib = _synthetic_calibration() q_start_deg = np.array([REFERENCE_URDF_DEG[j] for j in ARM_JOINTS]) # reference "zero" pose (0, 45, 90) start_ticks = urdf_vector_to_ticks(q_start_deg, calib) ticks = {j: float(start_ticks[j]) for j in ARM_JOINTS} ticks["gripper"] = float(sum(GRIP_RANGE) / 2) bus = _SimBus(ticks) pending = {"x": 0.0, "y": 0.0, "z": 0.0, "g": 0.0} state = {"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)} 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") reach = 0.32 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") def on_key(event): k = (event.key or "").lower() if k in ("q", "escape"): plt.close(fig) 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 fig.canvas.mpl_connect("key_press_event", on_key) def update(_frame): q_rad = np.deg2rad(ticks_to_urdf_vector({j: ticks[j] for j in ARM_JOINTS}, calib)) cmd = np.array([np.sign(pending["x"]), np.sign(pending["y"]), np.sign(pending["z"])], dtype=float) pending["x"] = pending["y"] = pending["z"] = 0.0 if np.any(cmd): q_dot = ee._joint_velocity(kin, q_rad, cmd, state["frame"], state["solver"]) ee_speed = float(np.linalg.norm(np.array(kin["pos_jac"](q_rad)) @ q_dot)) if ee_speed > 1e-6: q_step = q_dot * (args.speed / ee_speed) step_norm = np.linalg.norm(q_step) if step_norm > ee.MAX_JOINT_STEP_RAD: q_step *= ee.MAX_JOINT_STEP_RAD / step_norm q_target = np.clip(q_rad + q_step, kin["lower"], kin["upper"]) target_ticks = urdf_vector_to_ticks(np.rad2deg(q_target), calib) for j in ARM_JOINTS: c = calib[j] bus.write("Goal_Position", j, int(np.clip(target_ticks[j], c.range_min, c.range_max))) g = np.sign(pending["g"]); pending["g"] = 0.0 if g: gc = calib["gripper"] bus.write("Goal_Position", "gripper", int(np.clip(ticks["gripper"] + g * ee.GRIP_STEP_TICK, gc.range_min, gc.range_max))) pts = ee.chain_points(link_fks, np.deg2rad(ticks_to_urdf_vector({j: ticks[j] for j in ARM_JOINTS}, calib))) 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]) grip_frac = (ticks["gripper"] - GRIP_RANGE[0]) / (GRIP_RANGE[1] - GRIP_RANGE[0]) ax.set_title(f"frame={state['frame']} solver={state['solver']} gripper={grip_frac:.0%}\n" f"w/s/a/d/r/f=move o/c=grip t=frame m=solver q=quit") return chain_line, ee_pt anim = FuncAnimation(fig, update, interval=int(1000 / ee.FPS), blit=False, cache_frame_data=False) fig._anim = anim # keep a reference so it isn't garbage-collected plt.show() if __name__ == "__main__": main()