485 lines
16 KiB
Python
485 lines
16 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
4_robotState_estimation_v7.py
|
||
|
||
Estimate the Arm1 joint angle from ArUco marker positions.
|
||
|
||
What this script does
|
||
---------------------
|
||
1. Loads a robot configuration JSON and extracts the local marker positions of Arm1.
|
||
2. Loads an observation JSON with marker positions (e.g. aruco_positions_optimized.json).
|
||
3. Matches observed markers to Arm1 markers by marker_id.
|
||
4. Estimates the Arm1 angle in two ways:
|
||
|
||
Option 1 (single marker, optional):
|
||
----------------------------------
|
||
If a world-space joint origin is known, estimate angle from one marker by comparing
|
||
the vector from joint origin -> marker in the model and in the observation.
|
||
|
||
Option 2 (marker pairs, default and robust):
|
||
------------------------------------------
|
||
For every pair of Arm1 markers, estimate the rotation angle from the relative vector
|
||
between the two markers. This is translation-invariant and usually the preferred method.
|
||
|
||
5. Aggregates all valid estimates using a weighted circular mean and reports:
|
||
mean angle, circular variance, circular std, linear std on wrapped samples, and per-estimate details.
|
||
|
||
Coordinate conventions
|
||
----------------------
|
||
The provided robot config says:
|
||
- right-handed coordinate system
|
||
- x right, y backward, z up
|
||
|
||
Arm1 rotates around the x-axis (joint axis [-1, 0, 0] in the supplied config).
|
||
The script therefore works in 3D and automatically respects the joint axis sign.
|
||
|
||
Usage
|
||
-----
|
||
python 4_robotState_estimation_v7.py \
|
||
--robot-json robot_running_2026_05_30.json \
|
||
--aruco-json aruco_positions_optimized.json
|
||
|
||
Optional:
|
||
--joint-origin-world X Y Z # enables single-marker estimates
|
||
--output-json result.json
|
||
--arm-link Arm1
|
||
|
||
The script prints a concise report and can also write a JSON summary.
|
||
|
||
Notes
|
||
-----
|
||
- This script is intentionally defensive: it accepts positions in meters or millimeters.
|
||
- It does not require any additional robot state beyond the Arm1 marker geometry.
|
||
- If only a subset of Arm1 markers is observed, the pairwise estimator still works.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import math
|
||
from dataclasses import dataclass, asdict
|
||
from itertools import combinations
|
||
from pathlib import Path
|
||
from typing import Dict, Iterable, List, Optional, Sequence, Tuple
|
||
|
||
import numpy as np
|
||
|
||
|
||
# ----------------------------
|
||
# Data structures
|
||
# ----------------------------
|
||
|
||
@dataclass
|
||
class MarkerSpec:
|
||
marker_id: int
|
||
local_position_mm: np.ndarray
|
||
|
||
|
||
@dataclass
|
||
class MarkerObservation:
|
||
marker_id: int
|
||
position_mm: np.ndarray
|
||
link: Optional[str] = None
|
||
|
||
|
||
@dataclass
|
||
class AngleEstimate:
|
||
source: str
|
||
angle_rad: float
|
||
angle_deg: float
|
||
weight: float
|
||
marker_ids: Tuple[int, ...]
|
||
|
||
|
||
# ----------------------------
|
||
# Helpers
|
||
# ----------------------------
|
||
|
||
def _as_np3(value: Sequence[float]) -> np.ndarray:
|
||
arr = np.asarray(value, dtype=float).reshape(3)
|
||
return arr
|
||
|
||
|
||
def load_json(path: Path) -> dict:
|
||
with path.open("r", encoding="utf-8") as f:
|
||
return json.load(f)
|
||
|
||
|
||
def unit_scale_to_mm(units: Optional[str]) -> float:
|
||
if not units:
|
||
return 1.0
|
||
u = str(units).strip().lower()
|
||
if u in {"mm", "millimeter", "millimeters"}:
|
||
return 1.0
|
||
if u in {"m", "meter", "meters"}:
|
||
return 1000.0
|
||
if u in {"cm", "centimeter", "centimeters"}:
|
||
return 10.0
|
||
# conservative default
|
||
return 1.0
|
||
|
||
|
||
def normalize_angle_rad(angle: float) -> float:
|
||
"""Wrap to (-pi, pi]."""
|
||
return (angle + math.pi) % (2 * math.pi) - math.pi
|
||
|
||
|
||
def angle_between_vectors(v_from: np.ndarray, v_to: np.ndarray, axis: np.ndarray) -> float:
|
||
"""
|
||
Signed angle rotating v_from onto v_to around axis.
|
||
|
||
Uses atan2( axis_hat · (v_from × v_to), v_from · v_to ).
|
||
"""
|
||
a = np.asarray(axis, dtype=float).reshape(3)
|
||
an = np.linalg.norm(a)
|
||
if an == 0:
|
||
raise ValueError("Rotation axis must be non-zero.")
|
||
a = a / an
|
||
|
||
v1 = np.asarray(v_from, dtype=float).reshape(3)
|
||
v2 = np.asarray(v_to, dtype=float).reshape(3)
|
||
|
||
n1 = np.linalg.norm(v1)
|
||
n2 = np.linalg.norm(v2)
|
||
if n1 == 0 or n2 == 0:
|
||
raise ValueError("Cannot compute angle for zero-length vector.")
|
||
cross_term = float(np.dot(a, np.cross(v1, v2)))
|
||
dot_term = float(np.dot(v1, v2))
|
||
return math.atan2(cross_term, dot_term)
|
||
|
||
|
||
def weighted_circular_mean(angles_rad: Sequence[float], weights: Sequence[float]) -> Tuple[float, float, float]:
|
||
"""
|
||
Returns:
|
||
mean_angle_rad, circular_variance, circular_std_rad
|
||
|
||
circular_variance = 1 - R, where R is the normalized resultant length.
|
||
circular_std_rad = sqrt(-2 ln R) if R>0 else pi.
|
||
"""
|
||
if len(angles_rad) == 0:
|
||
raise ValueError("No angles provided.")
|
||
|
||
ang = np.asarray(angles_rad, dtype=float)
|
||
w = np.asarray(weights, dtype=float)
|
||
w = np.clip(w, 0.0, None)
|
||
|
||
if np.sum(w) <= 0:
|
||
w = np.ones_like(w)
|
||
|
||
s = np.sum(w * np.sin(ang))
|
||
c = np.sum(w * np.cos(ang))
|
||
mean = math.atan2(s, c)
|
||
|
||
R = math.sqrt(s * s + c * c) / float(np.sum(w))
|
||
R = min(max(R, 0.0), 1.0)
|
||
|
||
circular_variance = 1.0 - R
|
||
circular_std = math.sqrt(max(0.0, -2.0 * math.log(max(R, 1e-15)))) if R > 0 else math.pi
|
||
return mean, circular_variance, circular_std
|
||
|
||
|
||
def weighted_linear_std_rad(angles_rad: Sequence[float], mean_rad: float, weights: Sequence[float]) -> float:
|
||
if len(angles_rad) == 0:
|
||
return float("nan")
|
||
ang = np.asarray([normalize_angle_rad(a - mean_rad) for a in angles_rad], dtype=float)
|
||
w = np.asarray(weights, dtype=float)
|
||
w = np.clip(w, 0.0, None)
|
||
if np.sum(w) <= 0:
|
||
w = np.ones_like(w)
|
||
mu = np.sum(w * ang) / np.sum(w)
|
||
var = np.sum(w * (ang - mu) ** 2) / np.sum(w)
|
||
return float(math.sqrt(max(0.0, var)))
|
||
|
||
|
||
def extract_arm_markers(robot: dict, arm_link_name: str) -> Tuple[np.ndarray, List[MarkerSpec], np.ndarray]:
|
||
"""
|
||
Returns:
|
||
joint_origin_mm, markers, joint_axis
|
||
"""
|
||
links = robot.get("links", {})
|
||
if arm_link_name not in links:
|
||
# case-insensitive fallback
|
||
lookup = {k.lower(): k for k in links.keys()}
|
||
if arm_link_name.lower() not in lookup:
|
||
raise KeyError(f"Link '{arm_link_name}' not found in robot.json.")
|
||
arm_link_name = lookup[arm_link_name.lower()]
|
||
|
||
arm = links[arm_link_name]
|
||
jt = arm.get("jointToParent", {})
|
||
joint_origin_mm = _as_np3(jt.get("origin", [0, 0, 0]))
|
||
joint_axis = _as_np3(jt.get("axis", [1, 0, 0]))
|
||
|
||
markers = []
|
||
for m in arm.get("markers", []):
|
||
if "id" not in m or "position" not in m:
|
||
continue
|
||
markers.append(MarkerSpec(marker_id=int(m["id"]), local_position_mm=_as_np3(m["position"])))
|
||
|
||
if not markers:
|
||
raise ValueError(f"No markers found on link '{arm_link_name}'.")
|
||
|
||
return joint_origin_mm, markers, joint_axis
|
||
|
||
|
||
def load_observations(observations_json: dict) -> List[MarkerObservation]:
|
||
markers = observations_json.get("markers", [])
|
||
units = observations_json.get("units", {})
|
||
scale = 1.0
|
||
if isinstance(units, dict):
|
||
scale = unit_scale_to_mm(units.get("length"))
|
||
else:
|
||
scale = unit_scale_to_mm(units)
|
||
|
||
obs: List[MarkerObservation] = []
|
||
for m in markers:
|
||
if "marker_id" not in m and "id" not in m:
|
||
continue
|
||
marker_id = int(m.get("marker_id", m.get("id")))
|
||
if "position_mm" in m:
|
||
pos = _as_np3(m["position_mm"])
|
||
elif "position_m" in m:
|
||
pos = _as_np3(m["position_m"]) * 1000.0
|
||
elif "position" in m:
|
||
pos = _as_np3(m["position"]) * scale
|
||
else:
|
||
continue
|
||
obs.append(MarkerObservation(marker_id=marker_id, position_mm=pos, link=m.get("link")))
|
||
return obs
|
||
|
||
|
||
def match_markers(specs: Sequence[MarkerSpec], observations: Sequence[MarkerObservation]) -> Dict[int, Tuple[np.ndarray, np.ndarray]]:
|
||
"""
|
||
Returns marker_id -> (local_position_mm, observed_position_mm)
|
||
"""
|
||
spec_map = {m.marker_id: m.local_position_mm for m in specs}
|
||
out = {}
|
||
for o in observations:
|
||
if o.marker_id in spec_map:
|
||
out[o.marker_id] = (spec_map[o.marker_id], o.position_mm)
|
||
return out
|
||
|
||
|
||
def estimate_pairwise_angles(
|
||
matched: Dict[int, Tuple[np.ndarray, np.ndarray]],
|
||
axis: np.ndarray,
|
||
min_baseline_mm: float = 1.0,
|
||
) -> List[AngleEstimate]:
|
||
estimates: List[AngleEstimate] = []
|
||
ids = sorted(matched.keys())
|
||
|
||
for id1, id2 in combinations(ids, 2):
|
||
p1_local, p1_obs = matched[id1]
|
||
p2_local, p2_obs = matched[id2]
|
||
|
||
v_local = p2_local - p1_local
|
||
v_obs = p2_obs - p1_obs
|
||
|
||
baseline_local = float(np.linalg.norm(v_local))
|
||
baseline_obs = float(np.linalg.norm(v_obs))
|
||
if baseline_local < min_baseline_mm or baseline_obs < min_baseline_mm:
|
||
continue
|
||
|
||
try:
|
||
angle = angle_between_vectors(v_local, v_obs, axis)
|
||
except ValueError:
|
||
continue
|
||
|
||
# Weight long baselines a little more strongly; if a pair is very short, it becomes noisy.
|
||
weight = max(1e-6, baseline_local * baseline_obs)
|
||
estimates.append(
|
||
AngleEstimate(
|
||
source="pair",
|
||
angle_rad=normalize_angle_rad(angle),
|
||
angle_deg=math.degrees(normalize_angle_rad(angle)),
|
||
weight=weight,
|
||
marker_ids=(id1, id2),
|
||
)
|
||
)
|
||
return estimates
|
||
|
||
|
||
def estimate_single_marker_angles(
|
||
matched: Dict[int, Tuple[np.ndarray, np.ndarray]],
|
||
joint_origin_world_mm: np.ndarray,
|
||
axis: np.ndarray,
|
||
) -> List[AngleEstimate]:
|
||
"""
|
||
Single-marker estimates require a known world-space joint origin.
|
||
|
||
For each marker, compare:
|
||
v_local = marker_local - joint_origin_local
|
||
v_obs = marker_obs - joint_origin_world
|
||
|
||
In the provided Arm1 config the joint origin is the local origin of Arm1, so
|
||
joint_origin_local is taken as [0,0,0].
|
||
"""
|
||
estimates: List[AngleEstimate] = []
|
||
joint_origin_world_mm = _as_np3(joint_origin_world_mm)
|
||
|
||
for marker_id, (p_local, p_obs) in matched.items():
|
||
v_local = p_local # Arm1 local origin is at the joint
|
||
v_obs = p_obs - joint_origin_world_mm
|
||
|
||
if np.linalg.norm(v_local) < 1e-9 or np.linalg.norm(v_obs) < 1e-9:
|
||
continue
|
||
|
||
try:
|
||
angle = angle_between_vectors(v_local, v_obs, axis)
|
||
except ValueError:
|
||
continue
|
||
|
||
weight = float(np.linalg.norm(v_local))
|
||
estimates.append(
|
||
AngleEstimate(
|
||
source="single",
|
||
angle_rad=normalize_angle_rad(angle),
|
||
angle_deg=math.degrees(normalize_angle_rad(angle)),
|
||
weight=weight,
|
||
marker_ids=(marker_id,),
|
||
)
|
||
)
|
||
return estimates
|
||
|
||
|
||
def build_report(
|
||
arm_link_name: str,
|
||
joint_origin_mm: np.ndarray,
|
||
joint_axis: np.ndarray,
|
||
matched: Dict[int, Tuple[np.ndarray, np.ndarray]],
|
||
estimates: List[AngleEstimate],
|
||
) -> dict:
|
||
if not estimates:
|
||
raise ValueError("No valid angle estimates could be computed.")
|
||
|
||
mean_rad, circ_var, circ_std_rad = weighted_circular_mean(
|
||
[e.angle_rad for e in estimates],
|
||
[e.weight for e in estimates],
|
||
)
|
||
lin_std_rad = weighted_linear_std_rad(
|
||
[e.angle_rad for e in estimates],
|
||
mean_rad,
|
||
[e.weight for e in estimates],
|
||
)
|
||
|
||
# unwrap around mean for a stable arithmetic mean/std reporting
|
||
unwrapped_deg = []
|
||
for e in estimates:
|
||
d = math.degrees(normalize_angle_rad(e.angle_rad - mean_rad))
|
||
unwrapped_deg.append(d)
|
||
|
||
report = {
|
||
"link": arm_link_name,
|
||
"joint_origin_mm": joint_origin_mm.tolist(),
|
||
"joint_axis": joint_axis.tolist(),
|
||
"num_matched_markers": len(matched),
|
||
"num_estimates": len(estimates),
|
||
"mean_angle_rad": mean_rad,
|
||
"mean_angle_deg": math.degrees(mean_rad),
|
||
"circular_variance": circ_var,
|
||
"circular_std_rad": circ_std_rad,
|
||
"circular_std_deg": math.degrees(circ_std_rad),
|
||
"linear_std_rad": lin_std_rad,
|
||
"linear_std_deg": math.degrees(lin_std_rad),
|
||
"estimate_spread_deg_unwrapped": {
|
||
"min": float(np.min(unwrapped_deg)) if unwrapped_deg else None,
|
||
"max": float(np.max(unwrapped_deg)) if unwrapped_deg else None,
|
||
"mean_abs": float(np.mean(np.abs(unwrapped_deg))) if unwrapped_deg else None,
|
||
},
|
||
"estimates": [asdict(e) for e in estimates],
|
||
}
|
||
return report
|
||
|
||
|
||
def pretty_print_report(report: dict) -> None:
|
||
print("\nArm angle estimation report")
|
||
print("=" * 32)
|
||
print(f"Link : {report['link']}")
|
||
print(f"Matched markers : {report['num_matched_markers']}")
|
||
print(f"Valid estimates : {report['num_estimates']}")
|
||
print(f"Mean angle : {report['mean_angle_deg']:.3f} deg")
|
||
print(f"Circular variance : {report['circular_variance']:.6f}")
|
||
print(f"Circular std : {report['circular_std_deg']:.3f} deg")
|
||
print(f"Linear std (wrapped): {report['linear_std_deg']:.3f} deg")
|
||
|
||
print("\nIndividual estimates")
|
||
print("-" * 32)
|
||
for e in report["estimates"]:
|
||
mids = ",".join(str(x) for x in e["marker_ids"])
|
||
print(f"{e['source']:>6s} markers [{mids:<8s}] angle={e['angle_deg']:>8.3f} deg weight={e['weight']:.3f}")
|
||
|
||
|
||
def estimate_arm1_state(
|
||
robot_json: dict,
|
||
observations_json: dict,
|
||
arm_link_name: str = "Arm1",
|
||
joint_origin_world_mm: Optional[Sequence[float]] = None,
|
||
) -> dict:
|
||
joint_origin_mm, arm_markers, joint_axis = extract_arm_markers(robot_json, arm_link_name)
|
||
observations = load_observations(observations_json)
|
||
matched = match_markers(arm_markers, observations)
|
||
|
||
if not matched:
|
||
available = sorted(m.marker_id for m in arm_markers)
|
||
observed = sorted(o.marker_id for o in observations)
|
||
raise ValueError(
|
||
f"No Arm1 marker matches found. Arm markers: {available}. Observed IDs present in file: {observed[:25]}"
|
||
+ ("..." if len(observed) > 25 else "")
|
||
)
|
||
|
||
estimates = estimate_pairwise_angles(matched, axis=joint_axis)
|
||
|
||
# Optional single-marker estimates only if a world-space joint origin is supplied.
|
||
if joint_origin_world_mm is not None:
|
||
estimates.extend(estimate_single_marker_angles(matched, _as_np3(joint_origin_world_mm), axis=joint_axis))
|
||
|
||
report = build_report(
|
||
arm_link_name=arm_link_name,
|
||
joint_origin_mm=joint_origin_mm,
|
||
joint_axis=joint_axis,
|
||
matched=matched,
|
||
estimates=estimates,
|
||
)
|
||
return report
|
||
|
||
|
||
def main() -> int:
|
||
parser = argparse.ArgumentParser(description="Estimate Arm1 angle from ArUco marker positions.")
|
||
parser.add_argument("--robot-json", type=Path, default=Path("robot.json"), help="Path to robot config JSON.")
|
||
parser.add_argument("--aruco-json", type=Path, default=Path("aruco_positions_optimized.json"), help="Path to marker observation JSON.")
|
||
parser.add_argument("--arm-link", type=str, default="Arm1", help="Name of the link to estimate.")
|
||
parser.add_argument(
|
||
"--joint-origin-world",
|
||
type=float,
|
||
nargs=3,
|
||
default=None,
|
||
metavar=("X", "Y", "Z"),
|
||
help="Optional world-space joint origin for single-marker estimates.",
|
||
)
|
||
parser.add_argument("--output-json", type=Path, default=None, help="Optional path to write the result JSON.")
|
||
args = parser.parse_args()
|
||
|
||
robot_json = load_json(args.robot_json)
|
||
observations_json = load_json(args.aruco_json)
|
||
|
||
report = estimate_arm1_state(
|
||
robot_json=robot_json,
|
||
observations_json=observations_json,
|
||
arm_link_name=args.arm_link,
|
||
joint_origin_world_mm=args.joint_origin_world,
|
||
)
|
||
|
||
pretty_print_report(report)
|
||
|
||
if args.output_json:
|
||
args.output_json.parent.mkdir(parents=True, exist_ok=True)
|
||
with args.output_json.open("w", encoding="utf-8") as f:
|
||
json.dump(report, f, indent=2, ensure_ascii=False)
|
||
print(f"\nWrote JSON summary to: {args.output_json}")
|
||
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|