UFactory xArm 6 Robot Arm

"""Forward kinematics for the UFACTORY xArm 6 twin (models/kinematics.json, mm, Z-up).

Usage: python3 fk.py models/kinematics.json [j1 j2 j3 j4 j5 j6]   (degrees)
Kinematics data: UFACTORY xarm_ros2, BSD-3-Clause. This script: MIT, Adom.
"""
import json, math, sys
import numpy as np

def rpy(r, p, y):
    cr, sr, cp, sp, cy, sy = math.cos(r), math.sin(r), math.cos(p), math.sin(p), math.cos(y), math.sin(y)
    Rz = np.array([[cy, -sy, 0], [sy, cy, 0], [0, 0, 1]])
    Ry = np.array([[cp, 0, sp], [0, 1, 0], [-sp, 0, cp]])
    Rx = np.array([[1, 0, 0], [0, cr, -sr], [0, sr, cr]])
    return Rz @ Ry @ Rx

def fk(kin, q_deg):
    """Return {link_name: 4x4 world pose (mm)} for joint angles in degrees."""
    poses, qi = {}, 0
    for link in kin["links"]:
        if link["parent"] is None:
            poses[link["name"]] = np.eye(4); continue
        j = link["joint"]
        T = np.eye(4)
        T[:3, :3] = rpy(*j["origin_rpy"]); T[:3, 3] = j["origin_xyz_mm"]
        q = math.radians(q_deg[qi]); qi += 1
        Rq = np.eye(4); Rq[:3, :3] = rpy(0, 0, q)   # every axis is +Z
        poses[link["name"]] = poses[link["parent"]] @ T @ Rq
    return poses

if __name__ == "__main__":
    kin = json.load(open(sys.argv[1]))
    q = [float(a) for a in sys.argv[2:8]] if len(sys.argv) >= 8 else kin["home_deg"]
    np.set_printoptions(precision=2, suppress=True)
    print(fk(kin, q)["link6"])