Initial commit

This commit is contained in:
chk
2026-06-08 19:50:36 +02:00
commit 53db55ba36
39 changed files with 5860 additions and 0 deletions

View File

@@ -0,0 +1,310 @@
#!/usr/bin/env python3
"""
robot_fk.py
-----------
Minimal forward kinematics engine for the robot.json format.
Matches the Blender hierarchy used by render_robot.py exactly:
world_T_link = world_T_parent
@ Translate(mountPosition) @ Rotate_xyz(mountRotation)
@ Translate(jointOrigin) @ Rotate_xyz(joint.rotation)
@ T_motion
T_motion = Rotate(axis, value_deg) for revolute joints
Translate(axis * value_mm) for linear joints
Units throughout: millimetres, degrees.
Public API
----------
fk = RobotFK.from_file("robot.json")
transforms = fk.compute({"x": 180, "y": 86, "z": -120,
"a": -60, "b": 22, "c": 91, "e": 10})
# → dict link_name -> 4×4 np.ndarray (world frame, mm)
p_world = fk.marker_world(transforms, "Arm1", [0, -160, 35])
# → np.ndarray shape (3,), in mm
all_m = fk.all_markers_world(transforms)
# → dict marker_id -> {"world_mm", "link", "local_mm"}
# Cumulative x-offset for a link at all-zero state
# (useful for 4a: x_slider = world_x - local_x - link_x_at_zero)
x0 = fk.link_x_at_zero_state("Arm1") # → float mm
"""
from __future__ import annotations
import json
import math
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple
import numpy as np
STATE_KEYS = ("x", "y", "z", "a", "b", "c", "e")
# ──────────────────────────────────────────────────────────────
# Low-level matrix helpers
# ──────────────────────────────────────────────────────────────
def _rot_axis_angle(axis: Sequence[float], angle_deg: float) -> np.ndarray:
"""3×3 rotation matrix via Rodrigues (axis need not be normalised)."""
ax = np.asarray(axis, dtype=float)
n = float(np.linalg.norm(ax))
if n < 1e-12:
return np.eye(3)
ax = ax / n
c = math.cos(math.radians(angle_deg))
s = math.sin(math.radians(angle_deg))
t = 1.0 - c
x, y, z = ax
return np.array([
[t*x*x + c, t*x*y - s*z, t*x*z + s*y],
[t*x*y + s*z, t*y*y + c, t*y*z - s*x],
[t*x*z - s*y, t*y*z + s*x, t*z*z + c ],
])
def _rot_xyz_euler(rx: float, ry: float, rz: float) -> np.ndarray:
"""XYZ Euler angles (degrees) → 3×3 — matches Blender XYZ Euler mode."""
return (_rot_axis_angle([0, 0, 1], rz)
@ _rot_axis_angle([0, 1, 0], ry)
@ _rot_axis_angle([1, 0, 0], rx))
def make_T(R: np.ndarray, t: Sequence[float]) -> np.ndarray:
"""4×4 homogeneous transform."""
T = np.eye(4)
T[:3, :3] = R
T[:3, 3] = np.asarray(t, dtype=float)
return T
def transform_point(T: np.ndarray, p: Sequence[float]) -> np.ndarray:
"""Apply 4×4 transform to a 3-D point."""
h = np.array([p[0], p[1], p[2], 1.0])
return (T @ h)[:3]
# ──────────────────────────────────────────────────────────────
# FK engine
# ──────────────────────────────────────────────────────────────
class RobotFK:
"""Forward kinematics for the robot.json format."""
def __init__(self, robot_data: Dict[str, Any]):
self.robot = robot_data
self.links: Dict[str, Any] = robot_data.get("links", {})
self._topo: List[str] = self._topological_sort()
# ── construction ─────────────────────────────────────────
@classmethod
def from_file(cls, path) -> "RobotFK":
with open(path, "r", encoding="utf-8") as f:
return cls(json.load(f))
def _topological_sort(self) -> List[str]:
parent_map = {n: d.get("parent") for n, d in self.links.items()}
visited, order = set(), []
def visit(name: str) -> None:
if name in visited:
return
visited.add(name)
p = parent_map.get(name)
if p and p in self.links:
visit(p)
order.append(name)
for name in self.links:
visit(name)
return order
# ── core computation ──────────────────────────────────────
def compute(self, joint_values: Dict[str, float]) -> Dict[str, np.ndarray]:
"""
Compute link world transforms for the given joint state.
Parameters
----------
joint_values : dict variable_name -> value
Linear joints (x, e): mm
Revolute joints (y, z, a, b, c): degrees
Returns
-------
dict link_name -> 4×4 np.ndarray (world frame, mm)
"""
state = {k: 0.0 for k in STATE_KEYS}
for k, v in joint_values.items():
if k in state:
state[k] = float(v)
transforms: Dict[str, np.ndarray] = {}
for link_name in self._topo:
d = self.links[link_name]
parent = d.get("parent")
T_parent = transforms.get(parent, np.eye(4)) if parent else np.eye(4)
# 1 · Mount (static in parent frame)
mp = d.get("mountPosition", [0, 0, 0])
mr = d.get("mountRotation", [0, 0, 0])
T_m = make_T(_rot_xyz_euler(*mr), mp)
# 2 · Joint origin (pivot point in mount frame)
ji = d.get("jointToParent", {}) or {}
jp = ji.get("origin", [0, 0, 0])
jr = ji.get("rotation", [0, 0, 0])
T_j = make_T(_rot_xyz_euler(*jr), jp)
# 3 · Motion
jtype = str(ji.get("type", "fixed")).lower()
var = str(ji.get("variable", "")).lower()
axis = np.asarray(ji.get("axis", [1, 0, 0]), dtype=float)
val = state.get(var, 0.0)
if jtype == "revolute":
T_mot = make_T(_rot_axis_angle(axis, val), [0, 0, 0])
elif jtype == "linear":
n = float(np.linalg.norm(axis))
u = axis / n if n > 1e-12 else np.array([1.0, 0, 0])
T_mot = make_T(np.eye(3), u * val)
else:
T_mot = np.eye(4)
transforms[link_name] = T_parent @ T_m @ T_j @ T_mot
return transforms
# ── marker helpers ────────────────────────────────────────
@staticmethod
def marker_world(transforms: Dict[str, np.ndarray],
link_name: str,
local_pos: Sequence[float]) -> np.ndarray:
"""Transform a local marker position → world (mm)."""
return transform_point(transforms.get(link_name, np.eye(4)), local_pos)
def all_markers_world(self,
transforms: Dict[str, np.ndarray]
) -> Dict[int, Dict[str, Any]]:
"""
Returns
-------
dict marker_id -> {world_mm, local_mm, link, normal_world}
"""
result: Dict[int, Dict[str, Any]] = {}
for lname, ldata in self.links.items():
T = transforms.get(lname, np.eye(4))
R = T[:3, :3]
for m in ldata.get("markers", []):
mid = int(m.get("id", -1))
if mid < 0 or "position" not in m:
continue
loc = np.array(m["position"], dtype=float)
nor = np.array(m.get("normal", [0, 0, 1]), dtype=float)
result[mid] = {
"world_mm": transform_point(T, loc),
"local_mm": loc,
"link": lname,
"normal_world": (R @ nor) / max(np.linalg.norm(R @ nor), 1e-12),
}
return result
# ── x-axis invariant helpers (used by 4a) ────────────────
def link_x_at_zero_state(self, link_name: str) -> float:
"""
Return the world x-coordinate of the link-frame origin
when ALL joint variables are zero.
For any link reachable via only x-axis rotations (Arm1, Ellbow, Arm2),
this value is constant regardless of the actual revolute angles.
Adding the slider value x_mm gives the true link origin x:
link_origin_world_x = x_mm + link_x_at_zero_state(link_name)
And for any marker in that link:
marker_world_x = x_mm + link_x_at_zero_state(link_name) + marker_local_x
"""
T = self.compute({k: 0.0 for k in STATE_KEYS})
return float(T[link_name][0, 3])
def joint_origin_world(self,
link_name: str,
joint_state: Dict[str, float]) -> np.ndarray:
"""
World position of the pivot that link_name rotates / slides around.
"""
d = self.links[link_name]
parent = d.get("parent")
T_all = self.compute(joint_state)
T_parent = T_all.get(parent, np.eye(4)) if parent else np.eye(4)
mp = d.get("mountPosition", [0, 0, 0])
mr = d.get("mountRotation", [0, 0, 0])
T_m = make_T(_rot_xyz_euler(*mr), mp)
ji = d.get("jointToParent", {}) or {}
jp = ji.get("origin", [0, 0, 0])
jr = ji.get("rotation", [0, 0, 0])
T_j = make_T(_rot_xyz_euler(*jr), jp)
return transform_point(T_parent @ T_m @ T_j, [0, 0, 0])
def joint_axis_world(self,
link_name: str,
joint_state: Dict[str, float]) -> np.ndarray:
"""
Joint axis of link_name expressed in world frame.
"""
d = self.links[link_name]
parent = d.get("parent")
T_all = self.compute(joint_state)
T_parent = T_all.get(parent, np.eye(4)) if parent else np.eye(4)
mp = d.get("mountPosition", [0, 0, 0])
mr = d.get("mountRotation", [0, 0, 0])
T_m = make_T(_rot_xyz_euler(*mr), mp)
ji = d.get("jointToParent", {}) or {}
jp = ji.get("origin", [0, 0, 0])
jr = ji.get("rotation", [0, 0, 0])
T_j = make_T(_rot_xyz_euler(*jr), jp)
R_to_pivot = (T_parent @ T_m @ T_j)[:3, :3]
axis_local = np.asarray(ji.get("axis", [1, 0, 0]), dtype=float)
world = R_to_pivot @ axis_local
n = float(np.linalg.norm(world))
return world / n if n > 1e-12 else world
# ── utility ───────────────────────────────────────────────
def chain(self, link_name: str) -> List[str]:
"""Return chain from root to link_name (inclusive)."""
out, cur = [], link_name
while cur:
out.append(cur)
cur = self.links.get(cur, {}).get("parent")
return list(reversed(out))
def board_marker_world_positions(self, length_scale: float = 1.0) -> Dict[int, np.ndarray]:
"""
Return known world positions for all Board markers (in mm).
Board is the root, so its marker positions ARE world positions.
length_scale: 1/1000 if robot.json uses mm.
"""
board = self.links.get("Board", {})
result: Dict[int, np.ndarray] = {}
for m in board.get("markers", []):
mid = int(m.get("id", -1))
if mid >= 0 and "position" in m:
p = np.array(m["position"], dtype=float) * length_scale
result[mid] = p
return result