Initial commit
This commit is contained in:
310
scripts/pipeline/robot_fk.py
Normal file
310
scripts/pipeline/robot_fk.py
Normal 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
|
||||
Reference in New Issue
Block a user