From 87766114947bc5272a3e851dcec599b3e6ea6cca Mon Sep 17 00:00:00 2001 From: Filip Maric Date: Thu, 11 Jun 2026 07:25:44 +0200 Subject: [PATCH 1/9] fix: correctness issues from the 2026-06-10 code review Squash of the review-driven correctness branch: - manifolds: correct 2D Lyapunov matrix in PSDFixedRank projection - sdp: repair broken entry points; SDP tests run without Mosek - utils: RobotURDF raises ValueError instead of returning string literals - solvers: LocalSolver honors all goals and uses the true SE2 left Jacobian; convex iteration no longer mutates the caller's anchors dict - BOUNDED edge data consistently checked via constants with .get() defaults across sdp_snl, joint_angle_solver, graph, dgp - docs: fix malformed ICRA BibTeX entry; add_anchor_node typo and triple-computed norm --- .gitignore | 3 + README.md | 4 +- graphik/graphs/graph.py | 35 ++-- graphik/solvers/convex_iteration.py | 5 +- graphik/solvers/joint_angle_solver.py | 25 ++- graphik/solvers/sdp_snl.py | 28 ++- graphik/utils/dgp.py | 10 +- graphik/utils/manifolds/fixed_rank_psd_sym.py | 19 +- graphik/utils/roboturdf.py | 30 ++- tests/test_adjacency_matrix.py | 4 +- tests/test_convex_iteration.py | 30 +++ tests/test_distance_limits.py | 33 +++- tests/test_distance_matrix.py | 30 ++- tests/test_joint_angle_solver.py | 80 ++++++++ tests/test_psd_manifold.py | 55 ++++++ tests/test_roboturdf.py | 57 ++++++ tests/test_sdp_snl.py | 43 ++++- tests/test_sdp_snl_inequality_constraints.py | 176 ++++++++++++------ 18 files changed, 531 insertions(+), 136 deletions(-) create mode 100644 tests/test_convex_iteration.py create mode 100644 tests/test_joint_angle_solver.py create mode 100644 tests/test_psd_manifold.py create mode 100644 tests/test_roboturdf.py diff --git a/.gitignore b/.gitignore index 8ed1d70..d0e576e 100644 --- a/.gitignore +++ b/.gitignore @@ -137,5 +137,8 @@ dmypy.json # Local-only planning/design notes docs/superpowers/ +# Local Claude/Codex-adjacent session state +.claude/ + # Per-environment solver baselines (regenerate with tests/baselines/generate.py) tests/baselines/*.json diff --git a/README.md b/README.md index 4b2c18a..ab5ecf2 100644 --- a/README.md +++ b/README.md @@ -48,7 +48,7 @@ To enable the Mosek path: 3. Place the license file (`mosek.lic`) at the path Mosek expects (typically `~/mosek/mosek.lic`). -The SDP tests in `tests/test_sdp_snl*.py` are tuned to Mosek's tolerances and skip cleanly when it isn't installed. +The SDP tests in `tests/test_sdp_snl*.py` run without Mosek: constraint-construction tests are solver-free, and the end-to-end solves use whichever free solver cvxpy picks, with tolerances loose enough for it. ## Usage Use of GraphIK can be summarized by four key steps, which we'll walk through below (see the scripts in [experiments/](https://github.com/utiasSTARS/GraphIK/tree/main/experiments) for more details). @@ -176,7 +176,7 @@ arXiv: [Inverse Kinematics for Serial Kinematic Chains via Sum of Squares Optimi ```bibtex @misc{marić2022convex_arxiv, - author={Filip Marić and {Matthew Giamou and Soroush Khoubyarian and Ivan Petrović and Jonathan Kelly}, + author={Filip Marić and Matthew Giamou and Soroush Khoubyarian and Ivan Petrović and Jonathan Kelly}, title={Inverse Kinematics for Serial Kinematic Chains via Sum of Squares Optimization}, year={2020}, eprint={1909.09318}, diff --git a/graphik/graphs/graph.py b/graphik/graphs/graph.py index 0c20036..8598819 100644 --- a/graphik/graphs/graph.py +++ b/graphik/graphs/graph.py @@ -13,6 +13,10 @@ from graphik.utils import * # noqa: F401,F403 (existing convention; brings constants and helpers into scope: ROOT, POS, TYPE, BASE, ROBOT, END_EFFECTOR, OBSTACLE, BOUNDED, BELOW, ABOVE, LOWER, UPPER, DIST, MAIN_PREFIX, AUX_PREFIX, TRANSFORM, trans_axis, rot_axis, skew, normalize, wraptopi, R2, best_fit_transform, level2_descendants, max_min_distance_revolute, distance_matrix_from_graph, adjacency_matrix_from_graph, graph_complete_edges) +def _bounded_tags(limit) -> List[str]: + return [limit] if limit in (BELOW, ABOVE) else [] + + class ProblemGraph(nx.DiGraph): """ Graph with positions and distance bounds describing an IK problem. @@ -124,14 +128,15 @@ def from_pose(self, T_goal) -> nx.DiGraph: # ------------------------------------------------------------------ def add_anchor_node(self, name: str, data: Dict[str, Any]): if POS not in data: - raise KeyError("Node needs to gave a position to be added.") + raise KeyError("Node needs to have a position to be added.") self.add_nodes_from([(name, data)]) for nname, ndata in self.nodes(data=True): if POS in ndata and nname != name: + dist = la.norm(ndata[POS] - data[POS]) self.add_edge(nname, name) - self[nname][name][DIST] = la.norm(ndata[POS] - data[POS]) - self[nname][name][LOWER] = la.norm(ndata[POS] - data[POS]) - self[nname][name][UPPER] = la.norm(ndata[POS] - data[POS]) + self[nname][name][DIST] = dist + self[nname][name][LOWER] = dist + self[nname][name][UPPER] = dist self[nname][name][BOUNDED] = [] def add_spherical_obstacle(self, name: str, position: ArrayLike, radius: float): @@ -156,7 +161,8 @@ def check_distance_limits(self, G: nx.DiGraph, tol=1e-10) -> List[Dict[str, List typ = nx.get_node_attributes(self, name=TYPE) broken_limits = [] for u, v, data in self.edges(data=True): - if BELOW in data[BOUNDED] or ABOVE in data[BOUNDED]: + bounds = data.get(BOUNDED, []) + if BELOW in bounds or ABOVE in bounds: if G[u][v][DIST] < data[LOWER] - tol: bl = {} if (ROBOT in typ[u] and OBSTACLE in typ[v]) or ( @@ -199,10 +205,11 @@ def distance_bound_matrices(self) -> Tuple[ArrayLike, ArrayLike]: if BOUNDED in data: udx = self.node_ids.index(e1) vdx = self.node_ids.index(e2) - if BELOW in data[BOUNDED]: + bounds = data[BOUNDED] + if BELOW in bounds: L[udx, vdx] = data[LOWER] ** 2 L[vdx, udx] = L[udx, vdx] - if ABOVE in data[BOUNDED]: + if ABOVE in bounds: U[udx, vdx] = data[UPPER] ** 2 U[vdx, udx] = U[udx, vdx] return L, U @@ -277,6 +284,12 @@ def _joint_nodes(self, cur: str) -> List[Tuple[str, np.ndarray]]: (aux_name, (T0 @ trans_z)[:3, 3]), ] + def structure_graph(self) -> nx.DiGraph: + """Standalone copy of the robot structure subgraph (exact-distance + edges only, no joint-limit bound edges). Public API used by the SDP + solvers; equivalent to the pre-unification ``structure_graph()``.""" + return self._build_structure_subgraph() + def _build_structure_subgraph(self) -> nx.DiGraph: structure = nx.empty_graph(create_using=nx.DiGraph) end_effectors = self.robot.end_effectors @@ -364,7 +377,7 @@ def _root_angle_limits_2d(self): self[ax][node][LOWER] = sqrt( l1 ** 2 + l2 ** 2 - 2 * l1 * l2 * cos(pi - lim) ) - self[ax][node][BOUNDED] = BELOW + self[ax][node][BOUNDED] = [BELOW] def _set_limits_2d(self): S = self.structure @@ -384,7 +397,7 @@ def _set_limits_2d(self): self[u][v][LOWER] = sqrt( l1 ** 2 + l2 ** 2 - 2 * l1 * l2 * cos(pi - lim) ) - self[u][v][BOUNDED] = BELOW + self[u][v][BOUNDED] = [BELOW] # --- 3D bodies (verbatim from the old ProblemGraphRevolute) --- def _root_angle_limits_3d(self): @@ -442,7 +455,7 @@ def _root_angle_limits_3d(self): self.add_edge(base_node, node) if d_max == d_min: self[base_node][node][DIST] = d_max - self[base_node][node][BOUNDED] = [limit] + self[base_node][node][BOUNDED] = _bounded_tags(limit) self[base_node][node][UPPER] = d_max self[base_node][node][LOWER] = d_min @@ -506,7 +519,7 @@ def _set_limits_3d(self): self.add_edge(ids[0], ids[1]) if d_max == d_min: S[ids[0]][ids[1]][DIST] = d_max - self[ids[0]][ids[1]][BOUNDED] = [limit] + self[ids[0]][ids[1]][BOUNDED] = _bounded_tags(limit) self[ids[0]][ids[1]][UPPER] = d_max self[ids[0]][ids[1]][LOWER] = d_min diff --git a/graphik/solvers/convex_iteration.py b/graphik/solvers/convex_iteration.py index 60bceb5..565bf22 100644 --- a/graphik/solvers/convex_iteration.py +++ b/graphik/solvers/convex_iteration.py @@ -157,7 +157,7 @@ def get_sparsity_pattern(G, canonical_point_order: list) -> set: def convex_iterate_sdp_snl_graph( graph: ProblemGraph, - anchors: dict = {}, + anchors: dict = None, ranges=False, max_iters=10, sparse=False, @@ -181,6 +181,9 @@ def convex_iterate_sdp_snl_graph( d = robot.dim eig_value_sum_vs_iterations = [] + # Copy: anchors gets extended below and must not leak into the + # caller's dict (or across calls via a shared default). + anchors = dict(anchors) if anchors else {} # If a position is pre-defined for a node, set to anchor for node, data in G.nodes(data=True): if data.get(POS, None) is not None: diff --git a/graphik/solvers/joint_angle_solver.py b/graphik/solvers/joint_angle_solver.py index 5fead31..acbfe74 100644 --- a/graphik/solvers/joint_angle_solver.py +++ b/graphik/solvers/joint_angle_solver.py @@ -25,7 +25,7 @@ def __init__(self, robot_graph: ProblemGraph, params: Dict["str", Any]): typ = nx.get_node_attributes(self.graph, name=TYPE) pairs = [] for u, v, data in self.graph.edges(data=True): - if "below" in data[BOUNDED]: + if BELOW in data.get(BOUNDED, []): if ROBOT in typ[u] and OBSTACLE in typ[v] and u != ROOT: pairs += [(u, v)] self.m = len(pairs) @@ -143,7 +143,7 @@ def gen_cost_and_grad_ee(self, point: str, T_goal: np.ndarray): log = SE2.Log inverse = SE2.inverse adjoint = SE2.adjoint - inv_left_jacobian = lambda x: np.eye(3) + inv_left_jacobian = SE2.left_jacobian_inv def cost(q): q_dict = {joints[idx]: q[idx] for idx in range(n)} @@ -159,16 +159,27 @@ def cost(q): return cost def solve(self, goals: dict, q0: dict): - for node, goal in goals.items(): - cost_and_grad = self.gen_cost_and_grad_ee(node, goal) + # Each per-goal cost reads the leading entries of q along its + # root-to-goal chain (q0 must be ordered along the kinematic chain) + # and returns a full-width gradient (robot.jacobian pads with zero + # columns past the chain), so goals simply sum. + goal_costs = [ + self.gen_cost_and_grad_ee(node, goal) for node, goal in goals.items() + ] + + def cost_and_grad(q): + f = 0.0 + grad = np.zeros(self.robot.n) + for cg in goal_costs: + f_i, g_i = cg(q) + f += f_i + grad += g_i + return f, grad - # solve res = minimize( cost_and_grad, - # cost, np.asarray(list(q0.values())), jac=True, - # jac=grad, constraints=self.g, method="SLSQP", options={"ftol": 1e-7}, diff --git a/graphik/solvers/sdp_snl.py b/graphik/solvers/sdp_snl.py index 552cf31..3374415 100644 --- a/graphik/solvers/sdp_snl.py +++ b/graphik/solvers/sdp_snl.py @@ -14,6 +14,15 @@ from graphik.solvers.sdp_formulations import SdpSolverParams, solve_sdp +_INFEASIBLE_STATUSES = frozenset({cp.INFEASIBLE, cp.INFEASIBLE_INACCURATE}) + + +def status_is_infeasible(status) -> bool: + """True when a cvxpy problem status reports (possibly inaccurate) + infeasibility. Compares by equality; cvxpy statuses are plain strings.""" + return status in _INFEASIBLE_STATUSES + + def prepare_set_cover_problem( constraint_clique_dict: dict, nearest_points: dict, d: int ): @@ -30,7 +39,7 @@ def prepare_set_cover_problem( # We also only need to augment non-zero nearest points (important for nuclear norm case) for target in nearest_points.keys(): if target in targets_to_cover and np.all(nearest_points[target] == np.zeros(d)): - targets_to_cover.remove() + targets_to_cover.remove(target) return cliques_remaining, targets_to_cover @@ -269,20 +278,22 @@ def distance_clique_linear_map( def distance_constraints_graph( G: nx.Graph, - anchors: dict = {}, + anchors: dict = None, sparse: bool = False, ee_cost=None, angle_limits=False, ) -> dict: + anchors = anchors if anchors is not None else {} G = G.copy() if angle_limits: # Remove the edges that don't have a typ = nx.get_edge_attributes(G, name=BOUNDED) edges = [] for u, v, data in G.edges(data=True): + bounds = data.get(BOUNDED, []) if (DIST not in list(data.keys())) and ( - "below" not in data[BOUNDED] and "above" not in data[BOUNDED] + BELOW not in bounds and ABOVE not in bounds ): edges += [(u, v)] @@ -363,14 +374,13 @@ def distance_range_constraints( if ( u not in anchors.keys() or v not in anchors.keys() ): # If BOTH are anchors we can ignore - if data.get(BOUNDED, False): - if ( - "below" in data[BOUNDED] - ): # TODO: All the bounds are only listed as below, ask Filip! + bounds = data.get(BOUNDED, []) + if bounds: + if BELOW in bounds: # TODO: All the bounds are only listed as below, ask Filip! pairs += [frozenset((u, v))] dists += [data[LOWER]] upper += [False] - if "above" in data[BOUNDED]: + if ABOVE in bounds: pairs += [frozenset((u, v))] dists += [data[UPPER]] upper += [True] @@ -959,7 +969,7 @@ def solve_linear_cost_sdp( if solver_error: solution = SOLVER_ERROR - elif prob.status is INFEASIBLE: + elif status_is_infeasible(prob.status): solution = INFEASIBLE else: solution = extract_solution(constraint_clique_dict, sdp_variable_map, robot.dim) diff --git a/graphik/utils/dgp.py b/graphik/utils/dgp.py index 04e6cd8..f134b0d 100644 --- a/graphik/utils/dgp.py +++ b/graphik/utils/dgp.py @@ -101,7 +101,7 @@ def pos_from_graph(G: nx.DiGraph, node_ids=None) -> NDArray: def graph_from_pos( - P: NDArray, node_ids: Optional[List] = [], dist: bool = True + P: NDArray, node_ids: Optional[List] = None, dist: bool = True ) -> nx.DiGraph: """ Generates an nx.DiGraph object of the subclass type given @@ -242,10 +242,12 @@ def bound_smoothing(G: nx.DiGraph) -> tuple: def normalize_positions(Y: NDArray, scale=False): Y_c = Y - Y.mean(0) C = Y_c.T.dot(Y_c) - e, v = np.linalg.eig(C) + _, v = np.linalg.eigh(C) Y_cr = Y_c.dot(v) if scale: - Y_crs = Y_cr / (1 / abs(Y_cr).max()) - return Y_crs + max_abs = np.abs(Y_cr).max() + if max_abs > 0: + return Y_cr / max_abs + return Y_cr else: return Y_cr diff --git a/graphik/utils/manifolds/fixed_rank_psd_sym.py b/graphik/utils/manifolds/fixed_rank_psd_sym.py index 4e67388..e3cd96c 100644 --- a/graphik/utils/manifolds/fixed_rank_psd_sym.py +++ b/graphik/utils/manifolds/fixed_rank_psd_sym.py @@ -76,24 +76,9 @@ def typical_dist(self): @staticmethod def _build_lyap_matrix(X): # X = Y^T Y. Returns the (dim*dim) x (dim*dim) Lyapunov matrix - # whose linear system gives the quotient correction Omega. + # representing Omega -> X@Omega + Omega@X on row-major vec(Omega). dim = X.shape[0] - if dim == 3: - A = np.asarray([[X[0,0] + X[0,0], X[0,1] , X[0,2], X[1,0], 0, 0, X[2,0], 0, 0], - [X[1,0], X[1,1] + X[0,0], X[1,2], 0, X[1,0], 0, 0, X[2,0], 0], - [X[2,0], X[2,1], X[2,2] + X[0,0], 0, 0, X[1,0], 0, 0, X[2,0]], - [X[0,1], 0, 0, X[0,0] + X[1,1], X[0,1] , X[0,2], X[2,1], 0, 0], - [0, X[0,1], 0, X[1,0], X[1,1] + X[1,1], X[1,2], 0, X[2,1], 0], - [0, 0, X[0,1], X[2,0], X[2,1], X[2,2] + X[1,1], 0, 0, X[2,1]], - [X[0,2], 0, 0, X[1,2], 0, 0, X[0,0] + X[2,2], X[0,1] , X[0,2]], - [0, X[0,2], 0, 0, X[1,2], 0, X[1,0], X[1,1] + X[2,2], X[1,2]], - [0, 0, X[0,2], 0, 0, X[1,2], X[2,0], X[2,1], X[2,2] + X[2,2]]]) - else: # dim == 2 - A = np.asarray([[X[0,0] + X[0,0], X[0,1], X[0,1], 0], - [X[1,0], X[0,1] + X[0,0], 0, X[0,1]], - [X[0,1], 0, X[0,0] + X[1,1], X[0,1]], - [0, X[0,1], X[1,0], X[1,1]+ X[1,1]]]) - return A + return np.kron(X, np.eye(dim)) + np.kron(np.eye(dim), X) @staticmethod def _kron_with_eye(A, d): diff --git a/graphik/utils/roboturdf.py b/graphik/utils/roboturdf.py index 7cac559..3600bcb 100644 --- a/graphik/utils/roboturdf.py +++ b/graphik/utils/roboturdf.py @@ -106,7 +106,7 @@ def find_joints_child_joints_from_list(self, joint, joints): def get_parents(self, joints): base_joint = self.find_first_joint() if not (base_joint in joints): - raise ("Base joint not in joints") + raise ValueError("Base joint not in joints") label_base = "p{0}" # parents = {'p0': [label_base.format(joints.index(base_joint))]} @@ -127,7 +127,7 @@ def actuated_joint_index(self, joint): try: return self.urdf.actuated_joints.index(joint) except ValueError: - raise ("joint not an actuated joint") + raise ValueError(f"{joint.name} is not an actuated joint") from None def find_link_by_name(self, name): for link in self._links: @@ -276,15 +276,9 @@ def make_Revolute3d(self, ub, lb, randomized_links = False, randomize_percentage params["T_zero"] = T_zero params["num_joints"] = self.n_q_joints - l = 0 - for cl in self.parents.values(): - l += len(cl) - if l == len(self.parents.keys()) - 1: - params["joint_limits_upper"] = ub - params["joint_limits_lower"] = lb - return Robot({**params, "dim": 3}) - else: - return Robot({**params, "dim": 3}) + params["joint_limits_upper"] = ub + params["joint_limits_lower"] = lb + return Robot({**params, "dim": 3}) def get_T_from_joint_axis(axis: np.ndarray): """ @@ -392,7 +386,7 @@ def load_ur10(limits=None, randomized_links = False, randomize_percentage = 0.4) return robot, graph -def load_truncated_ur10(n: int): +def load_truncated_ur10(n: int, limits=None): """ Produce a robot and graph representing the first n links of a UR10. """ @@ -404,16 +398,20 @@ def load_truncated_ur10(n: int): d = d_full[0:n] al = al_full[0:n] th = th_full[0:n] - ub = (np.pi) * np.ones(n) - lb = -ub + if limits is None: + ub = np.pi * np.ones(n) + lb = -ub + else: + lb = limits[0] + ub = limits[1] modified_dh = False params = { "a": a[:n], "alpha": al[:n], "d": d[:n], "theta": th[:n], - "lb": lb[:n], - "ub": ub[:n], + "joint_limits_lower": lb[:n], + "joint_limits_upper": ub[:n], "modified_dh": modified_dh, "num_joints": n, } diff --git a/tests/test_adjacency_matrix.py b/tests/test_adjacency_matrix.py index 023904d..0b1686c 100644 --- a/tests/test_adjacency_matrix.py +++ b/tests/test_adjacency_matrix.py @@ -291,8 +291,8 @@ def test_planar_tree_pose_goal(self): "link_lengths": a, "theta": th, "parents": parents, - "ub": lim_u, - "lb": lim_l, + "joint_limits_upper": lim_u, + "joint_limits_lower": lim_l, "num_joints": n } diff --git a/tests/test_convex_iteration.py b/tests/test_convex_iteration.py new file mode 100644 index 0000000..4d79634 --- /dev/null +++ b/tests/test_convex_iteration.py @@ -0,0 +1,30 @@ +"""Argument-handling tests for convex_iterate_sdp_snl_graph. + +max_iters=0 exercises all the setup (anchor collection, constraint +construction) without running any SDP solve. +""" +import numpy as np + +from graphik.solvers.convex_iteration import convex_iterate_sdp_snl_graph +from graphik.utils.constants import POS +from graphik.utils.roboturdf import load_truncated_ur10 + + +def test_caller_anchors_dict_is_not_mutated(): + robot, graph = load_truncated_ur10(3) + # An obstacle node carries a POS attribute, so the anchor-collection + # loop would add it to the caller's dict if anchors weren't copied. + graph.add_spherical_obstacle("o0", np.array([10.0, 10.0, 10.0]), 0.5) + + n = robot.n + anchors = { + "p0": graph.nodes["p0"][POS], + "q0": graph.nodes["q0"][POS], + f"p{n}": np.array([0.5, 0.5, 0.5]), + f"q{n}": np.array([0.5, 0.5, 1.5]), + } + keys_before = set(anchors) + + convex_iterate_sdp_snl_graph(graph, anchors, max_iters=0) + + assert set(anchors) == keys_before diff --git a/tests/test_distance_limits.py b/tests/test_distance_limits.py index 4989712..433ee11 100644 --- a/tests/test_distance_limits.py +++ b/tests/test_distance_limits.py @@ -28,7 +28,7 @@ from graphik.graphs import ProblemGraph from graphik.robots import Robot from graphik.utils import list_to_variable_dict -from graphik.utils.constants import DIST +from graphik.utils.constants import ABOVE, BELOW, BOUNDED, DIST from graphik.utils.roboturdf import RobotURDF @@ -69,6 +69,37 @@ def _check_invariants(self, graph, robot, n_trials=N_TRIALS, seed=SEED, tol=TOL) f"expected {expected:.8f}, got {actual:.8f}" ) + def _check_bounded_tags(self, graph): + allowed = {BELOW, ABOVE} + for u, v, data in graph.edges(data=True): + self.assertIn(BOUNDED, data, msg=f"edge {(u, v)} missing {BOUNDED}") + self.assertIsInstance( + data[BOUNDED], + list, + msg=f"edge {(u, v)} stores non-list BOUNDED={data[BOUNDED]!r}", + ) + self.assertTrue( + set(data[BOUNDED]).issubset(allowed), + msg=f"edge {(u, v)} stores invalid BOUNDED={data[BOUNDED]!r}", + ) + + def test_bounded_attributes_are_normalized(self): + fname = graphik.__path__[0] + "/robots/urdfs/ur10_mod.urdf" + urdf_robot = RobotURDF(fname) + n = urdf_robot.n_q_joints + robot_3d = urdf_robot.make_Revolute3d(pi * np.ones(n), -pi * np.ones(n)) + self._check_bounded_tags(ProblemGraph(robot_3d)) + + n = 6 + params = { + "link_lengths": list_to_variable_dict(np.ones(n)), + "theta": list_to_variable_dict(np.zeros(n)), + "joint_limits_upper": list_to_variable_dict(0.8 * pi * np.ones(n)), + "joint_limits_lower": list_to_variable_dict(-0.8 * pi * np.ones(n)), + "num_joints": n, + } + self._check_bounded_tags(ProblemGraph(Robot({**params, "dim": 2}))) + def test_revolute_ur10(self): fname = graphik.__path__[0] + "/robots/urdfs/ur10_mod.urdf" urdf_robot = RobotURDF(fname) diff --git a/tests/test_distance_matrix.py b/tests/test_distance_matrix.py index 70e9c90..23afd19 100644 --- a/tests/test_distance_matrix.py +++ b/tests/test_distance_matrix.py @@ -6,9 +6,36 @@ from graphik.graphs import ProblemGraph from graphik.robots import Robot -from graphik.utils import best_fit_transform, list_to_variable_dict, MDS, gram_from_distance_matrix, pos_from_graph +from graphik.utils import ( + best_fit_transform, + list_to_variable_dict, + MDS, + gram_from_distance_matrix, + normalize_positions, + pos_from_graph, +) class TestDistanceMatrix(unittest.TestCase): + def test_normalize_positions_scale_uses_unit_max_abs(self): + Y = np.array([ + [2.0, 0.0], + [0.0, 4.0], + [-2.0, -2.0], + ]) + + P = normalize_positions(Y, scale=True) + + self.assertTrue(np.isrealobj(P)) + assert_allclose(P.mean(axis=0), np.zeros(2), atol=1e-12) + assert_allclose(np.abs(P).max(), 1.0) + + def test_normalize_positions_scale_handles_degenerate_input(self): + Y = np.ones((4, 3)) + + P = normalize_positions(Y, scale=True) + + assert_allclose(P, np.zeros_like(Y), atol=1e-12) + def test_special_case_3d_tree(self): print("\n Testing cases with zeroed out parameters ...") modified_dh = False @@ -200,4 +227,3 @@ def test_random_params_2d_tree(self): P_e = (R @ X.T + t.reshape(2, 1)).T self.assertIsNone(assert_allclose(P_e, Y, atol=1e-8)) - diff --git a/tests/test_joint_angle_solver.py b/tests/test_joint_angle_solver.py new file mode 100644 index 0000000..b5ec325 --- /dev/null +++ b/tests/test_joint_angle_solver.py @@ -0,0 +1,80 @@ +"""Behavior tests for LocalSolver (joint-space SLSQP solver).""" +import numpy as np +import pytest + +from graphik.robots import Robot +from graphik.graphs import ProblemGraph +from graphik.solvers.joint_angle_solver import LocalSolver +from graphik.utils.utils import list_to_variable_dict + + +def planar_chain(n): + params = { + "link_lengths": list_to_variable_dict(np.ones(n)), + "theta": list_to_variable_dict(np.zeros(n)), + "joint_limits_upper": list_to_variable_dict(np.pi * np.ones(n)), + "joint_limits_lower": list_to_variable_dict(-np.pi * np.ones(n)), + "num_joints": n, + "dim": 2, + } + robot = Robot(params) + return robot, ProblemGraph(robot) + + +class TestGradientConsistency2d: + def test_cost_and_grad_matches_finite_differences(self): + """The 2D branch must use the true SE2 inverse left Jacobian; a + stubbed identity makes the returned gradient inconsistent with the + returned cost away from the goal.""" + n = 4 + robot, graph = planar_chain(n) + solver = LocalSolver(graph, {}) + point = f"p{n}" + + q_goal = np.array([0.4, -0.3, 0.5, 0.2]) + T_goal = robot.pose(list_to_variable_dict(q_goal), point) + cost_and_grad = solver.gen_cost_and_grad_ee(point, T_goal) + + q = np.array([-0.2, 0.6, -0.1, 0.3]) + _, grad = cost_and_grad(q) + + eps = 1e-6 + fd = np.zeros(n) + for i in range(n): + dq = np.zeros(n) + dq[i] = eps + f_plus, _ = cost_and_grad(q + dq) + f_minus, _ = cost_and_grad(q - dq) + fd[i] = (f_plus - f_minus) / (2 * eps) + + np.testing.assert_allclose(grad, fd, rtol=1e-4, atol=1e-6) + + +class TestMultiGoalSolve: + def test_solve_optimizes_all_goals_not_just_the_last(self): + """Two consistent goals (intermediate point and end effector, both + from the same configuration). If solve() drops every goal but the + last dict entry, the end-effector goal here is silently ignored.""" + n = 6 + robot, graph = planar_chain(n) + solver = LocalSolver(graph, {}) + + q_goal = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6]) + q_goal_dict = list_to_variable_dict(q_goal) + T_ee = robot.pose(q_goal_dict, f"p{n}") + T_mid = robot.pose(q_goal_dict, "p3") + # End-effector goal first, intermediate goal last: the historical + # bug kept only the last entry. + goals = {f"p{n}": T_ee, "p3": T_mid} + + # Start with the first three joints already at the goal, the rest + # off: the p3 cost alone is zero here, so a solver that ignores the + # p{n} goal terminates immediately. + q0 = np.concatenate([q_goal[:3], np.zeros(3)]) + res = solver.solve(goals, list_to_variable_dict(q0)) + + q_sol = list_to_variable_dict(res.x) + err_ee = np.linalg.norm(robot.pose(q_sol, f"p{n}") - T_ee) + err_mid = np.linalg.norm(robot.pose(q_sol, "p3") - T_mid) + assert err_mid < 1e-3 + assert err_ee < 1e-3 diff --git a/tests/test_psd_manifold.py b/tests/test_psd_manifold.py new file mode 100644 index 0000000..ea03ad6 --- /dev/null +++ b/tests/test_psd_manifold.py @@ -0,0 +1,55 @@ +"""Contract tests for the PSDFixedRank quotient manifold. + +The projection onto the horizontal space at Y must be idempotent and land in +the horizontal space {U : Y^T U symmetric}; the hand-coded Lyapunov matrix must +represent the operator Omega -> X@Omega + Omega@X (row-major vec). Both are +checked for k=2 (planar problems) and k=3. +""" +import numpy as np +import pytest + +from graphik.utils.manifolds.fixed_rank_psd_sym import PSDFixedRank + + +@pytest.fixture(params=[2, 3], ids=["k2", "k3"]) +def k(request): + return request.param + + +def _random_Y(n, k, seed): + rng = np.random.default_rng(seed) + return rng.standard_normal((n, k)) + + +def test_lyapunov_matrix_represents_sylvester_operator(k): + rng = np.random.default_rng(0) + Y = rng.standard_normal((4, k)) + X = Y.T @ Y + A = PSDFixedRank._build_lyap_matrix(X) + for seed in range(3): + Omega = np.random.default_rng(seed).standard_normal((k, k)) + expected = X @ Omega + Omega @ X + np.testing.assert_allclose( + A @ Omega.ravel(), expected.ravel(), atol=1e-12, + err_msg=f"Lyapunov matrix wrong for k={k}", + ) + + +def test_projection_is_idempotent(k): + n = 8 + man = PSDFixedRank(n, k) + Y = _random_Y(n, k, seed=1) + Z = _random_Y(n, k, seed=2) + P1 = man.projection(Y, Z) + P2 = man.projection(Y, P1) + np.testing.assert_allclose(P2, P1, atol=1e-10) + + +def test_projection_output_is_horizontal(k): + n = 8 + man = PSDFixedRank(n, k) + Y = _random_Y(n, k, seed=3) + Z = _random_Y(n, k, seed=4) + P = man.projection(Y, Z) + skew_part = Y.T @ P - P.T @ Y + np.testing.assert_allclose(skew_part, np.zeros((k, k)), atol=1e-10) diff --git a/tests/test_roboturdf.py b/tests/test_roboturdf.py new file mode 100644 index 0000000..c6ef9c4 --- /dev/null +++ b/tests/test_roboturdf.py @@ -0,0 +1,57 @@ +"""Error-path tests for RobotURDF.""" +import numpy as np +import pytest +from numpy.testing import assert_allclose + +import graphik +from graphik.utils.roboturdf import RobotURDF, load_truncated_ur10, load_ur10 + + +@pytest.fixture(scope="module") +def ur10_urdf(): + return RobotURDF(graphik.__path__[0] + "/robots/urdfs/ur10_mod.urdf") + + +def test_get_parents_without_base_joint_raises_value_error(ur10_urdf): + joints = list(ur10_urdf.T_zero.keys()) + base_joint = ur10_urdf.find_first_joint() + joints_without_base = [j for j in joints if j is not base_joint] + with pytest.raises(ValueError): + ur10_urdf.get_parents(joints_without_base) + + +def test_actuated_joint_index_of_fixed_joint_raises_value_error(ur10_urdf): + fixed_joint = next( + j for j in ur10_urdf._joints if j not in ur10_urdf.urdf.actuated_joints + ) + with pytest.raises(ValueError): + ur10_urdf.actuated_joint_index(fixed_joint) + + +def _bounds_as_arrays(robot): + joint_ids = [f"p{idx}" for idx in range(1, robot.n + 1)] + lb = np.array([robot.lb[joint_id] for joint_id in joint_ids]) + ub = np.array([robot.ub[joint_id] for joint_id in joint_ids]) + return lb, ub + + +def test_urdf_loader_preserves_explicit_joint_limits(): + lb = np.array([-0.25, -0.35, -0.45, -0.55, -0.65, -0.75]) + ub = np.array([0.25, 0.35, 0.45, 0.55, 0.65, 0.75]) + + robot, _ = load_ur10(limits=(lb, ub)) + + robot_lb, robot_ub = _bounds_as_arrays(robot) + assert_allclose(robot_lb, lb) + assert_allclose(robot_ub, ub) + + +def test_truncated_ur10_loader_preserves_explicit_joint_limits(): + lb = np.array([-0.25, -0.35, -0.45]) + ub = np.array([0.25, 0.35, 0.45]) + + robot, _ = load_truncated_ur10(3, limits=(lb, ub)) + + robot_lb, robot_ub = _bounds_as_arrays(robot) + assert_allclose(robot_lb, lb) + assert_allclose(robot_ub, ub) diff --git a/tests/test_sdp_snl.py b/tests/test_sdp_snl.py index cb6c7dc..80261b0 100644 --- a/tests/test_sdp_snl.py +++ b/tests/test_sdp_snl.py @@ -1,6 +1,3 @@ -import pytest -pytest.importorskip("mosek") # SDP tests require commercial mosek; skip when unavailable - import numpy as np import unittest import networkx as nx @@ -15,6 +12,8 @@ evaluate_linear_map, constraints_and_nearest_points_to_sdp_vars, evaluate_cost, + prepare_set_cover_problem, + status_is_infeasible, ) from graphik.utils.constants import * from graphik.utils.roboturdf import load_ur10 @@ -139,6 +138,38 @@ def run_cost_test(test_case, robot, graph, sparse=False, ee_cost=False): test_case.assertAlmostEqual(cost_bad, cost_bad_explicit) +class TestPrepareSetCover(unittest.TestCase): + def test_zero_nearest_point_is_dropped_from_targets(self): + """Zero nearest points need no cost augmentation (nuclear-norm case): + they must be dropped from the cover targets, not crash.""" + clique = frozenset(("p1", "p2")) + constraint_clique_dict = { + clique: ([np.zeros((2, 2))], [0.0], {"p1": 0, "p2": 1}, False) + } + nearest_points = {"p1": np.zeros(3), "p2": np.ones(3)} + + cliques_remaining, targets_to_cover = prepare_set_cover_problem( + constraint_clique_dict, nearest_points, d=3 + ) + + self.assertEqual(targets_to_cover, ["p2"]) + self.assertEqual(cliques_remaining, {clique}) + + +class TestStatusIsInfeasible(unittest.TestCase): + def test_detects_infeasible_by_equality_not_identity(self): + # cvxpy returns plain strings; detection must not rely on interning. + non_interned = "".join(["in", "feasible"]) + self.assertTrue(status_is_infeasible(non_interned)) + + def test_detects_inaccurate_infeasible(self): + self.assertTrue(status_is_infeasible("infeasible_inaccurate")) + + def test_optimal_is_not_infeasible(self): + self.assertFalse(status_is_infeasible("optimal")) + self.assertFalse(status_is_infeasible("optimal_inaccurate")) + + class TestUR10(unittest.TestCase): def setUp(self): self.robot, self.graph = load_ur10() @@ -233,12 +264,14 @@ def setUp(self): "alpha": al[:n], "d": d[:n], "theta": th[:n], - "lb": lb[:n], - "ub": ub[:n], + "joint_limits_lower": lb[:n], + "joint_limits_upper": ub[:n], "modified_dh": modified_dh, "num_joints": n } self.robot = Robot({**params, "dim": 3}) + self.assertEqual(self.robot.lb["p1"], lb[0]) + self.assertEqual(self.robot.ub["p1"], ub[0]) self.graph = ProblemGraph(self.robot) def test_cost(self): diff --git a/tests/test_sdp_snl_inequality_constraints.py b/tests/test_sdp_snl_inequality_constraints.py index 4b8508f..008bfb1 100644 --- a/tests/test_sdp_snl_inequality_constraints.py +++ b/tests/test_sdp_snl_inequality_constraints.py @@ -1,67 +1,125 @@ -import pytest -pytest.importorskip("mosek") # SDP tests require commercial mosek; skip when unavailable +"""Inequality-constraint construction and an end-to-end nearest-point SDP. +These run on whatever SDP solver cvxpy has available (Clarabel by default); +no Mosek required. Tolerances are kept loose accordingly. +""" import numpy as np +import pytest + from graphik.solvers.sdp_snl import ( - cvxpy_inequality_constraints, - distance_inequality_constraint, + distance_constraints_graph, distance_range_constraints, solve_nearest_point_sdp, - distance_constraints, ) -from graphik.utils.roboturdf import load_ur10 -from graphik.solvers.sdp_formulations import SdpSolverParams from graphik.solvers.constraints import get_full_revolute_nearest_point +from graphik.utils.constants import POS +from graphik.utils.roboturdf import load_ur10, load_truncated_ur10 + +import networkx as nx + + +def _clique_gram(clique_entry, positions, d): + """Build Z = X^T X for a clique from a {node: position} dict, matching the + layout used by the constraint LMEs (augmented cliques carry eye(d)).""" + A, _, mapping, is_augmented = clique_entry + n_lme = A[0].shape[0] + n_vars = n_lme - d if is_augmented else n_lme + X = np.zeros((d, n_vars)) + for var, idx in mapping.items(): + if isinstance(var, str) and var in positions: + X[:, idx] = positions[var] + if is_augmented: + X = np.hstack([X, np.eye(d)]) + return X.T @ X + + +class TestObstacleInequalityConstruction: + def setup_method(self): + np.random.seed(7) + self.robot, self.graph = load_ur10() + # An obstacle far from the zero-configuration workspace boundary + # cannot be violated by the true configuration used below. + self.obstacle_pos = np.array([10.0, 10.0, 10.0]) + self.radius = 0.5 + self.graph.add_spherical_obstacle("o0", self.obstacle_pos, self.radius) + + n = self.robot.n + q = self.robot.random_configuration() + full_points = [f"p{i}" for i in range(n + 1)] + [f"q{i}" for i in range(n + 1)] + self.input_vals = get_full_revolute_nearest_point(self.graph, q, full_points) + + G = nx.DiGraph(self.graph) + G.remove_node("x") + G.remove_node("y") + self.anchors = { + key: self.input_vals[key] for key in ["p0", "q0", f"p{n}", f"q{n}"] + } + for node, data in G.nodes(data=True): + if data.get(POS, None) is not None: + self.anchors[node] = data[POS] + self.G = G + self.constraint_clique_dict = distance_constraints_graph( + G, self.anchors, sparse=False, angle_limits=True + ) + self.inequality_map = distance_range_constraints( + G, self.constraint_clique_dict, self.anchors + ) + + def test_obstacle_produces_inequality_constraints(self): + n_ineq = sum(len(v) for v in self.inequality_map.values()) + assert n_ineq > 0 + + def test_feasible_configuration_satisfies_inequalities(self): + positions = {**self.input_vals, "o0": self.obstacle_pos} + d = self.robot.dim + for clique, constraints in self.inequality_map.items(): + Z = _clique_gram(self.constraint_clique_dict[clique], positions, d) + for A_ineq, b in constraints: + assert np.trace(A_ineq @ Z) <= b + 1e-9 + + def test_point_inside_obstacle_violates_an_inequality(self): + # Move every robot point onto the obstacle centre: clearance bounds + # (lower bounds on distance) must be violated. + positions = {key: self.obstacle_pos.copy() for key in self.input_vals} + positions["o0"] = self.obstacle_pos + d = self.robot.dim + violations = 0 + for clique, constraints in self.inequality_map.items(): + Z = _clique_gram(self.constraint_clique_dict[clique], positions, d) + for A_ineq, b in constraints: + if np.trace(A_ineq @ Z) > b + 1e-9: + violations += 1 + assert violations > 0 + + +class TestNearestPointSdpFreeSolver: + def test_nuclear_norm_solve_satisfies_equality_constraints(self): + np.random.seed(11) + robot, graph = load_truncated_ur10(3) + n = robot.n + q = robot.random_configuration() + full_points = [f"p{i}" for i in range(n + 1)] + [f"q{i}" for i in range(n + 1)] + input_vals = get_full_revolute_nearest_point(graph, q, full_points) + end_effectors = { + key: input_vals[key] for key in ["p0", "q0", f"p{n}", f"q{n}"] + } + # Zero nearest points = nuclear-norm cost; also exercises the + # zero-target drop in prepare_set_cover_problem. + nearest_points = { + key: np.zeros(robot.dim) + for key in input_vals + if key not in end_effectors + } + + solution, prob, constraint_clique_dict, sdp_variable_map = ( + solve_nearest_point_sdp(nearest_points, end_effectors, graph) + ) -if __name__ == "__main__": - - sparse = False # Whether to exploit chordal sparsity in the SDP formulation - ee_cost = False # Whether to treat the end-effectors as variables with targets in the cost. - # If False, end-effectors are NOT variables (they are baked in to constraints as parameters) - conic_solver = "MOSEK" # One of "MOSEK", "CVXOPT" for now - solver_params = ( - SdpSolverParams() - ) # Use MOSEK settings that worked well for us before - - robot, graph = load_ur10() - n = robot.n - - # Generate a random feasible target - q = robot.random_configuration() - - # Extract the positions of the points - full_points = [f"p{idx}" for idx in range(0, graph.robot.n + 1)] + [ - f"q{idx}" for idx in range(0, graph.robot.n + 1) - ] - input_vals = get_full_revolute_nearest_point(graph, q, full_points) - - # End-effectors are 'generalized' to include the base pair ('p0', 'q0') - end_effectors = { - key: input_vals[key] for key in ["p0", "q0", f"p{robot.n}", f"q{robot.n}"] - } # NOTE not sure if using the term end-effector for the base is sensible - - # NOTE both functions should really take in the output of graph.complete_from_pos() and parse that - # Form the equality constraints - constraint_clique_dict = distance_constraints(graph, end_effectors, sparse, ee_cost) - - # Form the inequality constraints - inequality_map = distance_range_constraints( - graph.structure, constraint_clique_dict, end_effectors - ) # this will only add joint limits - - # Nuclear norm cost (simple) - nearest_points_nuclear = { - key: np.zeros(robot.dim) - for key in input_vals - if key not in ["p0", "q0", f"p{robot.n}", f"q{robot.n}"] - } - - solution, prob, constraint_clique_dict, sdp_variable_map = solve_nearest_point_sdp( - nearest_points_nuclear, - end_effectors, - graph, - sparse=False, - solver_params=None, - verbose=True, - inequality_constraints_map=inequality_map, - ) + assert prob.status in ("optimal", "optimal_inaccurate") + assert set(nearest_points.keys()) <= set(solution.keys()) + # The SDP variable must satisfy the distance LMEs to solver tolerance. + for clique, Z_var in sdp_variable_map.items(): + A, b, _, _ = constraint_clique_dict[clique] + Z = Z_var.value + residuals = [np.trace(A[i] @ Z) - b[i] for i in range(len(A))] + assert np.max(np.abs(residuals)) < 1e-5 From ef261284413cf28465deb84cc32fc739e6ed3891 Mon Sep 17 00:00:00 2001 From: Filip Maric Date: Wed, 10 Jun 2026 07:02:30 +0200 Subject: [PATCH 2/9] refactor: remove dead code identified in the 2026-06-10 review Deletions (no callers anywhere in package, tests, or experiments; verified by grep and a full-package import sweep): - graphik/utils/operators.py (entire module) and the dgp.py helpers that only it consumed (orthogonal_procrustes, factor, sample_matrix, incidence_matrix_from_adjacency) - graphik/utils/chordal.py: verbatim copy of networkx's complete_to_chordal_graph; sdp_snl now calls nx's directly - graphik/solvers/constraints.py: SOS-era sympy layer, broken since the 2D/3D unification (robot_graph.directed, parent_node_id). The one live function, get_full_revolute_nearest_point, moved to sdp_snl.py - utils.py: dZ, perpendicular_vector, norm_sq, list_to_variable_dict_- spherical, variable_dict_to_list, make_save_string, spherical_angle_- bounds_to_revolute, safe_arccos, bernoulli/wilson confidence helpers, measure_perturbation, constraint_violations, __main__ block - kinematics.py: fk_2d (only self-recursive) - robot.py: end_effector_pos (iterated characters of a node name), jacobian_geometric (referenced removed .parents), lambdified, spherical - graph.py: distance_bounds_from_sampling (dead, wrote garbage DIST) - joint_angle_solver.py: gen_objective_ee/gen_grad_ee (superseded by gen_cost_and_grad_ee), main() with ~60 lines of commented-out demo, unused imports - sdp_snl.py: sym_vec, constraints_list_to_matrix, ~130-line __main__ Mosek demo, unused typ assignment - convex_iteration.py: random_psd_matrix, solve_fantope_sdp (+ its only-callers fantope_constraints, solve_fantope_iterate), get_sparsity_pattern, ~150-line __main__ benchmark - robot_visualization.py: import-time rc('text', usetex=True) side effect (broke import without LaTeX) is now an opt-in use_latex_text() setup.py: sympy and progress dependencies dropped (only dead code used them). Suite: 91 passed. Co-Authored-By: Claude Fable 5 --- graphik/graphs/graph.py | 26 -- graphik/robots/robot.py | 45 ---- graphik/solvers/constraints.py | 265 ------------------- graphik/solvers/convex_iteration.py | 224 ---------------- graphik/solvers/joint_angle_solver.py | 119 +-------- graphik/solvers/sdp_snl.py | 192 ++------------ graphik/utils/chordal.py | 66 ----- graphik/utils/dgp.py | 39 --- graphik/utils/kinematics.py | 7 - graphik/utils/operators.py | 81 ------ graphik/utils/robot_visualization.py | 11 +- graphik/utils/utils.py | 152 ----------- setup.py | 2 - tests/test_pymlg_conventions.py | 12 +- tests/test_sdp_snl.py | 2 +- tests/test_sdp_snl_inequality_constraints.py | 2 +- 16 files changed, 31 insertions(+), 1214 deletions(-) delete mode 100644 graphik/solvers/constraints.py delete mode 100644 graphik/utils/chordal.py delete mode 100644 graphik/utils/operators.py diff --git a/graphik/graphs/graph.py b/graphik/graphs/graph.py index 8598819..ab0300e 100644 --- a/graphik/graphs/graph.py +++ b/graphik/graphs/graph.py @@ -647,29 +647,3 @@ def get_pose( T_trans = trans_axis(self.axis_length, "z") T = T @ T_trans return T - - # ------------------------------------------------------------------ - # 3D-only sampling helper (keep dim-agnostic — works for any dim) - # ------------------------------------------------------------------ - def distance_bounds_from_sampling(self): - robot = self.robot - ids = self.node_ids - q_rand = robot.random_configuration() - D_min = self.distance_matrix_from_joints(q_rand) - D_max = self.distance_matrix_from_joints(q_rand) - - for _ in range(2000): - q_rand = robot.random_configuration() - D_rand = self.distance_matrix_from_joints(q_rand) - D_max[D_rand > D_max] = D_rand[D_rand > D_max] - D_min[D_rand < D_min] = D_rand[D_rand < D_min] - - for idx in range(len(D_max)): - for jdx in range(len(D_max)): - e1 = ids[idx] - e2 = ids[jdx] - self.add_edge(e1, e2) - self[e1][e2][LOWER] = sqrt(D_min[idx, jdx]) - self[e1][e2][UPPER] = sqrt(D_max[idx, jdx]) - if abs(D_max[idx, jdx] - D_min[idx, jdx]) < 1e-5: - self[e1][e2][DIST] = abs(D_max[idx, jdx] - D_min[idx, jdx]) diff --git a/graphik/robots/robot.py b/graphik/robots/robot.py index 0071bdd..38fad23 100644 --- a/graphik/robots/robot.py +++ b/graphik/robots/robot.py @@ -36,7 +36,6 @@ def __init__(self, params: Dict): raise ValueError(f"dim must be 2 or 3, got {self.dim}") self._SE = SE2 if self.dim == 2 else SE3 - self.lambdified = False self.params = params self.n = params["num_joints"] @@ -127,9 +126,6 @@ def lb(self) -> Dict[str, Any]: def lb(self, lb: dict): self._lb = lb if type(lb) is dict else list_to_variable_dict(flatten([lb])) - @property - def spherical(self) -> bool: - return False # ------------------------------------------------------------------ # Configuration helpers @@ -155,13 +151,6 @@ def get_all_poses(self, joint_angles: Dict[str, Any]) -> Dict[str, np.ndarray]: T[node] = self.pose(joint_angles, node) return T - def end_effector_pos(self, q: Dict[str, float]) -> Dict[str, ArrayLike]: - goals = {} - for ee in self.end_effectors: - for node in ee: - goals[node] = self.pose(q, node)[:-1, -1] - return goals - # ------------------------------------------------------------------ # Dim-aware kinematics # ------------------------------------------------------------------ @@ -265,37 +254,3 @@ def jacobian( Ad = self._SE.adjoint(T) J[node][:, idx] = Ad @ self.nodes[pred]["S"] return J - - def jacobian_geometric( - self, - joint_angles: Dict[str, float], - nodes: Union[List[str], str], - Ts: Dict[str, np.ndarray] = None, - ) -> Dict[str, ArrayLike]: - """Geometric Jacobian (linear+angular) for end-effector p-nodes. 3D-only.""" - assert self.dim == 3, "jacobian_geometric is 3D-only" - kmap = self.kinematic_map[ROOT] - - if nodes is None: - nodes = [] - for ee in self.end_effectors: - if ee[0][0] == MAIN_PREFIX: - nodes += [ee[0]] - elif ee[1][0] == MAIN_PREFIX: - nodes += [ee[1]] - - if Ts is None: - Ts = self.get_all_poses(joint_angles) - - J = {} - for node in nodes: - path = kmap[node][1:] - p_ee = Ts[node][:3, 3] - J[node] = np.zeros([6, self.n]) - for idx, joint in enumerate(path): - T_0_i = Ts[list(self.parents.predecessors(joint))[0]] - z_hat_i = T_0_i[:3, 2] - p_i = T_0_i[:3, 3] - J[node][:3, idx] = np.cross(z_hat_i, p_ee - p_i) - J[node][3:, idx] = z_hat_i - return J diff --git a/graphik/solvers/constraints.py b/graphik/solvers/constraints.py deleted file mode 100644 index 14fb17c..0000000 --- a/graphik/solvers/constraints.py +++ /dev/null @@ -1,265 +0,0 @@ -from copy import deepcopy -from graphik.graphs import ProblemGraph -import networkx as nx -import numpy as np -import sympy as sp -import numpy.linalg as la -from numpy import pi -from graphik.utils.utils import norm_sq -from graphik.utils.dgp import pos_from_graph -from graphik.robots import Robot - - -def apply_angular_offset_2d(joint_angle_offset, z0): - """ - :param joint_angle_offset: - :param z0: - :return: - """ - R = np.zeros((2, 2)) - R[0, 0] = np.cos(joint_angle_offset) - R[0, 1] = np.sin(joint_angle_offset) - R[1, 0] = -np.sin(joint_angle_offset) - R[1, 1] = R[0, 0] - return R * z0 - - -def archimedean_constraint(variables, M_bounds): - return [var ** 2 <= M for var, M in zip(variables, M_bounds)] - - -def all_symbols(constraints): - all_vars = set() - for cons in constraints: - all_vars = all_vars.union(set(cons.free_symbols)) - return list(all_vars) - - -def generate_symbols( - robot_graph: ProblemGraph, nodes_with_angular_vars: dict = None -) -> nx.DiGraph: - """ - Make a copy of robot_graph's such that the nodes have "sym" fields populated. - Also populate the angular_variables field if angular residual constraint variables are needed. - - :param robot_graph: - :param nodes_with_angular_vars: dictionary whose keys are the nodes we want to have angular vars - :return: - """ - G = deepcopy(robot_graph.directed) # Make a copy of the problem graph - - # Assign symbolic variables where needed - for node_id in G: - if robot_graph.dim == 3: - G.nodes[node_id]["sym"] = np.array( - sp.symbols(f"{node_id}_x, {node_id}_y, {node_id}_z") - ) - elif robot_graph.dim == 2: - G.nodes[node_id]["sym"] = np.array(sp.symbols(f"{node_id}_x, {node_id}_y")) - if node_id in list(robot_graph.base.nodes()): - G.nodes[node_id]["root"] = True - - if nodes_with_angular_vars is not None: - G.angular_variables = {} - for node_id in nodes_with_angular_vars: - G.angular_variables[node_id] = sp.symbols("s_" + node_id) - # G.angular_variables = np.array(sp.symbols("s:"+str(n_angular))) - # else: - # angular_variables = - return G - - -def constraints_from_graph( - robot_graph: ProblemGraph, end_effector_assignments: dict, options: dict = None -) -> list: - """ - - :param robot_graph: - :param end_effector_assignments: - :param options: - :return: - """ - G = generate_symbols(robot_graph) - # Anchor end effectors - for node_id in end_effector_assignments: - G.nodes[node_id]["pos"] = end_effector_assignments[node_id] - - constraints = [] - - # define structural constraints (this excludes end-effector distances) - for u, v, weight in G.edges.data("weight"): - if weight: - # Ugly, but we don't need these (TODO: Porta case check?) - if ( - u not in ("x", "y") - and v not in ("x", "y") - and set((u, v)) != set(("p0", "q0")) - and ("bounded" not in G[u][v] or (G[u][v]["bounded"])) - ): - p_u = ( - G.nodes[u]["pos"] - if u in end_effector_assignments or u in ("p0", "q0") - else G.nodes[u]["sym"] - ) - p_v = ( - G.nodes[v]["pos"] - if v in end_effector_assignments or v in ("p0", "q0") - else G.nodes[v]["sym"] - ) - equality = sp.Eq(norm_sq(p_u - p_v), weight ** 2) - if ( - type(equality) == sp.Equality - ): # This will be a boolean expression if both values are set - constraints.append(equality) - - return constraints - - -def angular_constraints( - robot_graph: ProblemGraph, - angular_limits: dict, - end_effector_assignments: dict, - angular_offsets: dict = None, - as_equality: bool = False, -) -> list: - if as_equality: - G = generate_symbols(robot_graph, angular_limits) - else: - G = generate_symbols(robot_graph) - constraints = [] - for node in angular_limits: - limit = angular_limits[node] - if node in end_effector_assignments: - x2 = end_effector_assignments[node] - elif node in robot_graph.base.nodes: - x2 = G.nodes[node]["pos"] - else: - x2 = G.nodes[node]["sym"] - parent = robot_graph.parent_node_id(node) - if parent in end_effector_assignments: - x1 = end_effector_assignments[parent] - elif parent in robot_graph.base.nodes: - x1 = G.nodes[parent]["pos"] - else: - x1 = G.nodes[parent]["sym"] - - # TODO: devise a better system for this? p0 is no longer the root of the DiGraph (it's not a tree anymore) - if parent == "p0": - x0 = np.zeros((robot_graph.dim,)) - x0[0] = 1.0 # np.array([1., 0., 0.]) - x0 = G.nodes[parent]["pos"] - x0 - else: - grandparent = robot_graph.parent_node_id(parent) - if grandparent in end_effector_assignments: - x0 = end_effector_assignments[grandparent] - elif grandparent in robot_graph.base.nodes: - x0 = G.nodes[grandparent]["pos"] - else: - x0 = G.nodes[grandparent]["sym"] - - l1 = 1.0 if parent == "p0" else G.edges[grandparent, parent]["weight"] - l2 = G.edges[parent, node]["weight"] - z0 = (x1 - x0) / l1 - z1 = (x2 - x1) / l2 - - if angular_offsets is not None: - z0 = apply_angular_offset_2d(angular_offsets[node], z0) - - if as_equality: - constraints.append( - sp.Eq( - norm_sq(z1 - z0) + G.angular_variables[node] ** 2, - 2.0 * (1.0 - np.cos(limit)), - ) - ) - else: - constraints.append(norm_sq(z1 - z0) <= 2.0 * (1.0 - np.cos(limit))) - - return constraints - - -def nearest_neighbour_cost( - robot_graph: ProblemGraph, - nearest_neighbour_points: dict, - nearest_angular_residuals: dict = None, -): - """ - Produce a nearest-neighbour cost function. - TODO: add assertions to help with checking this. For graph case, need a (6,) ndarray with p before q. - - :param robot_graph: Graph object describing our robot - :param nearest_neighbour_points: dictionary mapping node ids (strings) to ndarrays - :param nearest_angular_residuals: list of angular residuals (floats) - :return: symbolic expression of cost - """ - if nearest_angular_residuals is None: - G = generate_symbols(robot_graph) # augment the nodes with a symbol attribute - else: - G = generate_symbols(robot_graph, nearest_angular_residuals) - cost = 0.0 - - for node_id in nearest_neighbour_points: - cost += norm_sq(G.nodes[node_id]["sym"] - nearest_neighbour_points[node_id]) - - if nearest_angular_residuals is not None: - for node_id in nearest_angular_residuals: - cost += ( - G.angular_variables[node_id] - nearest_angular_residuals[node_id] - ) ** 2 - # if self.as_equality and nearest_angular_residuals is not None: - # for idx, node in enumerate(self.nodes[1:]): - # cost += (node.angle_variable - nearest_angular_residuals[idx]) ** 2 - return cost - - -def get_full_revolute_nearest_points_pose(graph, q): - full_points = [f"p{idx}" for idx in range(1, graph.robot.n)] + [ - f"q{idx}" for idx in range(1, graph.robot.n) - ] - return get_full_revolute_nearest_point(graph, q, full_points) - - -def get_full_revolute_nearest_point(graph, q, full_points=None): - nearest_points = {} - if full_points is None: - # Assume only x, y, p0, q0, and the final p (p{n+1}) is constrained - full_points = [f"p{idx}" for idx in range(1, graph.robot.n)] + [ - f"q{idx}" for idx in range(1, graph.robot.n + 1) - ] - G = graph.realization(q) - P = pos_from_graph(G) - for idx, key in enumerate(graph.node_ids): - if key in full_points: - nearest_points[key] = P[idx, :] - return nearest_points - - -if __name__ == "__main__": - - # n = 5 - - # a = np.ones(n) - # th = np.zeros(n) - - # params = {"a": a, "theta": th} - - # robot = Robot({**params, "dim": 2}) - # graph = PlanarRobotGraph(robot) - # constraints = constraints_from_graph(graph) - - n = 5 - # Generate random DH parameters - a = np.random.rand(n) - d = np.random.rand(n + 1) - al = np.random.rand(n) * pi / 2 - 2 * np.random.rand(n) * pi / 2 - th = 0 * np.ones(n + 1) - ub = np.concatenate([np.ones(n) * pi - 2 * pi * np.random.rand(n), [0]]) - lb = -ub - - params = {"a": a, "alpha": al, "d": d, "theta": th, "lb": lb, "ub": ub} - robot = Robot({**params, "dim": 3}) # instantiate robot - end_effector_assignments = {"p6": np.array([1.0, 1.0, 1.0])} - graph = ProblemGraph(robot) - constraints = constraints_from_graph(graph, end_effector_assignments) - print(graph.directed.edges()) - print(constraints) diff --git a/graphik/solvers/convex_iteration.py b/graphik/solvers/convex_iteration.py index 565bf22..f28bd7b 100644 --- a/graphik/solvers/convex_iteration.py +++ b/graphik/solvers/convex_iteration.py @@ -6,7 +6,6 @@ import cvxpy as cp import networkx as nx from timeit import default_timer -from progress.bar import ShadyBar as Bar from graphik.solvers.sdp_formulations import SdpSolverParams, solve_sdp from graphik.solvers.sdp_snl import ( @@ -17,27 +16,10 @@ extract_solution, chordal_sparsity_overlap_constraints ) -from graphik.solvers.constraints import get_full_revolute_nearest_point -from graphik.utils.roboturdf import load_ur10 from graphik.utils.constants import * from graphik.graphs import ProblemGraph -def random_psd_matrix(N: int, d: int = None, normalize: bool = True) -> np.ndarray: - """ - Return a random PSD matrix. - - :param N: size of the matrix (NxN) - :param d: rank of the matrix (full rank if not specified) - """ - d = N if d is None else d - P = np.random.rand(N, d) - W = P @ P.T - if normalize: - W = W / np.linalg.norm(W, ord="fro") - return W - - def solve_fantope_closed_form(G: np.ndarray, d:int): """ @@ -51,40 +33,6 @@ def solve_fantope_closed_form(G: np.ndarray, d:int): return U@U.T, default_timer() - start -def solve_fantope_sdp(G: np.ndarray, d: int): - n = G.shape[0] - Z, cons = fantope_constraints(n, d) - solve_fantope_iterate(G, Z, cons) - return Z.value - - -def fantope_constraints(n: int, rank: int, sparsity_pattern: set = None): - assert rank < n, "Needs a desired rank less than the problem dimension." - Z = cp.Variable((n, n), PSD=True) - constraints = [cp.trace(Z) == float(n - rank), np.eye(Z.shape[0]) - Z >> 0] - - if sparsity_pattern is not None: - for idx in range(n): - for jdx in range(idx, n): - if (idx, jdx) not in sparsity_pattern: - constraints += [Z[idx, jdx] == 0.] - if idx != jdx: - constraints += [Z[jdx, idx] == 0.] - - return Z, constraints - - -def solve_fantope_iterate( - G: np.ndarray, Z: cp.Variable, constraints: list, verbose=False, solver_params=None -): - # TODO: templatize for speed? Ask Filip about that feature, I forget what it's called - prob = cp.Problem(cp.Minimize(cp.trace(G @ Z)), constraints) - if solver_params is None: - solver_params = SdpSolverParams() - solve_sdp(prob, solver_params, verbose=verbose) - return prob - - def solve_fantope_sparse(sdp_variable_map: dict, d: int): C_mapping = {} t_fantope = 0.0 @@ -136,25 +84,6 @@ def sparse_eigenvalue_sum(sdp_variable_map: dict, d: int): return running_eigenvalue_sum -def get_sparsity_pattern(G, canonical_point_order: list) -> set: - sparsity_pattern = set() - G = G.copy() - # remove the edges that don't have distances defined - edges = [] - for u, v, data in G.edges(data=True): - if not data.get(DIST, False): - edges += [(u, v)] - G.remove_edges_from(edges) - undirected = G.to_undirected() - for idx, point_idx in enumerate(canonical_point_order): - for jdx in range(idx, len(canonical_point_order)): - point_jdx = canonical_point_order[jdx] - if undirected.has_edge(point_idx, point_jdx): - sparsity_pattern.add((idx, jdx)) - - return sparsity_pattern - - def convex_iterate_sdp_snl_graph( graph: ProblemGraph, anchors: dict = None, @@ -319,156 +248,3 @@ def solve_with_cidgik(graph: ProblemGraph, T_goal: np.ndarray) -> (dict, dict): else: return None, None - -if __name__ == "__main__": - - # TODO: use the graph convex iteration from above and deprecate the old one - # UR10 Test - robot, graph = load_ur10() - d = robot.dim - full_points = [f"p{idx}" for idx in range(0, graph.robot.n + 1)] + [ - f"q{idx}" for idx in range(0, graph.robot.n + 1) - ] - n_runs = 1000 - - # Store results - final_eigvalue_sum_list = [] - primal_sdp_runtime = [] - fantope_runtime = [] - feasibility_list = [] - - final_eigvalue_sum_list_sparse_naive = [] - primal_sdp_runtime_sparse_naive = [] - fantope_runtime_sparse_naive = [] - feasibility_list_sparse_naive = [] - final_eigvalue_sum_list_sparse_sdp = [] - primal_sdp_runtime_sparse_sdp = [] - fantope_runtime_sparse_sdp = [] - feasibility_list_sparse_sdp = [] - - bar = Bar("", max=n_runs, check_tty=False, hide_cursor=False) - for idx in range(n_runs): - # Generate a random feasible target - q = robot.random_configuration() - input_vals = get_full_revolute_nearest_point(graph, q, full_points) - - # End-effectors don't include the base pair at this step, that's inside of convex_iterate_sdp_snl_graph() - end_effectors = { - key: input_vals[key] for key in [f"p{robot.n}", f"q{robot.n}"] - } - - # Run 'dense' default solver - ( - _, - _, - _, - _, - eig_value_sum_vs_iterations, - _, - t_primal, - t_fantope, - feasible - ) = convex_iterate_sdp_snl_graph(graph, anchors=end_effectors, ranges=False, max_iters=10, - sparse=False, verbose=False, closed_form=True) - # solution = extract_solution(constraint_clique_dict, sdp_variable_map, d) - if feasible is FEASIBLE: - final_eigvalue_sum_list.append(eig_value_sum_vs_iterations[-1]) - else: - final_eigvalue_sum_list.append(np.nan) - primal_sdp_runtime.append(t_primal) - fantope_runtime.append(t_fantope) - feasibility_list.append(feasible) - - # Run 'naive' sparse solver (Wang and Yu, 2018) - ( - _, - _, - _, - _, - eig_value_sum_vs_iterations, - _, - t_primal, - t_fantope, - feasible - ) = convex_iterate_sdp_snl_graph(graph, anchors=end_effectors, ranges=False, max_iters=10, - sparse=True, verbose=False, closed_form=True) - primal_sdp_runtime_sparse_naive.append(t_primal) - fantope_runtime_sparse_naive.append(t_fantope) - if feasible is FEASIBLE: - final_eigvalue_sum_list_sparse_naive.append(eig_value_sum_vs_iterations[-1]) - else: - final_eigvalue_sum_list_sparse_naive.append(np.nan) - - # Run sparse solver where Fantope program has a sparsity pattern matched to the primal SDP's - ( - _, - _, - _, - _, - eig_value_sum_vs_iterations, - _, - t_primal, - t_fantope, - feasible - ) = convex_iterate_sdp_snl_graph(graph, anchors=end_effectors, ranges=False, max_iters=10, - sparse=True, verbose=False, closed_form=False) - primal_sdp_runtime_sparse_sdp.append(t_primal) - fantope_runtime_sparse_sdp.append(t_fantope) - if feasible is FEASIBLE: - final_eigvalue_sum_list_sparse_sdp.append(eig_value_sum_vs_iterations[-1]) - else: - final_eigvalue_sum_list_sparse_sdp.append(np.nan) - - bar.next() - bar.finish() - - # Visualize results - from matplotlib import pyplot as plt - from matplotlib import rc - rc("font", **{"family": "serif", "serif": ["Computer Modern"]}) - rc("text", usetex=True) - - # Plot parameters - n_bins = 100 - colors = ['r', 'g', 'b'] - # Runtime histograms - primal_sdp_runtime_full = np.vstack([primal_sdp_runtime, primal_sdp_runtime_sparse_naive, primal_sdp_runtime_sparse_sdp]).T - primal_sdp_runtime_min = np.min(primal_sdp_runtime_full) - primal_sdp_runtime_max = np.max(primal_sdp_runtime_full) - primal_sdp_runtime_bins = np.linspace(primal_sdp_runtime_min, primal_sdp_runtime_max, n_bins) - plt.figure() - plt.hist(primal_sdp_runtime_full, bins=primal_sdp_runtime_bins, color=colors) - plt.legend(('Dense', 'Sparse (Naive)', 'Sparse (SDP)')) - plt.title('Distribution of Primal SDP Runtimes') - plt.xlabel('Runtime (s)') - plt.ylabel('Count') - plt.grid() - plt.show() - - fantope_runtime_full = np.vstack([fantope_runtime, fantope_runtime_sparse_naive, fantope_runtime_sparse_sdp]).T - fantope_runtime_min = np.min(fantope_runtime_full) - fantope_runtime_max = np.max(fantope_runtime_full) - fantope_runtime_bins = np.linspace(fantope_runtime_min, fantope_runtime_max, n_bins) - plt.figure() - plt.hist(fantope_runtime_full, bins=fantope_runtime_bins, color=colors) - plt.legend(('Dense', 'Sparse (Naive)', 'Sparse (SDP)')) - plt.title('Distribution of Fantope Solver Runtimes') - plt.xlabel('Runtime (s)') - plt.ylabel('Count') - plt.grid() - plt.show() - - # Eigenvalue sum histograms - eigenvalue_sum_full = np.vstack([final_eigvalue_sum_list, final_eigvalue_sum_list_sparse_naive, final_eigvalue_sum_list_sparse_sdp]).T - eigenvalue_sum_min = max(np.min(eigenvalue_sum_full), 1e-12) - eigenvalue_sum_max = np.max(eigenvalue_sum_full) - # eigenvalue_sum_bins = np.log10(np.linspace(eigenvalue_sum_min, eigenvalue_sum_max, n_bins)) - eigenvalue_sum_bins = np.log10(np.logspace(np.log10(eigenvalue_sum_min), np.log10(eigenvalue_sum_max), n_bins)) - plt.figure() - plt.hist(np.log10(eigenvalue_sum_full), bins=eigenvalue_sum_bins, color=colors) - plt.legend(('Dense', 'Sparse (Naive)', 'Sparse (SDP)')) - plt.title('Distribution of Excess Rank') - plt.xlabel('log$_{10}$ Excess Rank') - plt.ylabel('Count') - plt.grid() - plt.show() diff --git a/graphik/solvers/joint_angle_solver.py b/graphik/solvers/joint_angle_solver.py index acbfe74..03b1cd3 100644 --- a/graphik/solvers/joint_angle_solver.py +++ b/graphik/solvers/joint_angle_solver.py @@ -1,17 +1,12 @@ -from graphik.utils.geometry import skew -import graphik import numpy as np import networkx as nx -from numpy.typing import ArrayLike -from typing import Dict, List, Any, Union +from typing import Dict, Any from scipy.optimize import minimize from pymlg.numpy import SE3, SO3, SE2, SO2 -from numpy import pi -from graphik.utils.roboturdf import RobotURDF from graphik.utils.constants import * from graphik.graphs.graph import ProblemGraph from graphik.utils.utils import list_to_variable_dict -from graphik.utils.roboturdf import load_kuka, load_ur10 + class LocalSolver: def __init__(self, robot_graph: ProblemGraph, params: Dict["str", Any]): @@ -35,53 +30,6 @@ def __init__(self, robot_graph: ProblemGraph, params: Dict["str", Any]): jac = self.gen_obstacle_constraint_gradient(pairs) self.g = [{"type": "ineq", "fun": fun, "jac": jac}] - def gen_objective_ee(self, point: str, T_goal: np.ndarray): - joints = self.k_map[point][1:] - n = len(joints) - - if self.dim == 3: - log = SE3.Log - inverse = SE3.inverse - else: - log = SE2.Log - inverse = SE2.inverse - - def objective(q): - q_dict = {joints[idx]: q[idx] for idx in range(n)} - T = self.robot.pose(q_dict, point) - e = log(inverse(T) @ T_goal).ravel() # body frame - return e.T @ e - - return objective - - - def gen_grad_ee(self, point: str, T_goal: np.ndarray): - joints = self.k_map[point][1:] - n = len(joints) - - if self.dim == 3: - log = SE3.Log - inverse = SE3.inverse - adjoint = SE3.adjoint - inv_left_jacobian = SE3.left_jacobian_inv - else: - log = SE2.Log - inverse = SE2.inverse - adjoint = SE2.adjoint - inv_left_jacobian = SE2.left_jacobian_inv - - def gradient(q): - q_dict = {joints[idx]: q[idx] for idx in range(n)} - T = self.robot.pose(q_dict, point) - T_inv = inverse(T) - J = self.robot.jacobian(q_dict, [point]) - e = log(T_inv @ T_goal).ravel() # body frame - J_e = inv_left_jacobian(e) - J[point] = J_e @ adjoint(T_inv) @ J[point] - jac = -2 * J[point].T @ e - return jac - return gradient - def gen_obstacle_constraints(self, pairs: list): def obstacle_constraint(q): q_dict = list_to_variable_dict(q) @@ -185,66 +133,3 @@ def cost_and_grad(q): options={"ftol": 1e-7}, ) return res - - -def main(): - # - # Define the problem - # - - # robot, graph = load_ur10() - # scale = 0.75 - # radius = 0.4 - # obstacles = [ - # (scale * np.asarray([1, 1, 0]), radius), - # (scale * np.asarray([1, -1, 0]), radius), - # (scale * np.asarray([-1, 1, 0]), radius), - # (scale * np.asarray([-1, -1, 0]), radius), - # (scale * np.asarray([0, 0, 1]), radius), - # (scale * np.asarray([0, 0, -1]), radius), - # ] - - # for idx, obs in enumerate(obstacles): - # graph.add_spherical_obstacle(f"o{idx}", obs[0], obs[1]) - - # q_goal = robot.random_configuration() - # T_goal = robot.pose(q_goal, f"p{robot.n}") - # goals = {f"p{robot.n}": T_goal} - - # x0 = [0,0,0,0,0,0] - # problem = LocalSolver(graph,{}) - # sol = problem.solve(goals, robot.random_configuration()) - # # sol = problem.solve(goals, list_to_variable_dict(x0)) - # print(sol) - - from graphik.robots import Robot - from graphik.graphs import ProblemGraph - n = 10 - - a = list_to_variable_dict(np.ones(n)) - th = list_to_variable_dict(np.zeros(n)) - lim_u = list_to_variable_dict(np.pi * np.ones(n)) - lim_l = list_to_variable_dict(-np.pi * np.ones(n)) - params = { - "link_lengths": a, - "theta": th, - "ub": lim_u, - "lb": lim_l, - "num_joints": n - } - - robot = Robot({**params, "dim": 2}) - graph = ProblemGraph(robot) - q_goal = robot.random_configuration() - T_goal = robot.pose(q_goal, f"p{robot.n}") - goals = {f"p{robot.n}": T_goal} - problem = LocalSolver(graph,{}) - sol = problem.solve(goals, robot.random_configuration()) - q_sol = list_to_variable_dict(sol.x) - print(T_goal) - print(robot.pose(q_sol, robot.end_effectors[0])) - -if __name__ == '__main__': - - # np.random.seed(24) # TODO: this seems to have a significant effect on performance - main() diff --git a/graphik/solvers/sdp_snl.py b/graphik/solvers/sdp_snl.py index 3374415..e7b6126 100644 --- a/graphik/solvers/sdp_snl.py +++ b/graphik/solvers/sdp_snl.py @@ -5,15 +5,28 @@ import networkx as nx import cvxpy as cp -from graphik.utils.roboturdf import load_ur10, load_truncated_ur10 from graphik.utils.constants import * -from graphik.utils.chordal import complete_to_chordal_graph +from graphik.utils.dgp import pos_from_graph from graphik.robots import Robot from graphik.graphs import ProblemGraph -from graphik.solvers.constraints import get_full_revolute_nearest_point from graphik.solvers.sdp_formulations import SdpSolverParams, solve_sdp +def get_full_revolute_nearest_point(graph, q, full_points=None): + nearest_points = {} + if full_points is None: + # Assume only x, y, p0, q0, and the final p (p{n+1}) is constrained + full_points = [f"p{idx}" for idx in range(1, graph.robot.n)] + [ + f"q{idx}" for idx in range(1, graph.robot.n + 1) + ] + G = graph.realization(q) + P = pos_from_graph(G) + for idx, key in enumerate(graph.node_ids): + if key in full_points: + nearest_points[key] = P[idx, :] + return nearest_points + + _INFEASIBLE_STATUSES = frozenset({cp.INFEASIBLE, cp.INFEASIBLE_INACCURATE}) @@ -287,8 +300,7 @@ def distance_constraints_graph( anchors = anchors if anchors is not None else {} G = G.copy() if angle_limits: - # Remove the edges that don't have a - typ = nx.get_edge_attributes(G, name=BOUNDED) + # Remove the edges with neither a distance nor an active bound edges = [] for u, v, data in G.edges(data=True): bounds = data.get(BOUNDED, []) @@ -307,7 +319,7 @@ def distance_constraints_graph( undirected = G.to_undirected() if not nx.is_chordal(undirected): - undirected, _ = complete_to_chordal_graph(undirected) + undirected, _ = nx.complete_to_chordal_graph(undirected) equality_cliques = nx.chordal_graph_cliques( undirected ) # Returns maximal cliques (in spite of name) @@ -975,171 +987,3 @@ def solve_linear_cost_sdp( solution = extract_solution(constraint_clique_dict, sdp_variable_map, robot.dim) return solution, prob, sdp_variable_map, canonical_point_order - - -def sym_vec(A: np.ndarray) -> np.ndarray: - vec = [] - for idx in range(A.shape[0]): - vec_idx = A[idx, idx:] - vec_idx[1:] = 2.0 * vec_idx[1:] - vec.append(vec_idx) - return np.hstack(vec) - - -def constraints_list_to_matrix(A_list: list): - a_list = [sym_vec(A) for A in A_list] - A_symmetrized = np.vstack(a_list) - # TODO: does this capture a meaningful rank of the linear operator? I believe so! - return A_symmetrized - - -if __name__ == "__main__": - # Simple examples - sparse = False # Whether to exploit chordal sparsity in the SDP formulation - ee_cost = False # Whether to treat the end-effectors as variables with targets in the cost. - # If False, end-effectors are NOT variables (they are baked in to constraints as parameters) - conic_solver = "MOSEK" # One of "MOSEK", "CVXOPT" for now - solver_params = ( - SdpSolverParams() - ) # Use MOSEK settings that worked well for us before - # Truncated UR10 (only the first n joints) - n = 6 - if n == 6: - robot, graph = load_ur10() - else: - robot, graph = load_truncated_ur10(n) - - # Generate a random feasible target - q = robot.random_configuration() - - # Extract the positions of the points - full_points = [f"p{idx}" for idx in range(0, graph.robot.n + 1)] + [ - f"q{idx}" for idx in range(0, graph.robot.n + 1) - ] - input_vals = get_full_revolute_nearest_point(graph, q, full_points) - - # End-effectors are 'generalized' to include the base pair ('p0', 'q0') - end_effectors = { - key: input_vals[key] for key in ["p0", "q0", f"p{robot.n}", f"q{robot.n}"] - } - - # Form the constraints - constraint_clique_dict = distance_constraints(graph, end_effectors, sparse, ee_cost) - # A, b, mapping, _ = list(constraint_clique_dict.values())[0] - - # Different cost function options here - cost function is controlled by a dictionary mapping some subset of the keys - # of input_vals (all the points' positions) to some nearest point. - - # Nuclear norm - the nearest points are all zero - nearest_points_nuclear = { - key: np.zeros(robot.dim) - for key in input_vals - if key not in ["p0", "q0", f"p{robot.n}", f"q{robot.n}"] - } - ( - sdp_variable_map_nuclear, - sdp_constraints_map_nuclear, - sdp_cost_map_nuclear, - ) = constraints_and_nearest_points_to_sdp_vars( - constraint_clique_dict, nearest_points_nuclear, robot.dim - ) - prob_nuclear = form_sdp_problem( - constraint_clique_dict, - sdp_variable_map_nuclear, - sdp_constraints_map_nuclear, - sdp_cost_map_nuclear, - robot.dim, - ) - prob_nuclear.solve( - verbose=True, solver=conic_solver, mosek_params=solver_params.mosek_params - ) - # Analysis below assumes dense (sparse = False) case - Z_nuclear = list(sdp_variable_map_nuclear.values())[0].value - _, s_nuclear, _ = np.linalg.svd(Z_nuclear) - solution_rank_nuclear = np.linalg.matrix_rank(Z_nuclear, tol=1e-6, hermitian=True) - solution_nuclear = extract_solution( - constraint_clique_dict, sdp_variable_map_nuclear, robot.dim - ) - - # Feasibility (no nearest points means cost function is 0) - no_nearest_points = {} - ( - sdp_variable_map, - sdp_constraints_map, - sdp_cost_map, - ) = constraints_and_nearest_points_to_sdp_vars( - constraint_clique_dict, no_nearest_points, robot.dim - ) - prob_feas = form_sdp_problem( - constraint_clique_dict, - sdp_variable_map, - sdp_constraints_map, - sdp_cost_map, - robot.dim, - ) - prob_feas.solve( - verbose=True, solver=conic_solver, mosek_params=solver_params.mosek_params - ) - # Analysis below assumes dense (sparse = False) case - Z = list(sdp_variable_map.values())[0].value - _, s, _ = np.linalg.svd(Z) - solution_rank = np.linalg.matrix_rank(Z, tol=1e-6, hermitian=True) - solution = extract_solution(constraint_clique_dict, sdp_variable_map, robot.dim) - - # Exact nearest point - use the true value from q (don't perturb) - exact_nearest_points = { - key: input_vals[key] - for key in input_vals - if key not in ["p0", "q0", f"p{robot.n}", f"q{robot.n}"] - } - ( - sdp_variable_map_exact, - sdp_constraints_map_exact, - sdp_cost_map_exact, - ) = constraints_and_nearest_points_to_sdp_vars( - constraint_clique_dict, exact_nearest_points, robot.dim - ) - prob_exact = form_sdp_problem( - constraint_clique_dict, - sdp_variable_map_exact, - sdp_constraints_map_exact, - sdp_cost_map_exact, - robot.dim, - ) - prob_exact.solve( - verbose=True, solver=conic_solver, mosek_params=solver_params.mosek_params - ) - Z_exact = list(sdp_variable_map_exact.values())[0].value - _, s_exact, _ = np.linalg.svd(Z_exact) - solution_rank_exact = np.linalg.matrix_rank(Z_exact, tol=1e-6, hermitian=True) - solution_exact = extract_solution( - constraint_clique_dict, sdp_variable_map_exact, robot.dim - ) - - total_error_exact = 0.0 - total_error_nuclear = 0.0 - total_error_feas = 0.0 - for key in solution_exact: - print(f"{key}") - print(f"True value: {input_vals[key]}") - print(f"Nearest point value: {solution_exact[key]}") - print(f"Nuclear norm value: {solution_nuclear[key]}") - print(f"Feasibility value: {solution[key]}") - print( - "------------------------------------------------------------------------" - ) - total_error_exact += np.linalg.norm(input_vals[key] - solution_exact[key]) - total_error_nuclear += np.linalg.norm(input_vals[key] - solution_nuclear[key]) - total_error_feas += np.linalg.norm(input_vals[key] - solution[key]) - - # Compare the ranks - print(f"Feasibility formulation rank: {solution_rank}") - print(f"Nuclear norm rank: {solution_rank_nuclear}") - print(f"Exact nearest point rank: {solution_rank_exact}") - - # Print the total L2 error - print(f"Total error feas: {total_error_feas}") - print(f"Total error nuclear: {total_error_nuclear}") - print(f"Total error exact nearest: {total_error_exact}") - - # TODO: check constraint violations! See Filip's code (or just plug this in there) diff --git a/graphik/utils/chordal.py b/graphik/utils/chordal.py deleted file mode 100644 index d24449a..0000000 --- a/graphik/utils/chordal.py +++ /dev/null @@ -1,66 +0,0 @@ -import networkx as nx - - -def complete_to_chordal_graph(G): - """Return a copy of G completed to a chordal graph - Adds edges to a copy of G to create a chordal graph. A graph G=(V,E) is - called chordal if for each cycle with length bigger than 3, there exist - two non-adjacent nodes connected by an edge (called a chord). - Parameters - ---------- - G : NetworkX graph - Undirected graph - Returns - ------- - H : NetworkX graph - The chordal enhancement of G - alpha : Dictionary - The elimination ordering of nodes of G - Notes - ----- - There are different approaches to calculate the chordal - enhancement of a graph. The algorithm used here is called - MCS-M and gives at least minimal (local) triangulation of graph. Note - that this triangulation is not necessarily a global minimum. - https://en.wikipedia.org/wiki/Chordal_graph - References - ---------- - .. [1] Berry, Anne & Blair, Jean & Heggernes, Pinar & Peyton, Barry. (2004) - Maximum Cardinality Search for Computing Minimal Triangulations of - Graphs. Algorithmica. 39. 287-298. 10.1007/s00453-004-1084-3. - Examples - -------- - >> from networkx.algorithms.chordal import complete_to_chordal_graph - >> G = nx.wheel_graph(10) - >> H, alpha = complete_to_chordal_graph(G) - """ - H = G.copy() - alpha = {node: 0 for node in H} - if nx.is_chordal(H): - return H, alpha - chords = set() - weight = {node: 0 for node in H.nodes()} - unnumbered_nodes = list(H.nodes()) - for i in range(len(H.nodes()), 0, -1): - # get the node in unnumbered_nodes with the maximum weight - z = max(unnumbered_nodes, key=lambda node: weight[node]) - unnumbered_nodes.remove(z) - alpha[z] = i - update_nodes = [] - for y in unnumbered_nodes: - if G.has_edge(y, z): - update_nodes.append(y) - else: - # y_weight will be bigger than node weights between y and z - y_weight = weight[y] - lower_nodes = [ - node for node in unnumbered_nodes if weight[node] < y_weight - ] - if nx.has_path(H.subgraph(lower_nodes + [z, y]), y, z): - update_nodes.append(y) - chords.add((z, y)) - # during calculation of paths the weights should not be updated - for node in update_nodes: - weight[node] += 1 - H.add_edges_from(chords) - return H, alpha \ No newline at end of file diff --git a/graphik/utils/dgp.py b/graphik/utils/dgp.py index f134b0d..f058a44 100644 --- a/graphik/utils/dgp.py +++ b/graphik/utils/dgp.py @@ -4,24 +4,6 @@ import networkx as nx from numpy.typing import NDArray from graphik.utils.constants import * -from graphik.utils.geometry import best_fit_transform - - -def orthogonal_procrustes(G1: nx.DiGraph, G2: nx.DiGraph) -> nx.DiGraph: - """ - Aligns two point clouds represented by graphs by aligning nodes with - matching labels and returns a graph representing the aligned points. - - Keyword Arguments: - G1: nx.DiGraph -- Graph representing the point set to align with. - G2: nx.DiGraph -- Graph representing the point set to be aligned. - """ - Y = pos_from_graph(G1, node_ids=list(G1)) - X = pos_from_graph(G2, node_ids=list(G1)) - R, t = best_fit_transform(X, Y) - X = pos_from_graph(G2) - P_e = (R @ X.T + t.reshape(len(t), 1)).T - return graph_from_pos(P_e, list(G2)) # not really order-dependent def gram_from_distance_matrix(D: NDArray) -> NDArray: @@ -78,15 +60,6 @@ def adjacency_matrix_from_graph( G.edge_subgraph(selected_edges), weight="", nodelist=nodelist ) -def incidence_matrix_from_adjacency(A): - edges = np.nonzero(np.triu(A, 1)) # Get indices of upper triangle non-zeros - num_edges = len(edges[0]) - B = np.zeros((num_edges, A.shape[0])) - rows = np.arange(num_edges) - B[rows, edges[0]] = 1 - B[rows, edges[1]] = -1 - return B - def pos_from_graph(G: nx.DiGraph, node_ids=None) -> NDArray: """ Returns an n x m matrix of node positions from a given graph, @@ -162,14 +135,6 @@ def graph_complete_edges( return G -def factor(A: NDArray): - (evals, evecs) = np.linalg.eigh(A) - evals[evals < 0] = 0 # closest SDP matrix - X = evecs * np.sqrt(evals) - return np.fliplr(X) - - -## perform classic Multidimensional scaling def MDS(B: NDArray, eps: float = 1e-5): evals, evecs = np.linalg.eigh(B) # ascending evals = evals[::-1] @@ -191,10 +156,6 @@ def linear_projection(P: NDArray, F: NDArray, dim): return P @ np.fliplr(eigvec)[:, :dim] -def sample_matrix(lower_limit, upper_limit): - m, n = lower_limit.shape - return lower_limit + np.random.rand(m, n) * (upper_limit - lower_limit) - def bound_smoothing(G: nx.DiGraph) -> tuple: """ Given a graph with some edges containing upper and lower bounds on distance, diff --git a/graphik/utils/kinematics.py b/graphik/utils/kinematics.py index 95515e6..83adf29 100644 --- a/graphik/utils/kinematics.py +++ b/graphik/utils/kinematics.py @@ -38,13 +38,6 @@ def dh_to_se2(a: float, theta: float) -> np.ndarray: return T -def fk_2d(a: list, theta: list, q: list) -> np.ndarray: - """Forward kinematics from arrays of link lengths and angles. Returns 3x3 SE(2).""" - if len(a) > 1: - return dh_to_se2(a[0], theta[0] + q[0]) @ fk_2d(a[1:], theta[1:], q[1:]) - return dh_to_se2(a[0], theta[0] + q[0]) - - def fk_tree_2d(a: list, theta: list, q: list, path_indices: list) -> np.ndarray: """Forward kinematics along a path through a tree. Returns 3x3 SE(2).""" idx = path_indices[0] diff --git a/graphik/utils/operators.py b/graphik/utils/operators.py deleted file mode 100644 index 880c630..0000000 --- a/graphik/utils/operators.py +++ /dev/null @@ -1,81 +0,0 @@ -import numpy as np -from scipy.sparse import csc_array -from graphik.utils.dgp import incidence_matrix_from_adjacency - -def res_left_op_batched(omega, dim, vectorized=False, sparse=False): - # B is the incidence matrix; implements p_i - p_j for p -> (N, dim) - B = incidence_matrix_from_adjacency(omega) # (d N) - num_points = omega.shape[0] - num_dist = B.shape[0] - - if vectorized: - res_vec_all = [np.kron(B[i], np.eye(dim)).reshape(dim, num_points, dim) for i in range(num_dist)] - operator = np.stack(res_vec_all) # (d dim N dim) - else: - operator = B - - if sparse: - operator = csc_array(operator) - - return operator - -def sum_square_op(omega, dim, vectorized=False, sparse=False, reduced=False): - B = incidence_matrix_from_adjacency(omega) - - if vectorized: - B_vec = np.kron(B, np.eye(dim)) - operator = B_vec.T.dot(B_vec) - else: - operator = B.T.dot(B) - - if sparse: - operator = csc_array(operator) - - return operator - -def sum_square_op_batched(omega, dim, vectorized=False, sparse=False, flat=False): - B = incidence_matrix_from_adjacency(omega) # (d N) - num_points = omega.shape[0] - num_dist = B.shape[0] - - if vectorized: - # d[i] = (p[i] - p[j])**2 = p @ A[i] @ p, batched over i - res_sq_vec_all = [np.kron(np.outer(B[i], B[i]), np.eye(dim)) for i in range(num_dist)] - operator = np.stack(res_sq_vec_all) # (d N*dim N*dim) - - if flat: - operator = np.ascontiguousarray(operator.reshape(num_dist * num_points * dim, -1)) - else: - res_sq_all = [np.outer(B[i], B[i]) for i in range(num_dist)] - operator = np.stack(res_sq_all) - - if flat: - operator = np.ascontiguousarray(operator.reshape(num_dist * num_points, -1)) - - if sparse: - operator = csc_array(operator) - - return operator - -def diff_sum_square_op_batched(omega, dim, vectorized=False, sparse=False, flat=False): - B = incidence_matrix_from_adjacency(omega) - num_points = omega.shape[0] - num_dist = B.shape[0] - - if vectorized: - res_sq_vec_all = [np.kron(np.outer(B[i], B[i]), np.eye(dim)) for i in range(num_dist)] - operator = np.stack([M + M.T for M in res_sq_vec_all]) - - if flat: - operator = np.ascontiguousarray(operator.swapaxes(0, 1).reshape(num_points * dim, -1).T) - else: - res_sq_all = [np.outer(B.T[i], B.T[i]) for i in range(num_dist)] - operator = np.stack([M + M.T for M in res_sq_all]) - - if flat: - operator = np.ascontiguousarray(operator.swapaxes(0, 1).reshape(num_points, -1).T) - - if sparse: - operator = csc_array(operator) - - return operator diff --git a/graphik/utils/robot_visualization.py b/graphik/utils/robot_visualization.py index 093dfce..836ab47 100644 --- a/graphik/utils/robot_visualization.py +++ b/graphik/utils/robot_visualization.py @@ -10,11 +10,12 @@ # mpl.rcParams['legend.fontsize'] = 10 -# Use teX to render text -rc("font", **{"family": "serif", "serif": ["Computer Modern"]}) -rc("text", usetex=True) -# Change the sample figure directory to the local assets/ folder in this repo -# rc('examples', directory='assets') +# LaTeX text rendering is opt-in: enabling it at import time breaks the +# module for users without a TeX install. Call this before plotting if +# you want publication-style labels. +def use_latex_text(): + rc("font", **{"family": "serif", "serif": ["Computer Modern"]}) + rc("text", usetex=True) def plot_heatmap( diff --git a/graphik/utils/utils.py b/graphik/utils/utils.py index e3abd6b..4b22b2e 100644 --- a/graphik/utils/utils.py +++ b/graphik/utils/utils.py @@ -1,12 +1,9 @@ from typing import Any, Callable, TypeVar import numpy as np -import scipy as sp import networkx as nx from numpy import pi -dZ = 1 - T = TypeVar("T") R = TypeVar("R") @@ -29,18 +26,6 @@ def wrapped(arg: T) -> R: return wrapped -def perpendicular_vector(v): - """ Returns arbitrary perpendicular vector """ - - if np.isclose(v[1], 0.) and np.isclose(v[2], 0.): - - if np.isclose(v[0], 0.): - raise ValueError('Received zero vector.') - else: - return np.cross(v, np.array([0,1,0])) - - return np.cross(v, np.array([1,0,0])) - def level2_descendants(G: nx.DiGraph, node_id): successors = G.successors(node_id) @@ -51,11 +36,6 @@ def level2_descendants(G: nx.DiGraph, node_id): return flatten(desc) -def norm_sq(A: np.ndarray) -> np.ndarray: - """Returns the squared L2-norm of a symbolic vector""" - return A.transpose().dot(A) - - def wraptopi(e): return np.mod(e + pi, 2 * pi) - pi @@ -71,125 +51,6 @@ def list_to_variable_dict(l: list, label="p", index_start=1): return var_dict -def list_to_variable_dict_spherical(l: list, label="p", index_start=1, in_pairs=False): - var_dict = {} - if in_pairs: - for idx, val in enumerate(l): - if idx % 2 == 0: - var_dict[label + str(index_start + idx // 2)] = [val] - else: - var_dict[label + str(index_start + (idx - 1) // 2)].append(val) - else: - for idx, val in enumerate(l): - var_dict[ - label + str(index_start + idx // 2) + "_" + str(index_start + idx % 2) - ] = val - return var_dict - - -def variable_dict_to_list(d: dict, order: list = None) -> list: - if order is None: - return [d[item] for item in d] - else: - return [d[item] for item in order] - - - -def make_save_string(save_properties: list) -> str: - """ - - :param save_properties: list of tuples containing (property, val) - :return: save string with underscores delimiting values and properties - """ - save_string = "" - return save_string.join([p + "_" + str(v) + "_" for p, v in save_properties]) - - -def spherical_angle_bounds_to_revolute(ub_spherical, lb_spherical): - ub = {} - lb = {} - count_angle_bounds = 1 - for key in ub_spherical: # Assumes the keys are numerically sorted - ub[f"p{count_angle_bounds}"] = np.pi - lb[f"p{count_angle_bounds}"] = -np.pi - count_angle_bounds += 1 - ub[f"p{count_angle_bounds}"] = ub_spherical[key] - lb[f"p{count_angle_bounds}"] = lb_spherical[key] - count_angle_bounds += 1 - - return ub, lb - - -def safe_arccos(t): - t_in = min(max(-1, t), 1) - return np.arccos(t_in) - - -def bernoulli_confidence_normal_approximation(n, n_success, confidence=0.95): - """ - - :param n: - :param n_success: - :param confidence: - :return: - """ - alpha = 1.0 - confidence - z = sp.special.ndtri(1.0 - alpha / 2.0) - p_hat = n_success / n - rad = z * np.sqrt((p_hat * (1 - p_hat)) / n) - return p_hat, rad - - -def wilson(n, n_success, alpha=0.95): - p = n_success / n - z = sp.special.ndtri(1.0 - alpha / 2.0) - denominator = 1 + z ** 2 / n - centre_adjusted_probability = p + z * z / (2 * n) - adjusted_standard_deviation = np.sqrt((p * (1 - p) + z * z / (4 * n)) / n) - lower_bound = ( - centre_adjusted_probability - z * adjusted_standard_deviation - ) / denominator - upper_bound = ( - centre_adjusted_probability + z * adjusted_standard_deviation - ) / denominator - return (lower_bound, upper_bound) - - -def bernoulli_confidence_jeffreys(n, n_success, confidence=0.95): - alpha_low = (1.0 - confidence) / 2.0 - alpha_high = confidence + alpha_low - a = n_success + 0.5 - b = n - n_success + 0.5 - low_end = 0.0 if n_success == 0 else sp.special.btdtri(a, b, alpha_low) - high_end = 1.0 if n_success == n else sp.special.btdtri(a, b, alpha_high) - p_hat = (low_end + high_end) / 2.0 - rad = (high_end - low_end) / 2.0 - - return p_hat, rad - - -def measure_perturbation(points: dict, points_perturbed: dict) -> (float, float): - squared_sum = 0. - max_perturb = -np.inf - for key in points: - p = points[key] - p_perturbed = points_perturbed[key] - max_perturb = max(max_perturb, max(np.abs(p-p_perturbed))) - squared_sum += np.linalg.norm(p - p_perturbed)**2 - - return np.sqrt(squared_sum), max_perturb - - -def constraint_violations(constraints, solution_dict): - return [ - ( - con.args[1].subs(solution_dict) - con.args[0].subs(solution_dict), - con.is_Equality, - ) - for con in constraints - ] - - def normalize(v): norm = np.linalg.norm(v) if norm == 0: @@ -210,16 +71,3 @@ def table_environment(height=0.9, width=0.8, n_height=9, n_width=8, obs_inflatio leg4 = [(np.asarray([ width/2-radius, width/2-radius, (2*i+1)*radius]), obs_inflation*radius) for i in range(0, n_height)] return tabletop_obs + leg1 + leg2 + leg3 + leg4 - - -if __name__ == "__main__": - - # print("Bernoulli: ") - # print(bernoulli_confidence_normal_approximation(100, 100)) - # print("Jeffreys: ") - # print(bernoulli_confidence_jeffreys(100, 100)) - - print("Bernoulli: ") - print(bernoulli_confidence_normal_approximation(100, 99)) - print("Jeffreys: ") - print(bernoulli_confidence_jeffreys(100, 99)) diff --git a/setup.py b/setup.py index e840b13..bdbaebf 100644 --- a/setup.py +++ b/setup.py @@ -16,11 +16,9 @@ install_requires=[ "numpy >= 2", "scipy >= 1.17", - "sympy >= 1.14", "matplotlib >= 3.1", "cvxpy >= 1.8", "networkx >= 3.6", - "progress", "yourdfpy >= 0.0.60", "trimesh >= 4.6.1", "numba >= 0.65", diff --git a/tests/test_pymlg_conventions.py b/tests/test_pymlg_conventions.py index a2130bb..77107b3 100644 --- a/tests/test_pymlg_conventions.py +++ b/tests/test_pymlg_conventions.py @@ -113,15 +113,9 @@ def test_se3_left_jacobian_inv_at_origin_is_identity(): def test_se2_left_jacobian_inv_exists(): - """LocalSolver.gen_grad_ee uses SE2.left_jacobian_inv when dim==2. - If PyMLG doesn't expose it, Task 9 must use lambda x: np.eye(3) fallback - (matching gen_cost_and_grad_ee's existing pattern in the original code). - """ - assert hasattr(SE2, "left_jacobian_inv"), ( - "PyMLG SE2 has no left_jacobian_inv — joint_angle_solver Task 9 " - "must use lambda x: np.eye(3) fallback for the dim==2 branch of " - "gen_grad_ee" - ) + """LocalSolver.gen_cost_and_grad_ee uses SE2.left_jacobian_inv when + dim==2; its 2D gradient is wrong without it.""" + assert hasattr(SE2, "left_jacobian_inv") J = SE2.left_jacobian_inv(np.zeros(3)) np.testing.assert_allclose(J, np.eye(3), atol=1e-12) diff --git a/tests/test_sdp_snl.py b/tests/test_sdp_snl.py index 80261b0..d67897b 100644 --- a/tests/test_sdp_snl.py +++ b/tests/test_sdp_snl.py @@ -5,7 +5,7 @@ from graphik.graphs import ProblemGraph from graphik.robots import Robot -from graphik.solvers.constraints import get_full_revolute_nearest_point +from graphik.solvers.sdp_snl import get_full_revolute_nearest_point from graphik.solvers.sdp_snl import ( distance_constraints, distance_constraints_graph, diff --git a/tests/test_sdp_snl_inequality_constraints.py b/tests/test_sdp_snl_inequality_constraints.py index 008bfb1..ced547b 100644 --- a/tests/test_sdp_snl_inequality_constraints.py +++ b/tests/test_sdp_snl_inequality_constraints.py @@ -9,9 +9,9 @@ from graphik.solvers.sdp_snl import ( distance_constraints_graph, distance_range_constraints, + get_full_revolute_nearest_point, solve_nearest_point_sdp, ) -from graphik.solvers.constraints import get_full_revolute_nearest_point from graphik.utils.constants import POS from graphik.utils.roboturdf import load_ur10, load_truncated_ur10 From eec21edefd8dadf7e9f45218200ab2bf8e8e1cca Mon Sep 17 00:00:00 2001 From: Filip Maric Date: Wed, 10 Jun 2026 07:03:05 +0200 Subject: [PATCH 3/9] ci: add GitHub Actions workflow (pytest on 3.11/3.12) The full suite now runs without Mosek (free-solver SDP tests with loose tolerances) and without the optional numba AOT extension, so a plain 'pip install -e .[dev]' is all CI needs. Co-Authored-By: Claude Fable 5 --- .github/workflows/ci.yml | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) create mode 100644 .github/workflows/ci.yml diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..b72410b --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,29 @@ +name: CI + +on: + push: + branches: [main] + pull_request: + +jobs: + test: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.11", "3.12"] + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: pip + + - name: Install package + run: | + python -m pip install --upgrade pip + pip install -e .[dev] + + - name: Run tests + run: pytest -q From 7464c3a9e8a08ac820ff20f0475d30848db8b36d Mon Sep 17 00:00:00 2001 From: Filip Maric Date: Wed, 10 Jun 2026 08:07:44 +0200 Subject: [PATCH 4/9] perf(solvers): Gauss-Newton preconditioner for RTR's truncated CG The EDM loss's Gauss-Newton Hessian has graph-Laplacian block structure (rank-1 edge blocks 8w_ij uu^T, u = y_i - y_j), which is both why tCG needed ~40 HVPs per outer iteration and why a near-exact preconditioner is cheap: at IK sizes the N*d x N*d GN matrix factors in microseconds. make_gn_preconditioner inverts it with relative eigenvalue flooring (default 1e-3, swept empirically; the GN matrix is PSD with a translation nullspace) and one eigh per outer iteration, memoized per Y; the output is projected back to the horizontal space. A block-Jacobi variant is included for comparison. UR10 with joint limits, 60 random goals over 3 seeds (see experiments/rtr_preconditioner_study.py): precon=None: 54/60 success, ~40 HVPs/outer, median 77-147 ms/solve precon='gn': 57/60 success, ~4 HVPs/outer, median 29-66 ms/solve Opt-in via solve(..., precon='gn'|'jacobi'|callable) and solve_with_riemannian(..., precon=...): preconditioning changes the optimization trajectory, so flipping the default belongs with a baseline regeneration. Co-Authored-By: Claude Fable 5 --- experiments/rtr_preconditioner_study.py | 97 +++++++++++++++++++ graphik/solvers/riemannian_solver.py | 119 +++++++++++++++++++++++- tests/test_preconditioner.py | 84 +++++++++++++++++ 3 files changed, 298 insertions(+), 2 deletions(-) create mode 100644 experiments/rtr_preconditioner_study.py create mode 100644 tests/test_preconditioner.py diff --git a/experiments/rtr_preconditioner_study.py b/experiments/rtr_preconditioner_study.py new file mode 100644 index 0000000..c31fa03 --- /dev/null +++ b/experiments/rtr_preconditioner_study.py @@ -0,0 +1,97 @@ +"""Preconditioner study for the Riemannian solver's truncated CG. + +Compares no preconditioner vs block-Jacobi vs the full Gauss-Newton +Laplacian on UR10 with joint limits (the production configuration). +Counts HVPs (the profile's hottest operation), outer iterations, wall +time, and solution quality (final pose error after angle decoding). +""" +import time + +import numpy as np + +from graphik.solvers import rtr +from graphik.solvers.riemannian_solver import RiemannianSolver +from graphik.utils.dgp import ( + adjacency_matrix_from_graph, + bound_smoothing, + distance_matrix_from_graph, + graph_from_pos, +) +from graphik.utils.roboturdf import load_ur10 + + +def run_case(graph, T_goal, precon): + solver = RiemannianSolver(graph) + G = graph.from_pose(T_goal) + D_goal = distance_matrix_from_graph(G) + omega = adjacency_matrix_from_graph(G) + lb, ub = bound_smoothing(G) + + n_hvp = [0] + orig_for_riemannian = __import__( + "graphik.solvers.loss", fromlist=["for_riemannian"] + ).for_riemannian + + import graphik.solvers.loss as loss_mod + + def counting_for_riemannian(*args, **kwargs): + cost, egrad, ehvp = orig_for_riemannian(*args, **kwargs) + + def counted_ehvp(Y, Z): + n_hvp[0] += 1 + return ehvp(Y, Z) + + return cost, egrad, counted_ehvp + + loss_mod.for_riemannian, saved = counting_for_riemannian, loss_mod.for_riemannian + try: + t0 = time.perf_counter() + out = solver.solve( + D_goal, omega, use_limits=True, bounds=(lb, ub), precon=precon + ) + wall = time.perf_counter() - t0 + finally: + loss_mod.for_riemannian = saved + + G_sol = graph_from_pos(out["x"], graph.node_ids) + q_sol = graph.joint_variables(G_sol, {f"p{graph.robot.n}": T_goal}) + T_sol = graph.robot.pose(q_sol, f"p{graph.robot.n}") + trans_err = np.linalg.norm(T_sol[:3, 3] - T_goal[:3, 3]) + rot_err = np.linalg.norm(T_sol[:3, :3] - T_goal[:3, :3]) + limits_ok = ( + len(graph.check_distance_limits(graph.realization(q_sol), tol=1e-6)) == 0 + ) + return { + "iters": out["iterations"], + "hvps": n_hvp[0], + "time": wall, + "trans_err": trans_err, + "rot_err": rot_err, + "success": trans_err < 1e-2 and rot_err < 1e-2 and limits_ok, + } + + +def main(n_goals=20, seed=0): + np.random.seed(seed) + robot, graph = load_ur10() + goals = [] + for _ in range(n_goals): + q = robot.random_configuration() + goals.append(robot.pose(q, f"p{robot.n}")) + + for precon in (None, "jacobi", "gn"): + rows = [run_case(graph, T, precon) for T in goals] + succ = sum(r["success"] for r in rows) + med = lambda k: np.median([r[k] for r in rows]) + tot = lambda k: np.sum([r[k] for r in rows]) + print( + f"{str(precon):>8}: success {succ}/{n_goals} | " + f"median outer {med('iters'):.0f} | median HVPs {med('hvps'):.0f} | " + f"HVPs/outer {tot('hvps')/tot('iters'):.1f} | " + f"median time {med('time')*1e3:.0f} ms | " + f"median trans_err {med('trans_err'):.2e}" + ) + + +if __name__ == "__main__": + main() diff --git a/graphik/solvers/riemannian_solver.py b/graphik/solvers/riemannian_solver.py index 18d3424..b18836b 100644 --- a/graphik/solvers/riemannian_solver.py +++ b/graphik/solvers/riemannian_solver.py @@ -9,12 +9,101 @@ MDS, linear_projection, gram_from_distance_matrix, + distance_matrix_from_pos, + memoize_last, ) from graphik.utils.manifolds.fixed_rank_psd_sym import PSDFixedRank from graphik.graphs.graph import ProblemGraph from graphik.solvers import rtr, loss +def _gn_weights(omega, psi_L, psi_U): + """Active-edge weight matrix for the Gauss-Newton model: equality edges + plus whichever bound edges are active at the current distances.""" + if psi_L is None: + return lambda D: omega + diff = psi_L != psi_U + LL = diff * (psi_L > 0) + UU = diff * (psi_U > 0) + + def weights(D): + return omega + LL * (D < psi_L) + UU * (D > psi_U) + + return weights + + +def _gn_blocks(Y, W): + """Per-edge Gauss-Newton blocks of the EDM loss at Y. + + The cost is sum over (undirected, active) edges of + (w_ij * (d_goal² - ||y_i - y_j||²))², whose GN Hessian has graph-Laplacian + block structure with edge blocks 8 w_ij u u^T, u = y_i - y_j. Returns the + (N, N, d, d) edge-block tensor E (E[i, i] = 0). + """ + Y_diff = Y[:, None, :] - Y[None, :, :] + return (8.0 * W)[:, :, None, None] * ( + Y_diff[..., :, None] * Y_diff[..., None, :] + ) + + +def make_gn_preconditioner(manifold, omega, psi_L=None, psi_U=None, + floor_rel=1e-3): + """Preconditioner for RTR's truncated CG: the inverse of the (eigenvalue- + floored) Gauss-Newton Hessian of the EDM loss. + + The GN matrix is N*d x N*d — small for IK-sized graphs — so one eigh per + outer iteration (memoized per Y) buys a near-exact Newton preconditioner. + The flooring keeps it positive definite: the GN matrix is PSD with a + nullspace containing the translation directions. Output is projected back + to the tangent (horizontal) space. + """ + weights = _gn_weights(omega, psi_L, psi_U) + + @memoize_last + def factor(Y): + N, d = Y.shape + D = distance_matrix_from_pos(Y) + E = _gn_blocks(Y, weights(D)) + B = -E + idx = np.arange(N) + B[idx, idx] = E.sum(axis=1) + H = B.transpose(0, 2, 1, 3).reshape(N * d, N * d) + lam, V = np.linalg.eigh(H) + lam = np.maximum(lam, max(floor_rel * lam[-1], 1e-12)) + return V, lam + + def precondition(Y, r): + V, lam = factor(Y) + z = (V @ ((V.T @ r.ravel()) / lam)).reshape(r.shape) + return manifold.projection(Y, z) + + return precondition + + +def make_jacobi_preconditioner(manifold, omega, psi_L=None, psi_U=None, + floor_rel=1e-3): + """Block-Jacobi variant of `make_gn_preconditioner`: keeps only the N + diagonal d x d blocks of the GN Hessian. Cheaper apply, weaker model.""" + weights = _gn_weights(omega, psi_L, psi_U) + + @memoize_last + def factor(Y): + D = distance_matrix_from_pos(Y) + E = _gn_blocks(Y, weights(D)) + Bd = E.sum(axis=1) # (N, d, d) diagonal blocks + lam, V = np.linalg.eigh(Bd) # batched + lam = np.maximum(lam, np.maximum(floor_rel * lam[:, -1:], 1e-12)) + return V, lam + + def precondition(Y, r): + V, lam = factor(Y) + rv = np.einsum("nij,ni->nj", V, r) + z = np.einsum("nij,nj->ni", V, rv / lam) + return manifold.projection(Y, z) + + return precondition + + class RiemannianSolver: def __init__( self, @@ -102,6 +191,7 @@ def solve( Y_init=None, method=None, output_log=True, + precon=None, ): manifold = PSDFixedRank( self.N, self.dim, @@ -109,6 +199,7 @@ def solve( cache_projection=self.cache, ) + psi_L = psi_U = None if not use_limits: cost, egrad, ehess = self.create_cost(D_goal, omega) else: @@ -117,6 +208,21 @@ def solve( D_goal, omega, psi_L, psi_U ) + if precon is None: + preconditioner = None + elif precon == "gn": + preconditioner = make_gn_preconditioner( + manifold, omega, psi_L=psi_L, psi_U=psi_U + ) + elif precon == "jacobi": + preconditioner = make_jacobi_preconditioner( + manifold, omega, psi_L=psi_L, psi_U=psi_U + ) + elif callable(precon): + preconditioner = precon + else: + raise ValueError("precon must be None, 'gn', 'jacobi', or callable") + if Y_init is None: Y_init = self.generate_initialization(D_goal, omega, bounds=bounds) @@ -131,6 +237,7 @@ def rhess(Y, U): min_gradient_norm=self.params.get("mingradnorm", 1e-8), theta=self.params.get("theta", 1.0), kappa=self.params.get("kappa", 0.1), + preconditioner=preconditioner, ) for k in ("rho_prime", "rho_regularization", "Delta_bar", "Delta0", "mininner", "maxinner"): @@ -149,7 +256,7 @@ def rhess(Y, U): } return res.point -def solve_with_riemannian(graph, T_goal, use_jit=False, cache=True): +def solve_with_riemannian(graph, T_goal, use_jit=False, cache=True, precon=None): """One-shot wrapper: build the IK problem from ``T_goal``, solve via ``RiemannianSolver``, decode joint angles from the recovered point cloud. @@ -162,6 +269,12 @@ def solve_with_riemannian(graph, T_goal, use_jit=False, cache=True): cache : bool, default True Memoize per-Y cost state and projection ops across cost/grad/HVP calls within a single iteration. Pure perf, identical trajectory. + precon : None | "gn" | "jacobi" | callable, default None + tCG preconditioner. "gn" (inverse Gauss-Newton Hessian, one small + eigh per outer iteration) cuts inner HVPs ~10x and wall time 2-5x + at equal-or-better success rate on the UR10 benchmark — see + experiments/rtr_preconditioner_study.py. Off by default because it + changes the optimization trajectory (baselines would shift). Uses the bsmooth initialization (the ``RiemannianSolver`` default), so the triangle-inequality bounds from ``bound_smoothing(G)`` are computed here @@ -172,7 +285,9 @@ def solve_with_riemannian(graph, T_goal, use_jit=False, cache=True): D_goal = distance_matrix_from_graph(G) omega = adjacency_matrix_from_graph(G) lb, ub = bound_smoothing(G) - sol_info = solver.solve(D_goal, omega, use_limits=True, bounds=(lb, ub)) + sol_info = solver.solve( + D_goal, omega, use_limits=True, bounds=(lb, ub), precon=precon + ) G_sol = graph_from_pos(sol_info["x"], graph.node_ids) q_sol = graph.joint_variables(G_sol, {f"p{graph.robot.n}": T_goal}) diff --git a/tests/test_preconditioner.py b/tests/test_preconditioner.py new file mode 100644 index 0000000..c9e7315 --- /dev/null +++ b/tests/test_preconditioner.py @@ -0,0 +1,84 @@ +"""Contract and end-to-end tests for the tCG preconditioners.""" +import numpy as np +import pytest + +from graphik.solvers.riemannian_solver import ( + RiemannianSolver, + make_gn_preconditioner, + make_jacobi_preconditioner, +) +from graphik.utils.dgp import ( + adjacency_matrix_from_graph, + bound_smoothing, + distance_matrix_from_graph, + graph_from_pos, +) +from graphik.utils.manifolds.fixed_rank_psd_sym import PSDFixedRank +from graphik.utils.roboturdf import load_ur10 + + +FACTORIES = [make_gn_preconditioner, make_jacobi_preconditioner] + + +def _setup(seed=3, N=10, d=3): + rng = np.random.default_rng(seed) + manifold = PSDFixedRank(N, d) + Y = rng.standard_normal((N, d)) + # A connected random edge mask (symmetric, zero diagonal) + omega = (rng.random((N, N)) < 0.5).astype(float) + omega = np.triu(omega, 1) + omega += omega.T + return manifold, Y, omega, rng + + +@pytest.mark.parametrize("factory", FACTORIES) +def test_output_is_tangent(factory): + manifold, Y, omega, rng = _setup() + precon = factory(manifold, omega) + r = manifold.projection(Y, rng.standard_normal(Y.shape)) + z = precon(Y, r) + # Horizontal-space condition for the PSDFixedRank quotient: Y^T Z symmetric + asym = Y.T @ z - z.T @ Y + np.testing.assert_allclose(asym, 0.0, atol=1e-9) + + +@pytest.mark.parametrize("factory", FACTORIES) +def test_symmetric_positive_definite_on_tangent_space(factory): + manifold, Y, omega, rng = _setup() + precon = factory(manifold, omega) + u = manifold.projection(Y, rng.standard_normal(Y.shape)) + v = manifold.projection(Y, rng.standard_normal(Y.shape)) + # Symmetry: == + np.testing.assert_allclose( + np.sum(u * precon(Y, v)), np.sum(precon(Y, u) * v), rtol=1e-9 + ) + # Positive definiteness along tangent directions + assert np.sum(u * precon(Y, u)) > 0 + assert np.sum(v * precon(Y, v)) > 0 + + +def test_gn_solve_matches_unpreconditioned_quality(): + np.random.seed(5) + robot, graph = load_ur10() + q = robot.random_configuration() + T_goal = robot.pose(q, f"p{robot.n}") + + G = graph.from_pose(T_goal) + D_goal = distance_matrix_from_graph(G) + omega = adjacency_matrix_from_graph(G) + bounds = bound_smoothing(G) + + solver = RiemannianSolver(graph) + out = solver.solve(D_goal, omega, use_limits=True, bounds=bounds, precon="gn") + G_sol = graph_from_pos(out["x"], graph.node_ids) + q_sol = graph.joint_variables(G_sol, {f"p{robot.n}": T_goal}) + T_sol = robot.pose(q_sol, f"p{robot.n}") + assert np.linalg.norm(T_sol[:3, 3] - T_goal[:3, 3]) < 1e-2 + assert np.linalg.norm(T_sol[:3, :3] - T_goal[:3, :3]) < 1e-2 + + +def test_unknown_precon_raises(): + robot, graph = load_ur10() + solver = RiemannianSolver(graph) + with pytest.raises(ValueError): + solver.solve(np.eye(4), np.zeros((4, 4)), precon="bogus") From c2b280656ba3facfd99dfa94a3fcd05b695377f3 Mon Sep 17 00:00:00 2001 From: Filip Maric Date: Thu, 11 Jun 2026 06:32:13 +0200 Subject: [PATCH 5/9] feat(solvers): zero-initialization option for the Riemannian solver Adds init='zero' alongside the spectral and bound-smoothing initializations; exposed in rtr_profile for comparison. --- experiments/rtr_profile.py | 11 +++--- graphik/solvers/riemannian_solver.py | 53 ++++++++++++++++++++++------ 2 files changed, 50 insertions(+), 14 deletions(-) diff --git a/experiments/rtr_profile.py b/experiments/rtr_profile.py index 728a145..19c7452 100644 --- a/experiments/rtr_profile.py +++ b/experiments/rtr_profile.py @@ -5,7 +5,7 @@ whole run sorted by cumulative time. Usage: - python experiments/rtr_profile.py [--n-poses 20] [--init spectral|bsmooth] + python experiments/rtr_profile.py [--n-poses 20] [--init spectral|bsmooth|zero] [--robot ur10] [--seed 0] [--top 20] [--jit] """ @@ -32,7 +32,10 @@ load_schunk_lwa4p, load_ur10, ) -from graphik.solvers.riemannian_solver import RiemannianSolver +from graphik.solvers.riemannian_solver import ( + RIEMANNIAN_INIT_STRATEGIES, + RiemannianSolver, +) ROBOTS = { @@ -48,7 +51,7 @@ def _run_one(graph, T_goal, init: str, jit: bool, params: dict) -> dict: G_partial = graph.from_pose(T_goal) D_goal = distance_matrix_from_graph(G_partial) omega = adjacency_matrix_from_graph(G_partial) - bounds = bound_smoothing(G_partial) if init in ("bsmooth", "bspectral") else None + bounds = bound_smoothing(G_partial) if init == "bsmooth" else None t0 = time.perf_counter() solver = RiemannianSolver(graph, jit=jit, init=init) @@ -118,7 +121,7 @@ def main() -> None: p = argparse.ArgumentParser(description=__doc__) p.add_argument("--n-poses", type=int, default=20, help="number of IK problems (default: 20)") p.add_argument("--robot", choices=tuple(ROBOTS), default="ur10") - p.add_argument("--init", choices=("spectral", "bsmooth"), default="spectral") + p.add_argument("--init", choices=RIEMANNIAN_INIT_STRATEGIES, default="spectral") p.add_argument("--seed", type=int, default=0) p.add_argument("--top", type=int, default=20, help="cProfile rows to print (default: 20)") p.add_argument("--jit", action="store_true", help="use AOT-compiled costgrd kernels") diff --git a/graphik/solvers/riemannian_solver.py b/graphik/solvers/riemannian_solver.py index b18836b..d20720e 100644 --- a/graphik/solvers/riemannian_solver.py +++ b/graphik/solvers/riemannian_solver.py @@ -11,12 +11,16 @@ gram_from_distance_matrix, distance_matrix_from_pos, memoize_last, + POS, ) from graphik.utils.manifolds.fixed_rank_psd_sym import PSDFixedRank from graphik.graphs.graph import ProblemGraph from graphik.solvers import rtr, loss +RIEMANNIAN_INIT_STRATEGIES = ("spectral", "bsmooth", "zero") + + def _gn_weights(omega, psi_L, psi_U): """Active-edge weight matrix for the Gauss-Newton model: equality edges plus whichever bound edges are active at the current distances.""" @@ -131,21 +135,36 @@ def __init__( # the solver's cost-state builder and the manifold's projection # operator. Both are perf-only — same algorithmic trajectory. self.cache = cache - if init not in ("spectral", "bsmooth"): - raise ValueError("init must be 'spectral' or 'bsmooth'") + if init not in RIEMANNIAN_INIT_STRATEGIES: + raise ValueError( + "init must be one of " + + ", ".join(repr(strategy) for strategy in RIEMANNIAN_INIT_STRATEGIES) + ) self.init = init + def zero_configuration_initialization(self): + """Return an RTR point-cloud factor from the robot zero configuration.""" + G_zero = self.graph.realization(self.graph.robot.zero_configuration()) + return np.stack( + [ + np.asarray(G_zero.nodes[node][POS], dtype=float) + for node in self.graph.node_ids + ] + ) + def generate_initialization(self, D_goal, omega, bounds=None): """Build a starting point for RTR. Dispatches on `self.init`: - - 'spectral' (default): Smith–Cai–Tasissa style. Center the partially- + - 'spectral': Smith–Cai–Tasissa style. Center the partially- observed distance matrix (zeros for unknowns), eigendecompose, take the top-d eigenvectors weighted by sqrt(eigvals). One N×N symmetric eigendecomposition. - - 'bsmooth': legacy. Triangle-inequality smoothing of distance bounds + - 'bsmooth' (default): legacy. Triangle-inequality smoothing of distance bounds (networkx all-pairs Bellman-Ford), sample an EDM at 0.9 of (lb, ub), classical MDS, then project to dim along the known-edge structure. Requires a ``bounds=(lb, ub)`` tuple (e.g. from ``bound_smoothing``). + - 'zero': realize the graph at ``robot.zero_configuration()`` and use + the resulting node positions directly, ordered by ``graph.node_ids``. """ if self.init == "spectral": n = D_goal.shape[0] @@ -161,7 +180,9 @@ def generate_initialization(self, D_goal, omega, bounds=None): lam = np.maximum(eigvals[top], 1e-12) return eigvecs[:, top] * np.sqrt(lam) - # bsmooth + if self.init == "zero": + return self.zero_configuration_initialization() + if bounds is None: raise ValueError("init='bsmooth' requires bounds=(lb, ub)") lb, ub = bounds @@ -256,7 +277,14 @@ def rhess(Y, U): } return res.point -def solve_with_riemannian(graph, T_goal, use_jit=False, cache=True, precon=None): +def solve_with_riemannian( + graph, + T_goal, + use_jit=False, + cache=True, + precon=None, + init="bsmooth", +): """One-shot wrapper: build the IK problem from ``T_goal``, solve via ``RiemannianSolver``, decode joint angles from the recovered point cloud. @@ -275,18 +303,23 @@ def solve_with_riemannian(graph, T_goal, use_jit=False, cache=True, precon=None) at equal-or-better success rate on the UR10 benchmark — see experiments/rtr_preconditioner_study.py. Off by default because it changes the optimization trajectory (baselines would shift). + init : {"spectral", "bsmooth", "zero"}, default "bsmooth" + Initialization strategy forwarded to ``RiemannianSolver``. Uses the bsmooth initialization (the ``RiemannianSolver`` default), so the triangle-inequality bounds from ``bound_smoothing(G)`` are computed here - and forwarded as the ``bounds`` argument to ``solve``. + and forwarded as the ``bounds`` argument to ``solve`` when needed. """ G = graph.from_pose(T_goal) - solver = RiemannianSolver(graph, jit=use_jit, cache=cache) + solver = RiemannianSolver(graph, jit=use_jit, cache=cache, init=init) D_goal = distance_matrix_from_graph(G) omega = adjacency_matrix_from_graph(G) - lb, ub = bound_smoothing(G) + bounds = None + if init == "bsmooth": + lb, ub = bound_smoothing(G) + bounds = (lb, ub) sol_info = solver.solve( - D_goal, omega, use_limits=True, bounds=(lb, ub), precon=precon + D_goal, omega, use_limits=True, bounds=bounds, precon=precon ) G_sol = graph_from_pos(sol_info["x"], graph.node_ids) q_sol = graph.joint_variables(G_sol, {f"p{graph.robot.n}": T_goal}) From 00c51265b7231460442b988655d3603af064efb0 Mon Sep 17 00:00:00 2001 From: Filip Maric Date: Thu, 11 Jun 2026 07:05:02 +0200 Subject: [PATCH 6/9] refactor(solvers)!: remove numba dependency (JIT/AOT cost kernels) The Gauss-Newton preconditioner removed the performance motivation for the AOT-compiled costgrd kernels; the dense NumPy loss backend (already the default and FD-verified) is now the only backend. - delete graphik/solvers/costs.py (numba.pycc AOT module) and the _numba_* wrappers in loss.py; drop the jit kwarg from the loss factories - hard-remove jit/use_jit from RiemannianSolver, NonlinearSolver and solve_with_riemannian; stale kwargs now raise TypeError instead of being silently swallowed by the **kwargs setattr loop - drop the njit branch from PSDFixedRank - drop numba/llvmlite from setup.py, the macOS lockfile, and the README AOT build instructions Default solver trajectories are unchanged (dense kernels untouched). Co-Authored-By: Claude Opus 4.8 --- README.md | 10 +- experiments/riemannian_example.py | 4 +- experiments/rtr_profile.py | 9 +- .../test_chain_2d_limits_new.py | 2 +- .../simple_ik_examples/test_chain_2d_new.py | 2 +- graphik/solvers/costs.py | 209 ------------------ graphik/solvers/loss.py | 83 +------ graphik/solvers/nonlinear_solver.py | 27 +-- graphik/solvers/riemannian_solver.py | 29 +-- graphik/utils/manifolds/fixed_rank_psd_sym.py | 11 +- requirements-macos-arm64.txt | 2 - setup.py | 1 - tests/baselines/runner.py | 4 +- tests/test_loss.py | 142 +----------- tests/test_removed_kwargs.py | 32 +++ 15 files changed, 82 insertions(+), 485 deletions(-) delete mode 100644 graphik/solvers/costs.py create mode 100644 tests/test_removed_kwargs.py diff --git a/README.md b/README.md index ab5ecf2..3417c81 100644 --- a/README.md +++ b/README.md @@ -23,14 +23,6 @@ pip install -r requirements-macos-arm64.txt pip install -e . --no-deps ``` -**Optional: numba AOT compilation.** GraphIK's Riemannian solver has an AOT-compiled cost-gradient kernel for speed. The default solver path runs without it (`use_jit=False`); for the JIT path, build the extension once after install: - -```bash -cd graphik/solvers && python costs.py -``` - -This produces a per-platform `.so` that's gitignored and env-local — regenerate it after each environment recreation. - ## SDP solvers (Mosek recommended) GraphIK's SDP-relaxation solvers (`solve_with_cidgik` in `graphik/solvers/convex_iteration.py`, the SDP formulations in `graphik/solvers/sdp_*.py`) run out of the box on a free solver (Clarabel by default, falling back to SCS or CVXOPT) bundled with `cvxpy`. `experiments/cidgik_example.py` works without any extra setup. @@ -86,7 +78,7 @@ The main purpose of our graphical interpretation of robot kinematics is to devel ```python from graphik.solvers.riemannian_solver import solve_with_riemannian -q_sol, solution_points = solve_with_riemannian(graph, T_goal, use_jit=False) # Returns None if infeasible or didn't solve +q_sol, solution_points = solve_with_riemannian(graph, T_goal) # Returns None if infeasible or didn't solve ``` For a similar example using [`CIDGIK`](https://arxiv.org/abs/2109.03374), a convex optimization-based approach, please see [experiments/cidgik_example.py](https://github.com/utiasSTARS/GraphIK/blob/main/experiments/cidgik_example.py). diff --git a/experiments/riemannian_example.py b/experiments/riemannian_example.py index 8208cfb..2ec5a08 100644 --- a/experiments/riemannian_example.py +++ b/experiments/riemannian_example.py @@ -22,9 +22,9 @@ q_goal = robot.random_configuration() T_goal = robot.pose(q_goal, f"p{robot.n}") # Can be any desired pose, this is just a simple example - # Run Riemannian solver without jit compiled cost function and gradient computations + # Run the Riemannian solver t0 = perf_counter() - q_sol, solution_points = solve_with_riemannian(graph, T_goal, use_jit=False) # Returns None if infeasible or didn't solve + q_sol, solution_points = solve_with_riemannian(graph, T_goal) # Returns None if infeasible or didn't solve solve_time = perf_counter() - t0 # Compare the solution's end effector pose to the goal. diff --git a/experiments/rtr_profile.py b/experiments/rtr_profile.py index 19c7452..62e55de 100644 --- a/experiments/rtr_profile.py +++ b/experiments/rtr_profile.py @@ -7,7 +7,6 @@ Usage: python experiments/rtr_profile.py [--n-poses 20] [--init spectral|bsmooth|zero] [--robot ur10] [--seed 0] [--top 20] - [--jit] """ from __future__ import annotations @@ -47,14 +46,14 @@ } -def _run_one(graph, T_goal, init: str, jit: bool, params: dict) -> dict: +def _run_one(graph, T_goal, init: str, params: dict) -> dict: G_partial = graph.from_pose(T_goal) D_goal = distance_matrix_from_graph(G_partial) omega = adjacency_matrix_from_graph(G_partial) bounds = bound_smoothing(G_partial) if init == "bsmooth" else None t0 = time.perf_counter() - solver = RiemannianSolver(graph, jit=jit, init=init) + solver = RiemannianSolver(graph, init=init) solver.params.update(params) sol = solver.solve(D_goal, omega, use_limits=True, bounds=bounds) wall = time.perf_counter() - t0 @@ -91,7 +90,6 @@ def _print_summary(results: list[dict], total_wall: float, args) -> None: print(f"=== RTR profile on {args.robot} ===") print(f" init: {args.init}") - print(f" jit: {args.jit}") print(f" n_poses: {args.n_poses}") print(f" seed: {args.seed}") print(f" total wall (incl prof): {total_wall*1000:.0f}ms") @@ -124,7 +122,6 @@ def main() -> None: p.add_argument("--init", choices=RIEMANNIAN_INIT_STRATEGIES, default="spectral") p.add_argument("--seed", type=int, default=0) p.add_argument("--top", type=int, default=20, help="cProfile rows to print (default: 20)") - p.add_argument("--jit", action="store_true", help="use AOT-compiled costgrd kernels") # Solver hyperparameter overrides (forwarded to RiemannianSolver.params). p.add_argument("--mingradnorm", type=float, default=None) p.add_argument("--maxiter", type=int, default=None) @@ -164,7 +161,7 @@ def main() -> None: profiler.enable() try: for T_goal in poses: - results.append(_run_one(graph, T_goal, args.init, args.jit, params)) + results.append(_run_one(graph, T_goal, args.init, params)) finally: profiler.disable() total_wall = time.perf_counter() - t0 diff --git a/experiments/simple_ik_examples/test_chain_2d_limits_new.py b/experiments/simple_ik_examples/test_chain_2d_limits_new.py index 44735ba..ec5cd96 100644 --- a/experiments/simple_ik_examples/test_chain_2d_limits_new.py +++ b/experiments/simple_ik_examples/test_chain_2d_limits_new.py @@ -56,7 +56,7 @@ def random_problem_2d_chain(): D_goal = distance_matrix_from_graph(G) omega = adjacency_matrix_from_graph(G) lb, ub = bound_smoothing(G) - sol_info = solver.solve(D_goal, omega, Y_init =Y_init, use_limits=True, bounds=(lb,ub), jit=False) + sol_info = solver.solve(D_goal, omega, Y_init =Y_init, use_limits=True, bounds=(lb,ub)) G_sol = graph_from_pos(sol_info["x"], graph.node_ids) q_sol = graph.joint_variables(G_sol, {f"p{graph.robot.n}": T_goal}) diff --git a/experiments/simple_ik_examples/test_chain_2d_new.py b/experiments/simple_ik_examples/test_chain_2d_new.py index 4ad58ff..0e6a832 100644 --- a/experiments/simple_ik_examples/test_chain_2d_new.py +++ b/experiments/simple_ik_examples/test_chain_2d_new.py @@ -55,7 +55,7 @@ def random_problem_2d_chain(): D_goal = distance_matrix_from_graph(G) omega = adjacency_matrix_from_graph(G) lb, ub = bound_smoothing(G) - sol_info = solver.solve(D_goal, omega, Y_init =Y_init, bounds=(lb,ub), jit=False) + sol_info = solver.solve(D_goal, omega, Y_init =Y_init, bounds=(lb,ub)) G_sol = graph_from_pos(sol_info["x"], graph.node_ids) q_sol = graph.joint_variables(G_sol, {f"p{graph.robot.n}": T_goal}) diff --git a/graphik/solvers/costs.py b/graphik/solvers/costs.py deleted file mode 100644 index 4f3f4c9..0000000 --- a/graphik/solvers/costs.py +++ /dev/null @@ -1,209 +0,0 @@ -import numpy as np -from numba import jit -from numba.pycc import CC - -cc = CC("costgrd") - -@cc.export("jcost", "f8(f8[:,:],f8[:,:],UniTuple(u8[:],2))") -def jcost(Y, D_goal, inds): - cost = 0 - dim = Y.shape[1] - for (idx, jdx) in zip(*inds): - nrm = 0 - for kdx in range(dim): - nrm += (Y[idx, kdx] - Y[jdx, kdx]) ** 2 - cost += 2 * (D_goal[idx, jdx] - nrm) ** 2 - return 0.5 * cost - - -@cc.export("jgrad", "f8[:,:](f8[:,:],f8[:,:],UniTuple(u8[:],2))") -def jgrad(Y, D_goal, inds): - num_el = Y.shape[0] - dim = Y.shape[1] - grad = np.zeros((num_el, dim)) - for (idx, jdx) in zip(*inds): - nrm = 0 - for kdx in range(dim): - nrm += (Y[idx, kdx] - Y[jdx, kdx]) ** 2 - for kdx in range(dim): - grad[idx, kdx] += ( - -4 * (D_goal[idx, jdx] - nrm) * (Y[idx, kdx] - Y[jdx, kdx]) - ) - grad[jdx, kdx] += ( - -4 * (D_goal[jdx, idx] - nrm) * (Y[jdx, kdx] - Y[idx, kdx]) - ) - return grad - - -@cc.export("jhess", "f8[:,:](f8[:,:],f8[:,:],f8[:,:],UniTuple(u8[:],2))") -def jhess(Y, w, D_goal, inds): - num_el = Y.shape[0] - dim = Y.shape[1] - hess = np.zeros((num_el, dim)) - for (idx, jdx) in zip(*inds): - nrm = 0 - sc = 0 - for kdx in range(dim): - sc += (Y[idx, kdx] - Y[jdx, kdx]) * (w[idx, kdx] - w[jdx, kdx]) - nrm += (Y[idx, kdx] - Y[jdx, kdx]) ** 2 - for kdx in range(dim): - hess[idx, kdx] += 4 * ( - 2 * sc * (Y[idx, kdx] - Y[jdx, kdx]) - + (nrm - D_goal[idx, jdx]) * (w[idx, kdx] - w[jdx, kdx]) - ) - hess[jdx, kdx] += 4 * ( - 2 * sc * (Y[jdx, kdx] - Y[idx, kdx]) - + (nrm - D_goal[jdx, idx]) * (w[jdx, kdx] - w[idx, kdx]) - ) - return hess - -@cc.export("jcost_and_grad", "Tuple((f8, f8[:,:]))(f8[:,:],f8[:,:],UniTuple(u8[:],2))") -def jcost_and_grad(Y, D_goal, inds): - cost = 0 - dim = Y.shape[1] - grad = np.zeros((Y.shape[0], dim)) - for (idx, jdx) in zip(*inds): - nrm = 0 - for kdx in range(dim): - nrm += (Y[idx, kdx] - Y[jdx, kdx]) ** 2 - for kdx in range(dim): - grad[idx, kdx] += ( - -4 * (D_goal[idx, jdx] - nrm) * (Y[idx, kdx] - Y[jdx, kdx]) - ) - grad[jdx, kdx] += ( - -4 * (D_goal[jdx, idx] - nrm) * (Y[jdx, kdx] - Y[idx, kdx]) - ) - cost += 2 * (D_goal[idx, jdx] - nrm) ** 2 - return 0.5*cost, grad - -@cc.export("lcost", "f8(f8[:,:],f8[:,:],f8[:,:],f8[:,:],f8[:,:],UniTuple(u8[:],2))") -def lcost(Y, D_goal, omega, psi_L, psi_U, inds): - cost = 0 - dim = Y.shape[1] - for (idx, jdx) in zip(*inds): - nrm = 0 - for kdx in range(dim): - nrm += (Y[idx, kdx] - Y[jdx, kdx]) ** 2 - if omega[idx, jdx]>0: - cost += (D_goal[idx, jdx] - nrm) ** 2 - if psi_L[idx, jdx]>0: - cost += max((psi_L[idx, jdx] - nrm), 0) ** 2 - if psi_U[idx, jdx]>0: - cost += max((-psi_U[idx, jdx] + nrm), 0) ** 2 - return cost - -@cc.export( - "lgrad", "f8[:,:](f8[:,:],f8[:,:],f8[:,:],f8[:,:],f8[:,:],UniTuple(u8[:],2))" -) -def lgrad(Y, D_goal, omega, psi_L, psi_U, inds): - num_el = Y.shape[0] - dim = Y.shape[1] - grad = np.zeros((num_el, dim)) - for (idx, jdx) in zip(*inds): - nrm = 0 - for kdx in range(dim): - nrm += (Y[idx, kdx] - Y[jdx, kdx]) ** 2 - if omega[idx, jdx]: - for kdx in range(dim): # NOTE inside to avoid needless triggering - a = (nrm - D_goal[idx, jdx]) * (Y[idx, kdx] - Y[jdx, kdx]) - grad[idx, kdx] += a - grad[jdx, kdx] += -a - if psi_L[idx, jdx]: - if max(psi_L[idx, jdx] - nrm, 0) > 0: - for kdx in range(dim): - a = (nrm - psi_L[idx, jdx]) * (Y[idx, kdx] - Y[jdx, kdx]) - grad[idx, kdx] += a - grad[jdx, kdx] += -a - if psi_U[idx, jdx]: - if max(-psi_U[idx, jdx] + nrm, 0) > 0: - for kdx in range(dim): - a = (nrm - psi_U[idx, jdx]) * (Y[idx, kdx] - Y[jdx, kdx]) - grad[idx, kdx] += a - grad[jdx, kdx] += -a - return 4*grad - -@cc.export("lcost_and_grad", "Tuple((f8, f8[:,:]))(f8[:,:],f8[:,:],f8[:,:],f8[:,:],f8[:,:],UniTuple(u8[:],2))") -def lcost_and_grad(Y, D_goal, omega, psi_L, psi_U, inds): - cost = 0 - dim = Y.shape[1] - num_el = Y.shape[0] - grad = np.zeros((num_el, dim)) - - for (idx, jdx) in zip(*inds): - nrm = 0 - for kdx in range(dim): - nrm += (Y[idx, kdx] - Y[jdx, kdx]) ** 2 - if omega[idx, jdx]>0: - cost += 2 * (D_goal[idx, jdx] - nrm) ** 2 - for kdx in range(dim): - grad[idx, kdx] += ( - 4 * (nrm - D_goal[idx, jdx]) * (Y[idx, kdx] - Y[jdx, kdx]) - ) - grad[jdx, kdx] += ( - 4 * (nrm - D_goal[jdx, idx]) * (Y[jdx, kdx] - Y[idx, kdx]) - ) - if psi_L[idx, jdx]>0: - cost += 2 * max((psi_L[idx, jdx] - nrm), 0) ** 2 - if max(psi_L[idx, jdx] - nrm, 0) > 0: - for kdx in range(dim): - grad[idx, kdx] += ( - 4 * (nrm - psi_L[idx, jdx]) * (Y[idx, kdx] - Y[jdx, kdx]) - ) - grad[jdx, kdx] += ( - 4 * (nrm - psi_L[jdx, idx]) * (Y[jdx, kdx] - Y[idx, kdx]) - ) - if psi_U[idx, jdx]>0: - cost += 2 * max((-psi_U[idx, jdx] + nrm), 0) ** 2 - if max(-psi_U[idx, jdx] + nrm, 0) > 0: - for kdx in range(dim): - grad[idx, kdx] += ( - 4 - * (nrm - psi_U[idx, jdx]) - * (Y[idx, kdx] - Y[jdx, kdx]) - ) - grad[jdx, kdx] += ( - 4 - * (nrm - psi_U[jdx, idx]) - * (Y[jdx, kdx] - Y[idx, kdx]) - ) - return 0.5*cost, grad - -@cc.export( - "lhess", - "f8[:,:](f8[:,:],f8[:,:],f8[:,:],f8[:,:],f8[:,:],f8[:,:],UniTuple(u8[:],2))", -) -def lhess(Y, w, D_goal, omega, psi_L, psi_U, inds): - num_el = Y.shape[0] - dim = Y.shape[1] - hess = np.zeros((num_el, dim)) - for (idx, jdx) in zip(*inds): - nrm = 0 - sc = 0 - for kdx in range(dim): - nrm += (Y[idx, kdx] - Y[jdx, kdx]) ** 2 - sc += (Y[idx, kdx] - Y[jdx, kdx]) * (w[idx, kdx] - w[jdx, kdx]) - if omega[idx, jdx]: - for kdx in range(dim): - a = 2 * sc * (Y[idx, kdx] - Y[jdx, kdx]) - b = (nrm - D_goal[idx, jdx]) * (w[idx, kdx] - w[jdx, kdx]) - c = a + b - hess[idx, kdx] += c - hess[jdx, kdx] += -c - if psi_L[idx, jdx] and max(psi_L[idx, jdx] - nrm, 0) > 0: - for kdx in range(dim): - a = 2 * sc * (Y[idx, kdx] - Y[jdx, kdx]) - b = (nrm - psi_L[idx, jdx]) * (w[idx, kdx] - w[jdx, kdx]) - c = a + b - hess[idx, kdx] += c - hess[jdx, kdx] += -c - if psi_U[idx, jdx] and max(-psi_U[idx, jdx] + nrm, 0) > 0: - for kdx in range(dim): - a = 2 * sc * (Y[idx, kdx] - Y[jdx, kdx]) - b = (nrm - psi_U[idx, jdx]) * (w[idx, kdx] - w[jdx, kdx]) - c = a + b - hess[idx, kdx] += c - hess[jdx, kdx] += -c - - return 4 * hess -if __name__ == "__main__": - cc.compile() diff --git a/graphik/solvers/loss.py b/graphik/solvers/loss.py index 0ba4876..bb70f03 100644 --- a/graphik/solvers/loss.py +++ b/graphik/solvers/loss.py @@ -1,15 +1,12 @@ """Single-source EDM loss kernels for the non-SDP IK solvers. -Two backends: -- NumPy dense (default): ``_dense_equality`` and ``_dense_limits`` build an S - matrix once per Y and reuse it across cost/egrad/ehvp via ``memoize_last``. -- Numba AOT: ``_numba_equality`` and ``_numba_limits`` wrap the kernels - exported from ``graphik/solvers/costs.py`` (compiled to ``costgrd.so``). +NumPy dense backend: ``_dense_equality`` and ``_dense_limits`` build an S +matrix once per Y and reuse it across cost/egrad/ehvp via ``memoize_last``. Two public factories: -- ``for_riemannian(D_goal, omega, *, psi_L, psi_U, jit, cache)`` returns +- ``for_riemannian(D_goal, omega, *, psi_L, psi_U, cache)`` returns ``(cost, egrad, ehvp)`` operating on ``(N, dim)`` Y arrays — RTR consumer. -- ``for_minimize(D_goal, omega, *, psi_L, psi_U, dim, jit, cache)`` returns +- ``for_minimize(D_goal, omega, *, psi_L, psi_U, dim, cache)`` returns ``(cost_and_grad, hessp)`` operating on flat ``Y_flat`` of length ``N*dim`` — scipy.optimize.minimize consumer. @@ -28,24 +25,6 @@ ) -def _import_costgrd(): - """Lazy-import the AOT-compiled cost kernels. Only invoked when jit=True.""" - try: - from graphik.solvers.costgrd import ( - jcost, jgrad, jhess, jcost_and_grad, - lcost, lgrad, lhess, lcost_and_grad, - ) - except ModuleNotFoundError as e: - raise ModuleNotFoundError( - "jit=True requires the AOT-compiled costgrd module. " - "Build it via 'cd graphik/solvers && python costs.py'." - ) from e - return ( - jcost, jgrad, jhess, jcost_and_grad, - lcost, lgrad, lhess, lcost_and_grad, - ) - - def _dense_equality(D_goal, omega, cache=True): """NumPy-dense (cost, egrad, ehvp) for the equality-only EDM loss. @@ -134,71 +113,24 @@ def ehvp(Y, Z): return cost, egrad, ehvp -def _numba_equality(D_goal, omega): - """Wraps costgrd.{jcost, jgrad, jhess}. Pure call-through; no Python state.""" - inds = np.nonzero(np.triu(omega)) - jcost, jgrad, jhess, *_ = _import_costgrd() - - def cost(Y): - return jcost(Y, D_goal, inds) - - def egrad(Y): - return jgrad(Y, D_goal, inds) - - def ehvp(Y, Z): - return jhess(Y, Z, D_goal, inds) - - return cost, egrad, ehvp - - -def _numba_limits(D_goal, omega, psi_L, psi_U): - """Wraps costgrd.{lcost, lgrad, lhess}. Pure call-through. - - The inds set is the union of the equality, lower-bound, and upper-bound - triangle masks; the AOT kernel decides per-pair which residual is active. - """ - diff = psi_L != psi_U - inds = np.nonzero( - np.triu(omega) + np.triu(diff * (psi_L > 0)) + np.triu(diff * (psi_U > 0)) - ) - *_, lcost, lgrad, lhess, _ = _import_costgrd() - - def cost(Y): - return lcost(Y, D_goal, omega, psi_L, psi_U, inds) - - def egrad(Y): - return lgrad(Y, D_goal, omega, psi_L, psi_U, inds) - - def ehvp(Y, Z): - return lhess(Y, Z, D_goal, omega, psi_L, psi_U, inds) - - return cost, egrad, ehvp - - def for_riemannian( D_goal, omega, *, psi_L=None, psi_U=None, - jit=False, cache=True, ): """Build (cost, egrad, ehvp) closures for an RTR-style consumer. All three accept Y of shape (N, dim) and return scalars / (N, dim) arrays. - State is memoized per-Y identity when ``cache=True``; this matters for - the dense backend (RTR's inner CG calls cost / egrad / ehvp at the same - Y), and is a no-op for numba_aot (which is stateless from Python). + State is memoized per-Y identity when ``cache=True``; this matters because + RTR's inner CG calls cost / egrad / ehvp at the same Y. Dispatch key is ``psi_L is None``: if ``psi_L is None``, the equality backend is used regardless of ``psi_U`` (callers pass both or neither in practice). """ - if jit: - if psi_L is None: - return _numba_equality(D_goal, omega) - return _numba_limits(D_goal, omega, psi_L, psi_U) if psi_L is None: return _dense_equality(D_goal, omega, cache=cache) return _dense_limits(D_goal, omega, psi_L, psi_U, cache=cache) @@ -211,7 +143,6 @@ def for_minimize( dim, psi_L=None, psi_U=None, - jit=False, cache=True, ): """Build (cost_and_grad, hessp) closures for scipy.optimize.minimize. @@ -225,7 +156,7 @@ def for_minimize( which lets the dense backend's ``memoize_last`` state cache hit. """ cost, egrad, ehvp = for_riemannian( - D_goal, omega, psi_L=psi_L, psi_U=psi_U, jit=jit, cache=cache + D_goal, omega, psi_L=psi_L, psi_U=psi_U, cache=cache ) last: list[Any] = [None, None] # [id(Y_flat), reshaped view] diff --git a/graphik/solvers/nonlinear_solver.py b/graphik/solvers/nonlinear_solver.py index 92f2ada..15677c6 100644 --- a/graphik/solvers/nonlinear_solver.py +++ b/graphik/solvers/nonlinear_solver.py @@ -20,25 +20,23 @@ class NonlinearSolver: - def __init__(self, graph: ProblemGraph, jit=False, *args, **kwargs): + def __init__(self, graph: ProblemGraph, *args, **kwargs): """Distance-based IK solver wrapping ``scipy.optimize.minimize``. - Cost / gradient / HVP come from ``graphik.solvers.loss.for_minimize``, - which dispatches between a NumPy-dense backend and the AOT-compiled - ``costgrd`` kernels based on ``jit``. + Cost / gradient / HVP come from ``graphik.solvers.loss.for_minimize`` + (dense NumPy backend). """ - if "cost_type" in kwargs: - raise TypeError( - "cost_type was removed in the loss-module consolidation. " - "The single dense backend is selected by default; pass " - "jit=True for the AOT-compiled kernels." - ) + for removed in ("cost_type", "jit"): + if removed in kwargs: + raise TypeError( + f"{removed} was removed: the numba/AOT backend is gone " + "and the dense NumPy loss backend is the only backend." + ) for key in kwargs: setattr(self, key, kwargs[key]) self.graph = graph self.dim = graph.dim self.N = graph.number_of_nodes() - self.jit = jit def generate_initialization(self, bounds, dim, omega): """Sample an EDM within the supplied (lb, ub) bounds, then MDS + project.""" @@ -49,14 +47,11 @@ def generate_initialization(self, bounds, dim, omega): return linear_projection(X_rand, omega, dim) def create_cost(self, D_goal, omega): - return loss.for_minimize( - D_goal, omega, dim=self.dim, jit=self.jit, - ) + return loss.for_minimize(D_goal, omega, dim=self.dim) def create_cost_limits(self, D_goal, omega, psi_L, psi_U): return loss.for_minimize( - D_goal, omega, psi_L=psi_L, psi_U=psi_U, - dim=self.dim, jit=self.jit, + D_goal, omega, psi_L=psi_L, psi_U=psi_U, dim=self.dim, ) def position_constraints(self): diff --git a/graphik/solvers/riemannian_solver.py b/graphik/solvers/riemannian_solver.py index d20720e..fb15650 100644 --- a/graphik/solvers/riemannian_solver.py +++ b/graphik/solvers/riemannian_solver.py @@ -112,25 +112,23 @@ class RiemannianSolver: def __init__( self, graph: ProblemGraph, - jit=False, cache=True, init="bsmooth", *args, **kwargs, ): - if "cost_type" in kwargs: - raise TypeError( - "cost_type was removed in the loss-module consolidation. " - "The single dense backend is selected by default; pass " - "jit=True for the AOT-compiled kernels." - ) + for removed in ("cost_type", "jit"): + if removed in kwargs: + raise TypeError( + f"{removed} was removed: the numba/AOT backend is gone " + "and the dense NumPy loss backend is the only backend." + ) self.params = {} for key in kwargs: setattr(self, key, kwargs[key]) self.graph = graph self.dim = graph.dim self.N = graph.number_of_nodes() - self.jit = jit # Single switch that enables both per-Y memoization caches: # the solver's cost-state builder and the manifold's projection # operator. Both are perf-only — same algorithmic trajectory. @@ -193,14 +191,11 @@ def generate_initialization(self, D_goal, omega, bounds=None): return linear_projection(X_rand, omega, self.dim) def create_cost(self, D_goal, omega): - return loss.for_riemannian( - D_goal, omega, jit=self.jit, cache=self.cache, - ) + return loss.for_riemannian(D_goal, omega, cache=self.cache) def create_cost_limits(self, D_goal, omega, psi_L, psi_U): return loss.for_riemannian( - D_goal, omega, psi_L=psi_L, psi_U=psi_U, - jit=self.jit, cache=self.cache, + D_goal, omega, psi_L=psi_L, psi_U=psi_U, cache=self.cache, ) def solve( @@ -216,7 +211,6 @@ def solve( ): manifold = PSDFixedRank( self.N, self.dim, - jit=self.jit, cache_projection=self.cache, ) @@ -280,7 +274,6 @@ def rhess(Y, U): def solve_with_riemannian( graph, T_goal, - use_jit=False, cache=True, precon=None, init="bsmooth", @@ -290,10 +283,6 @@ def solve_with_riemannian( Parameters ---------- - use_jit : bool, default False - Use the AOT-compiled cost kernels when True. False keeps the pure - numpy path; build the kernels first via ``python -m graphik.solvers.costs`` - before enabling this. cache : bool, default True Memoize per-Y cost state and projection ops across cost/grad/HVP calls within a single iteration. Pure perf, identical trajectory. @@ -311,7 +300,7 @@ def solve_with_riemannian( and forwarded as the ``bounds`` argument to ``solve`` when needed. """ G = graph.from_pose(T_goal) - solver = RiemannianSolver(graph, jit=use_jit, cache=cache, init=init) + solver = RiemannianSolver(graph, cache=cache, init=init) D_goal = distance_matrix_from_graph(G) omega = adjacency_matrix_from_graph(G) bounds = None diff --git a/graphik/utils/manifolds/fixed_rank_psd_sym.py b/graphik/utils/manifolds/fixed_rank_psd_sym.py index e3cd96c..919572c 100644 --- a/graphik/utils/manifolds/fixed_rank_psd_sym.py +++ b/graphik/utils/manifolds/fixed_rank_psd_sym.py @@ -42,7 +42,7 @@ class PSDFixedRank: "Low-Rank Optimization on the Cone of Positive Semidefinite Matrices". """ - def __init__(self, n, k, jit=False, cache_projection=False): + def __init__(self, n, k, cache_projection=False): self._n = n self._k = k self.name = f"YY' quotient manifold of {n}x{n} psd matrices of rank {k}" @@ -60,13 +60,8 @@ def __init__(self, n, k, jit=False, cache_projection=False): else: self._projection_op = self._build_projection_op - if jit: - from numba import njit - self.inner_product = njit(cache=True)(_inner_product_np) - self.norm = njit(cache=True)(_norm_np) - else: - self.inner_product = _inner_product_np - self.norm = _norm_np + self.inner_product = _inner_product_np + self.norm = _norm_np @property def typical_dist(self): diff --git a/requirements-macos-arm64.txt b/requirements-macos-arm64.txt index 6f12bad..d2a9fad 100644 --- a/requirements-macos-arm64.txt +++ b/requirements-macos-arm64.txt @@ -43,7 +43,6 @@ joblib==1.5.3 jsonschema==4.26.0 jsonschema-specifications==2025.9.1 kiwisolver==1.4.5 -llvmlite==0.47.0 lxml==5.2.1 manifold3d==3.4.1 mapbox_earcut==2.0.0 @@ -51,7 +50,6 @@ MarkupSafe==3.0.3 matplotlib==3.8.4 mpmath==1.3.0 networkx==3.6.1 -numba==0.65.1 numpy==2.4.4 osqp==1.1.1 packaging==24.0 diff --git a/setup.py b/setup.py index bdbaebf..d60f8d8 100644 --- a/setup.py +++ b/setup.py @@ -21,7 +21,6 @@ "networkx >= 3.6", "yourdfpy >= 0.0.60", "trimesh >= 4.6.1", - "numba >= 0.65", "pymlg @ git+https://github.com/decargroup/pymlg@8e6dc5ea61327ddfc2c8c1d16f276ae829f22db8#egg=pymlg", ], extras_require={ diff --git a/tests/baselines/runner.py b/tests/baselines/runner.py index 303c21b..ddf68cd 100644 --- a/tests/baselines/runner.py +++ b/tests/baselines/runner.py @@ -39,7 +39,7 @@ def _setup(case: Case): def _solve_riemannian(graph, T_goal): - q_sol, _ = solve_with_riemannian(graph, T_goal, use_jit=False) + q_sol, _ = solve_with_riemannian(graph, T_goal) return q_sol @@ -48,7 +48,7 @@ def _solve_nonlinear(graph, T_goal, method: str): D_goal = distance_matrix_from_graph(G) omega = adjacency_matrix_from_graph(G) lb, ub = bound_smoothing(G) - solver = NonlinearSolver(graph, jit=False) + solver = NonlinearSolver(graph) sol = solver.solve(D_goal, omega, use_limits=True, bounds=(lb, ub), method=method) G_sol = graph_from_pos(sol["x"], graph.node_ids) q_sol = graph.joint_variables(G_sol, {f"p{graph.robot.n}": T_goal}) diff --git a/tests/test_loss.py b/tests/test_loss.py index 889a4dc..95abd2b 100644 --- a/tests/test_loss.py +++ b/tests/test_loss.py @@ -1,9 +1,9 @@ """Finite-difference parity tests for graphik.solvers.loss. -Each backend (dense, numba_aot) is checked against an FD approximation -of its own gradient and HVP. egrad equals the analytical gradient of -cost; ehvp equals the analytical Hessian-vector product of cost (verified -by FD on egrad). +The dense backend is checked against an FD approximation of its own +gradient and HVP. egrad equals the analytical gradient of cost; ehvp +equals the analytical Hessian-vector product of cost (verified by FD +on egrad). """ from __future__ import annotations @@ -132,123 +132,17 @@ def test_active_set_exercises_both_branches(self): self.assertTrue(a2_seen, "no trial activated upper-bound branch (A2 < 0)") -class TestNumbaEquality(unittest.TestCase): - def setUp(self): - self.D_goal, self.omega, self.n, self.d, self.rng = _make_problem(seed=2) - self.dense = loss._dense_equality(self.D_goal, self.omega, cache=False) - self.numba = loss._numba_equality(self.D_goal, self.omega) - - def test_cost_matches_dense(self): - for trial in range(3): - Y = self.rng.standard_normal((self.n, self.d)) - np.testing.assert_allclose( - self.numba[0](Y), self.dense[0](Y), atol=1e-12, rtol=0, - err_msg=f"trial {trial}: numba cost != dense cost", - ) - - def test_egrad_matches_dense(self): - for trial in range(3): - Y = self.rng.standard_normal((self.n, self.d)) - np.testing.assert_allclose( - self.numba[1](Y), self.dense[1](Y), atol=1e-12, rtol=0, - err_msg=f"trial {trial}: numba egrad != dense egrad", - ) - - def test_ehvp_matches_dense(self): - for trial in range(3): - Y = self.rng.standard_normal((self.n, self.d)) - w = self.rng.standard_normal((self.n, self.d)) - np.testing.assert_allclose( - self.numba[2](Y, w), self.dense[2](Y, w), atol=1e-12, rtol=0, - err_msg=f"trial {trial}: numba ehvp != dense ehvp", - ) - - -class TestNumbaLimits(unittest.TestCase): - def setUp(self): - self.D_goal, self.omega, self.n, self.d, self.rng = _make_problem(seed=3) - # Two distinct pairs to exercise both active-set branches (mirrors - # TestDenseLimits.setUp's structure after the Task 3 fix). - psi_L = np.zeros_like(self.D_goal) - psi_U = np.zeros_like(self.D_goal) - pair_L = (0, self.n - 1) - pair_U = (1, self.n - 2) - psi_L[pair_L] = psi_L[pair_L[::-1]] = 5.0 * self.D_goal[pair_L] - psi_U[pair_U] = psi_U[pair_U[::-1]] = 0.2 * self.D_goal[pair_U] - self.psi_L, self.psi_U = psi_L, psi_U - self.dense = loss._dense_limits( - self.D_goal, self.omega, self.psi_L, self.psi_U, cache=False - ) - self.numba = loss._numba_limits( - self.D_goal, self.omega, self.psi_L, self.psi_U - ) - - def test_cost_matches_dense(self): - for trial in range(3): - Y = self.rng.standard_normal((self.n, self.d)) - np.testing.assert_allclose( - self.numba[0](Y), self.dense[0](Y), atol=1e-10, rtol=0, - err_msg=f"trial {trial}: numba cost != dense cost", - ) - - def test_egrad_matches_dense(self): - for trial in range(3): - Y = self.rng.standard_normal((self.n, self.d)) - np.testing.assert_allclose( - self.numba[1](Y), self.dense[1](Y), atol=1e-10, rtol=0, - err_msg=f"trial {trial}: numba egrad != dense egrad", - ) - - def test_ehvp_matches_dense(self): - for trial in range(3): - Y = self.rng.standard_normal((self.n, self.d)) - w = self.rng.standard_normal((self.n, self.d)) - np.testing.assert_allclose( - self.numba[2](Y, w), self.dense[2](Y, w), atol=1e-10, rtol=0, - err_msg=f"trial {trial}: numba ehvp != dense ehvp", - ) - - def test_active_set_exercises_both_branches(self): - """Sanity check: across 3 trials, both A1 (lower-bound) and A2 - (upper-bound) branches are active at least once. If this regresses, - the cost/egrad/ehvp parity tests above stop covering the m4 / m5 - paths and a sign bug in either AOT branch becomes invisible. - """ - from graphik.utils.dgp import distance_matrix_from_pos - rng = np.random.default_rng(99) # independent rng so test ordering doesn't matter - a1_seen = a2_seen = False - LL = (self.psi_L != self.psi_U) * (self.psi_L > 0) - UU = (self.psi_L != self.psi_U) * (self.psi_U > 0) - for _ in range(3): - Y = rng.standard_normal((self.n, self.d)) - D = distance_matrix_from_pos(Y) - A1 = np.maximum(self.psi_L - LL * D, 0) - A2 = -np.maximum(-self.psi_U + UU * D, 0) - if (A1 > 0).any(): - a1_seen = True - if (A2 < 0).any(): - a2_seen = True - self.assertTrue(a1_seen, "no trial activated lower-bound branch (A1 > 0)") - self.assertTrue(a2_seen, "no trial activated upper-bound branch (A2 < 0)") - - class TestForRiemannian(unittest.TestCase): def setUp(self): self.D_goal, self.omega, self.n, self.d, self.rng = _make_problem(seed=4) def test_dense_dispatch_equality(self): - cost, egrad, _ = loss.for_riemannian(self.D_goal, self.omega, jit=False) + cost, egrad, _ = loss.for_riemannian(self.D_goal, self.omega) ref_cost, ref_egrad, _ = loss._dense_equality(self.D_goal, self.omega, cache=True) Y = self.rng.standard_normal((self.n, self.d)) np.testing.assert_allclose(cost(Y), ref_cost(Y), atol=1e-12, rtol=0) np.testing.assert_allclose(egrad(Y), ref_egrad(Y), atol=1e-12, rtol=0) - def test_numba_dispatch_equality(self): - cost, _, _ = loss.for_riemannian(self.D_goal, self.omega, jit=True) - ref_cost, _, _ = loss._numba_equality(self.D_goal, self.omega) - Y = self.rng.standard_normal((self.n, self.d)) - np.testing.assert_allclose(cost(Y), ref_cost(Y), atol=1e-12, rtol=0) - def test_dense_dispatch_limits(self): psi_L = np.zeros_like(self.D_goal) psi_U = np.zeros_like(self.D_goal) @@ -257,7 +151,7 @@ def test_dense_dispatch_limits(self): psi_L[pair_L] = psi_L[pair_L[::-1]] = 5.0 * self.D_goal[pair_L] psi_U[pair_U] = psi_U[pair_U[::-1]] = 0.2 * self.D_goal[pair_U] cost, _, _ = loss.for_riemannian( - self.D_goal, self.omega, psi_L=psi_L, psi_U=psi_U, jit=False + self.D_goal, self.omega, psi_L=psi_L, psi_U=psi_U ) ref_cost, _, _ = loss._dense_limits( self.D_goal, self.omega, psi_L, psi_U, cache=True @@ -265,22 +159,6 @@ def test_dense_dispatch_limits(self): Y = self.rng.standard_normal((self.n, self.d)) np.testing.assert_allclose(cost(Y), ref_cost(Y), atol=1e-12, rtol=0) - def test_numba_dispatch_limits(self): - psi_L = np.zeros_like(self.D_goal) - psi_U = np.zeros_like(self.D_goal) - pair_L = (0, self.n - 1) - pair_U = (1, self.n - 2) - psi_L[pair_L] = psi_L[pair_L[::-1]] = 5.0 * self.D_goal[pair_L] - psi_U[pair_U] = psi_U[pair_U[::-1]] = 0.2 * self.D_goal[pair_U] - cost, _, _ = loss.for_riemannian( - self.D_goal, self.omega, psi_L=psi_L, psi_U=psi_U, jit=True - ) - ref_cost, _, _ = loss._numba_limits( - self.D_goal, self.omega, psi_L, psi_U - ) - Y = self.rng.standard_normal((self.n, self.d)) - np.testing.assert_allclose(cost(Y), ref_cost(Y), atol=1e-12, rtol=0) - class TestForMinimize(unittest.TestCase): def setUp(self): @@ -288,7 +166,7 @@ def setUp(self): def test_cost_and_grad_shapes(self): cost_and_grad, hessp = loss.for_minimize( - self.D_goal, self.omega, dim=self.d, jit=False + self.D_goal, self.omega, dim=self.d ) Y_flat = self.rng.standard_normal(self.n * self.d) f, g = cost_and_grad(Y_flat) @@ -300,9 +178,9 @@ def test_cost_and_grad_shapes(self): def test_cost_and_grad_matches_riemannian(self): # for_minimize must return the same numbers as for_riemannian, just flattened. - cost, egrad, _ = loss.for_riemannian(self.D_goal, self.omega, jit=False) + cost, egrad, _ = loss.for_riemannian(self.D_goal, self.omega) cost_and_grad, _ = loss.for_minimize( - self.D_goal, self.omega, dim=self.d, jit=False + self.D_goal, self.omega, dim=self.d ) Y = self.rng.standard_normal((self.n, self.d)) f_ref = cost(Y) @@ -320,7 +198,7 @@ def test_limits_dispatch_shapes(self): psi_U[pair_U] = psi_U[pair_U[::-1]] = 0.2 * self.D_goal[pair_U] cost_and_grad, hessp = loss.for_minimize( self.D_goal, self.omega, dim=self.d, - psi_L=psi_L, psi_U=psi_U, jit=False, + psi_L=psi_L, psi_U=psi_U, ) Y_flat = self.rng.standard_normal(self.n * self.d) f, g = cost_and_grad(Y_flat) diff --git a/tests/test_removed_kwargs.py b/tests/test_removed_kwargs.py new file mode 100644 index 0000000..e909121 --- /dev/null +++ b/tests/test_removed_kwargs.py @@ -0,0 +1,32 @@ +"""Removed-kwarg guards on the solver constructors. + +Both solver __init__s setattr unknown kwargs onto self, so without an +explicit guard a stale ``jit=True`` from a pre-removal caller would be +silently swallowed instead of erroring. +""" +import unittest + +from graphik.solvers.nonlinear_solver import NonlinearSolver +from graphik.solvers.riemannian_solver import RiemannianSolver + + +class TestRemovedKwargs(unittest.TestCase): + def test_riemannian_rejects_jit(self): + with self.assertRaisesRegex(TypeError, "jit"): + RiemannianSolver(None, jit=True) + + def test_riemannian_rejects_cost_type(self): + with self.assertRaisesRegex(TypeError, "cost_type"): + RiemannianSolver(None, cost_type="dense") + + def test_nonlinear_rejects_jit(self): + with self.assertRaisesRegex(TypeError, "jit"): + NonlinearSolver(None, jit=True) + + def test_nonlinear_rejects_cost_type(self): + with self.assertRaisesRegex(TypeError, "cost_type"): + NonlinearSolver(None, cost_type="dense") + + +if __name__ == "__main__": + unittest.main() From d8dbe910e8d4d5edefb2b3a439942aebfc02deb9 Mon Sep 17 00:00:00 2001 From: Filip Maric Date: Thu, 11 Jun 2026 08:49:00 +0200 Subject: [PATCH 7/9] refactor(solvers)!: consolidate local solver API Add IKSolver and IKResult, move shared EDM plumbing into DistanceSolver, rename loss.py to costs.py, single-source initialization strategies, and replace the old local solver modules with RiemannianSolver, ScipySolver, and JointAngleSolver on solve(T_goal). Migrate tests, baselines runner, benchmarks, examples, and README to the problem-level API. SDP solvers are unchanged. --- README.md | 9 +- experiments/riemannian_example.py | 11 +- experiments/rtr_preconditioner_study.py | 51 +-- experiments/rtr_profile.py | 99 +++--- .../test_chain_2d_limits_new.py | 43 +-- .../simple_ik_examples/test_chain_2d_new.py | 45 +-- experiments/solver_benchmark.py | 86 ++--- graphik/solvers/__init__.py | 12 + graphik/solvers/base.py | 57 ++++ graphik/solvers/{loss.py => costs.py} | 92 ++++- graphik/solvers/distance_solver.py | 113 +++++++ graphik/solvers/initializations.py | 49 +++ graphik/solvers/joint_angle.py | 80 +++++ graphik/solvers/joint_angle_solver.py | 135 -------- graphik/solvers/nonlinear_solver.py | 122 ------- graphik/solvers/riemannian.py | 171 ++++++++++ graphik/solvers/riemannian_solver.py | 320 ------------------ graphik/solvers/scipy_solver.py | 87 +++++ tests/baselines/runner.py | 41 +-- tests/helpers.py | 20 ++ tests/{test_loss.py => test_costs.py} | 107 +++++- tests/test_distance_solver.py | 66 ++++ tests/test_initializations.py | 49 +++ tests/test_joint_angle_solver.py | 130 +++---- tests/test_nonlinear_solver_gradients.py | 81 ----- tests/test_preconditioner.py | 29 +- tests/test_pymlg_conventions.py | 5 +- tests/test_removed_kwargs.py | 26 +- tests/test_scipy_solver.py | 36 ++ tests/test_solver_contract.py | 112 ++++++ 30 files changed, 1259 insertions(+), 1025 deletions(-) create mode 100644 graphik/solvers/base.py rename graphik/solvers/{loss.py => costs.py} (65%) create mode 100644 graphik/solvers/distance_solver.py create mode 100644 graphik/solvers/initializations.py create mode 100644 graphik/solvers/joint_angle.py delete mode 100644 graphik/solvers/joint_angle_solver.py delete mode 100644 graphik/solvers/nonlinear_solver.py create mode 100644 graphik/solvers/riemannian.py delete mode 100644 graphik/solvers/riemannian_solver.py create mode 100644 graphik/solvers/scipy_solver.py create mode 100644 tests/helpers.py rename tests/{test_loss.py => test_costs.py} (67%) create mode 100644 tests/test_distance_solver.py create mode 100644 tests/test_initializations.py delete mode 100644 tests/test_nonlinear_solver_gradients.py create mode 100644 tests/test_scipy_solver.py create mode 100644 tests/test_solver_contract.py diff --git a/README.md b/README.md index 3417c81..c050343 100644 --- a/README.md +++ b/README.md @@ -74,11 +74,14 @@ T_goal = robot.pose(q_goal, f"p{robot.n}") ``` ### 4. Solve the IK Problem -The main purpose of our graphical interpretation of robot kinematics is to develop distance-geometric IK solvers. One example is the [Riemannian optimization-based solver](https://arxiv.org/abs/2011.04850) implemented in [`RiemannianSolver`](https://github.com/utiasSTARS/GraphIK/blob/main/graphik/solvers/riemannian_solver.py). +The main purpose of our graphical interpretation of robot kinematics is to develop distance-geometric IK solvers. One example is the [Riemannian optimization-based solver](https://arxiv.org/abs/2011.04850) implemented in [`RiemannianSolver`](https://github.com/utiasSTARS/GraphIK/blob/main/graphik/solvers/riemannian.py). ```python -from graphik.solvers.riemannian_solver import solve_with_riemannian -q_sol, solution_points = solve_with_riemannian(graph, T_goal) # Returns None if infeasible or didn't solve +from graphik.solvers import RiemannianSolver + +solver = RiemannianSolver(graph) +result = solver.solve(T_goal) +q_sol = result.q if result.feasible else None # feasible == within joint limits ``` For a similar example using [`CIDGIK`](https://arxiv.org/abs/2109.03374), a convex optimization-based approach, please see [experiments/cidgik_example.py](https://github.com/utiasSTARS/GraphIK/blob/main/experiments/cidgik_example.py). diff --git a/experiments/riemannian_example.py b/experiments/riemannian_example.py index 2ec5a08..953fcce 100644 --- a/experiments/riemannian_example.py +++ b/experiments/riemannian_example.py @@ -1,7 +1,7 @@ from time import perf_counter from graphik.utils.utils import table_environment -from graphik.solvers.riemannian_solver import solve_with_riemannian +from graphik.solvers import RiemannianSolver # Multiple robot models to try out, or you can implement your own from graphik.utils.roboturdf import load_ur10 @@ -23,8 +23,9 @@ T_goal = robot.pose(q_goal, f"p{robot.n}") # Can be any desired pose, this is just a simple example # Run the Riemannian solver + solver = RiemannianSolver(graph) t0 = perf_counter() - q_sol, solution_points = solve_with_riemannian(graph, T_goal) # Returns None if infeasible or didn't solve + result = solver.solve(T_goal) solve_time = perf_counter() - t0 # Compare the solution's end effector pose to the goal. @@ -36,10 +37,10 @@ print(q_goal) print("--------------------------------------------") print(f"Solve time: {solve_time:.4f} s") - if q_sol: + if result.feasible: print("Riemannian solution's pose: ") - print(robot.pose(q_sol, f"p{robot.n}")) + print(robot.pose(result.q, f"p{robot.n}")) print("Riemannian configuration: ") - print(q_sol) + print(result.q) else: print("Riemannian did not return a feasible solution.") diff --git a/experiments/rtr_preconditioner_study.py b/experiments/rtr_preconditioner_study.py index c31fa03..9e7e54a 100644 --- a/experiments/rtr_preconditioner_study.py +++ b/experiments/rtr_preconditioner_study.py @@ -1,41 +1,21 @@ -"""Preconditioner study for the Riemannian solver's truncated CG. - -Compares no preconditioner vs block-Jacobi vs the full Gauss-Newton -Laplacian on UR10 with joint limits (the production configuration). -Counts HVPs (the profile's hottest operation), outer iterations, wall -time, and solution quality (final pose error after angle decoding). -""" +"""Preconditioner study for the Riemannian solver's truncated CG.""" import time import numpy as np -from graphik.solvers import rtr -from graphik.solvers.riemannian_solver import RiemannianSolver -from graphik.utils.dgp import ( - adjacency_matrix_from_graph, - bound_smoothing, - distance_matrix_from_graph, - graph_from_pos, -) +from graphik.solvers import RiemannianSolver +import graphik.solvers.costs as costs_mod from graphik.utils.roboturdf import load_ur10 def run_case(graph, T_goal, precon): - solver = RiemannianSolver(graph) - G = graph.from_pose(T_goal) - D_goal = distance_matrix_from_graph(G) - omega = adjacency_matrix_from_graph(G) - lb, ub = bound_smoothing(G) + solver = RiemannianSolver(graph, precon=precon) n_hvp = [0] - orig_for_riemannian = __import__( - "graphik.solvers.loss", fromlist=["for_riemannian"] - ).for_riemannian - - import graphik.solvers.loss as loss_mod + saved = costs_mod.for_riemannian def counting_for_riemannian(*args, **kwargs): - cost, egrad, ehvp = orig_for_riemannian(*args, **kwargs) + cost, egrad, ehvp = saved(*args, **kwargs) def counted_ehvp(Y, Z): n_hvp[0] += 1 @@ -43,31 +23,24 @@ def counted_ehvp(Y, Z): return cost, egrad, counted_ehvp - loss_mod.for_riemannian, saved = counting_for_riemannian, loss_mod.for_riemannian + costs_mod.for_riemannian = counting_for_riemannian try: t0 = time.perf_counter() - out = solver.solve( - D_goal, omega, use_limits=True, bounds=(lb, ub), precon=precon - ) + res = solver.solve(T_goal) wall = time.perf_counter() - t0 finally: - loss_mod.for_riemannian = saved + costs_mod.for_riemannian = saved - G_sol = graph_from_pos(out["x"], graph.node_ids) - q_sol = graph.joint_variables(G_sol, {f"p{graph.robot.n}": T_goal}) - T_sol = graph.robot.pose(q_sol, f"p{graph.robot.n}") + T_sol = graph.robot.pose(res.q, f"p{graph.robot.n}") trans_err = np.linalg.norm(T_sol[:3, 3] - T_goal[:3, 3]) rot_err = np.linalg.norm(T_sol[:3, :3] - T_goal[:3, :3]) - limits_ok = ( - len(graph.check_distance_limits(graph.realization(q_sol), tol=1e-6)) == 0 - ) return { - "iters": out["iterations"], + "iters": res.iterations, "hvps": n_hvp[0], "time": wall, "trans_err": trans_err, "rot_err": rot_err, - "success": trans_err < 1e-2 and rot_err < 1e-2 and limits_ok, + "success": trans_err < 1e-2 and rot_err < 1e-2 and res.feasible, } diff --git a/experiments/rtr_profile.py b/experiments/rtr_profile.py index 62e55de..d0cd428 100644 --- a/experiments/rtr_profile.py +++ b/experiments/rtr_profile.py @@ -1,12 +1,7 @@ """Profile the in-house RTR solver on a single robot. -Runs N IK problems on ur10 (or another robot via --robot), reports per- -pose wall-clock and convergence stats, and dumps a cProfile of the -whole run sorted by cumulative time. - -Usage: - python experiments/rtr_profile.py [--n-poses 20] [--init spectral|bsmooth|zero] - [--robot ur10] [--seed 0] [--top 20] +Runs N IK problems, reports per-pose wall-clock/convergence stats, and dumps +a cProfile of the whole run sorted by cumulative time. """ from __future__ import annotations @@ -18,12 +13,8 @@ import numpy as np -from graphik.utils import ( - adjacency_matrix_from_graph, - bound_smoothing, - distance_matrix_from_graph, - graph_from_pos, -) +from graphik.solvers import RiemannianSolver +from graphik.solvers.initializations import INIT_STRATEGIES from graphik.utils.roboturdf import ( load_kuka, load_panda, @@ -31,10 +22,6 @@ load_schunk_lwa4p, load_ur10, ) -from graphik.solvers.riemannian_solver import ( - RIEMANNIAN_INIT_STRATEGIES, - RiemannianSolver, -) ROBOTS = { @@ -47,32 +34,31 @@ def _run_one(graph, T_goal, init: str, params: dict) -> dict: - G_partial = graph.from_pose(T_goal) - D_goal = distance_matrix_from_graph(G_partial) - omega = adjacency_matrix_from_graph(G_partial) - bounds = bound_smoothing(G_partial) if init == "bsmooth" else None - + solver = RiemannianSolver(graph, init=init, rtr_params=params) t0 = time.perf_counter() - solver = RiemannianSolver(graph, init=init) - solver.params.update(params) - sol = solver.solve(D_goal, omega, use_limits=True, bounds=bounds) - wall = time.perf_counter() - t0 + try: + res = solver.solve(T_goal) + except Exception: + return { + "wall_s": time.perf_counter() - t0, + "iterations": 0, + "stopping_criterion": "solve/decode raised", + "fx": float("nan"), + "q_sol": None, + "solver_feasible": False, + } return { - "wall_s": wall, - "rtr_time_s": sol["time"], - "iterations": sol["iterations"], - "stopping_criterion": sol["stopping_criterion"], - "gradnorm": sol["gradnorm"], - "fx": sol["f(x)"], - "Y_sol": sol["x"], + "wall_s": time.perf_counter() - t0, + "iterations": res.iterations, + "stopping_criterion": res.status, + "fx": res.cost, + "q_sol": res.q, + "solver_feasible": res.feasible, } -def _pose_error(graph, robot, ee: str, T_goal: np.ndarray, Y_sol: np.ndarray) -> tuple[bool, float, float]: - G_sol = graph_from_pos(Y_sol, graph.node_ids) - try: - q_sol = graph.joint_variables(G_sol, {ee: T_goal}) - except Exception: +def _pose_error(robot, ee: str, T_goal: np.ndarray, q_sol) -> tuple[bool, float, float]: + if q_sol is None: return False, float("nan"), float("nan") T_got = np.asarray(robot.pose(q_sol, ee)) pos = float(np.linalg.norm(T_goal[:3, 3] - T_got[:3, 3])) @@ -84,7 +70,6 @@ def _pose_error(graph, robot, ee: str, T_goal: np.ndarray, Y_sol: np.ndarray) -> def _print_summary(results: list[dict], total_wall: float, args) -> None: walls = np.array([r["wall_s"] for r in results]) - rtr_times = np.array([r["rtr_time_s"] for r in results]) iters = np.array([r["iterations"] for r in results]) feas = np.array([r["feasible"] for r in results], dtype=bool) @@ -94,9 +79,14 @@ def _print_summary(results: list[dict], total_wall: float, args) -> None: print(f" seed: {args.seed}") print(f" total wall (incl prof): {total_wall*1000:.0f}ms") print(f" feasible: {feas.sum()}/{args.n_poses} ({feas.mean()*100:.1f}%)") - print(f" iterations: mean {iters.mean():.1f} median {np.median(iters):.0f} max {iters.max()}") - print(f" per-pose wall: mean {walls.mean()*1000:.1f}ms median {np.median(walls)*1000:.1f}ms max {walls.max()*1000:.1f}ms") - print(f" rtr-internal time: mean {rtr_times.mean()*1000:.1f}ms (vs {walls.mean()*1000:.1f}ms wall — overhead {(walls.mean()-rtr_times.mean())*1000:.1f}ms)") + print( + f" iterations: mean {iters.mean():.1f} " + f"median {np.median(iters):.0f} max {iters.max()}" + ) + print( + f" per-pose wall: mean {walls.mean()*1000:.1f}ms " + f"median {np.median(walls)*1000:.1f}ms max {walls.max()*1000:.1f}ms" + ) if feas.any(): pos = np.array([r["pos_err"] for r in results])[feas] rot = np.array([r["rot_err"] for r in results])[feas] @@ -117,12 +107,11 @@ def _print_profile(profiler: cProfile.Profile, top: int) -> None: def main() -> None: p = argparse.ArgumentParser(description=__doc__) - p.add_argument("--n-poses", type=int, default=20, help="number of IK problems (default: 20)") + p.add_argument("--n-poses", type=int, default=20, help="number of IK problems") p.add_argument("--robot", choices=tuple(ROBOTS), default="ur10") - p.add_argument("--init", choices=RIEMANNIAN_INIT_STRATEGIES, default="spectral") + p.add_argument("--init", choices=INIT_STRATEGIES, default="spectral") p.add_argument("--seed", type=int, default=0) - p.add_argument("--top", type=int, default=20, help="cProfile rows to print (default: 20)") - # Solver hyperparameter overrides (forwarded to RiemannianSolver.params). + p.add_argument("--top", type=int, default=20) p.add_argument("--mingradnorm", type=float, default=None) p.add_argument("--maxiter", type=int, default=None) p.add_argument("--theta", type=float, default=None) @@ -146,14 +135,17 @@ def main() -> None: ("delta_bar", "Delta_bar"), ("delta0", "Delta0"), ]: - v = getattr(args, cli_name) - if v is not None: - params[key] = v + value = getattr(args, cli_name) + if value is not None: + params[key] = value np.random.seed(args.seed) robot, graph = ROBOTS[args.robot]() ee = f"p{robot.n}" - poses = [np.asarray(robot.pose(robot.random_configuration(), ee)) for _ in range(args.n_poses)] + poses = [ + np.asarray(robot.pose(robot.random_configuration(), ee)) + for _ in range(args.n_poses) + ] profiler = cProfile.Profile() results: list[dict] = [] @@ -166,11 +158,10 @@ def main() -> None: profiler.disable() total_wall = time.perf_counter() - t0 - # Feasibility check is outside the profile so it doesn't pollute solver - # function timings; it's a tiny pose-FK roundtrip per pose. + # Decode is inside solve; this loop only computes pose-FK error. for r, T_goal in zip(results, poses): - ok, pos_err, rot_err = _pose_error(graph, robot, ee, T_goal, r["Y_sol"]) - r["feasible"] = ok + ok, pos_err, rot_err = _pose_error(robot, ee, T_goal, r["q_sol"]) + r["feasible"] = r["solver_feasible"] and ok r["pos_err"] = pos_err r["rot_err"] = rot_err diff --git a/experiments/simple_ik_examples/test_chain_2d_limits_new.py b/experiments/simple_ik_examples/test_chain_2d_limits_new.py index ec5cd96..9a4e53b 100644 --- a/experiments/simple_ik_examples/test_chain_2d_limits_new.py +++ b/experiments/simple_ik_examples/test_chain_2d_limits_new.py @@ -1,30 +1,21 @@ #!/usr/bin/env python3 import numpy as np -from numpy.testing import assert_array_less -import time + from graphik.graphs import ProblemGraph from graphik.robots import Robot -from graphik.utils.dgp import bound_smoothing -from graphik.utils import ( - pos_from_graph, - list_to_variable_dict, - distance_matrix_from_graph, - adjacency_matrix_from_graph, - graph_from_pos - ) -from graphik.solvers.riemannian_solver import RiemannianSolver +from graphik.solvers import RiemannianSolver +from graphik.utils import list_to_variable_dict, pos_from_graph + def random_problem_2d_chain(): e_rot = [] e_pos = [] - t_sol = [] n = 10 fails = 0 a = list_to_variable_dict(np.ones(n)) th = list_to_variable_dict(np.zeros(n)) - # angular_limits = np.minimum(np.random.rand(n) * np.pi + 0.20, np.pi) - angular_limits = np.array((n-1)*[np.pi/2] + [np.pi]) + angular_limits = np.array((n - 1) * [np.pi / 2] + [np.pi]) upper_angular_limits = list_to_variable_dict(angular_limits) lower_angular_limits = list_to_variable_dict(-angular_limits) robot_params = { @@ -39,33 +30,25 @@ def random_problem_2d_chain(): graph = ProblemGraph(robot) solver = RiemannianSolver(graph) - # solver = RiemannianSolver(solver_params) n_tests = 100 - q_init = list_to_variable_dict(n * [0]) G_init = graph.realization(q_init) Y_init = pos_from_graph(G_init) t_sol = [] for idx in range(n_tests): - q_goal = robot.random_configuration() - T_goal = robot.pose(q_goal, f"p{robot.n}") + T_goal = np.asarray(robot.pose(q_goal, f"p{robot.n}")) - G = graph.from_pose(T_goal) - D_goal = distance_matrix_from_graph(G) - omega = adjacency_matrix_from_graph(G) - lb, ub = bound_smoothing(G) - sol_info = solver.solve(D_goal, omega, Y_init =Y_init, use_limits=True, bounds=(lb,ub)) - G_sol = graph_from_pos(sol_info["x"], graph.node_ids) - q_sol = graph.joint_variables(G_sol, {f"p{graph.robot.n}": T_goal}) + sol = solver.solve(T_goal, Y_init=Y_init) + q_sol = sol.q - T_riemannian = robot.pose(q_sol, f"p{robot.n}") - err_riemannian = (T_goal.dot(T_riemannian.inv())).log() - err_riemannian_pos = np.linalg.norm(T_goal.trans - T_riemannian.trans) - err_riemannian_rot = np.linalg.norm(err_riemannian[2]) + T_riemannian = np.asarray(robot.pose(q_sol, f"p{robot.n}")) + err_riemannian_pos = np.linalg.norm(T_goal[:2, 2] - T_riemannian[:2, 2]) + R_delta = T_goal[:2, :2] @ T_riemannian[:2, :2].T + err_riemannian_rot = abs(np.arctan2(R_delta[1, 0], R_delta[0, 0])) - t_sol.append(sol_info['time']) + t_sol.append(sol.time) e_rot.append(err_riemannian_rot) e_pos.append(err_riemannian_pos) q_abs = np.abs(np.array(list(q_sol.values()))) diff --git a/experiments/simple_ik_examples/test_chain_2d_new.py b/experiments/simple_ik_examples/test_chain_2d_new.py index 0e6a832..e1de26d 100644 --- a/experiments/simple_ik_examples/test_chain_2d_new.py +++ b/experiments/simple_ik_examples/test_chain_2d_new.py @@ -1,25 +1,15 @@ #!/usr/bin/env python3 import numpy as np -from numpy.testing import assert_array_less -import networkx as nx -import time + from graphik.graphs import ProblemGraph from graphik.robots import Robot -from graphik.utils import ( - pos_from_graph, - list_to_variable_dict, - distance_matrix_from_graph, - adjacency_matrix_from_graph, - graph_from_pos - ) -from graphik.solvers.riemannian_solver import RiemannianSolver -from graphik.utils.dgp import bound_smoothing +from graphik.solvers import RiemannianSolver +from graphik.utils import list_to_variable_dict, pos_from_graph def random_problem_2d_chain(): e_rot = [] e_pos = [] - t_sol = [] fails = 0 n = 10 @@ -37,34 +27,27 @@ def random_problem_2d_chain(): robot = Robot({**robot_params, "dim": 2}) graph = ProblemGraph(robot) - solver = RiemannianSolver(graph) + solver = RiemannianSolver(graph, use_limits=False) n_tests = 100 - q_init = list_to_variable_dict(n * [0]) G_init = graph.realization(q_init) Y_init = pos_from_graph(G_init) t_sol = [] for idx in range(n_tests): - q_goal = robot.random_configuration() - T_goal = robot.pose(q_goal, f"p{robot.n}") + T_goal = np.asarray(robot.pose(q_goal, f"p{robot.n}")) - G = graph.from_pose(T_goal) - D_goal = distance_matrix_from_graph(G) - omega = adjacency_matrix_from_graph(G) - lb, ub = bound_smoothing(G) - sol_info = solver.solve(D_goal, omega, Y_init =Y_init, bounds=(lb,ub)) - G_sol = graph_from_pos(sol_info["x"], graph.node_ids) - q_sol = graph.joint_variables(G_sol, {f"p{graph.robot.n}": T_goal}) + sol = solver.solve(T_goal, Y_init=Y_init) + q_sol = sol.q - T_riemannian = robot.pose(q_sol, f"p{robot.n}") - err_riemannian = (T_goal.dot(T_riemannian.inv())).log() - err_riemannian_pos = np.linalg.norm(T_goal.trans - T_riemannian.trans) - err_riemannian_rot = np.linalg.norm(err_riemannian[2]) + T_riemannian = np.asarray(robot.pose(q_sol, f"p{robot.n}")) + err_riemannian_pos = np.linalg.norm(T_goal[:2, 2] - T_riemannian[:2, 2]) + R_delta = T_goal[:2, :2] @ T_riemannian[:2, :2].T + err_riemannian_rot = abs(np.arctan2(R_delta[1, 0], R_delta[0, 0])) - t_sol.append(sol_info['time']) + t_sol.append(sol.time) e_rot.append(err_riemannian_rot) e_pos.append(err_riemannian_pos) if err_riemannian_pos > 0.01 or err_riemannian_rot > 0.01: @@ -77,8 +60,8 @@ def random_problem_2d_chain(): print("Average pos error {:}".format(np.average(np.array(e_pos)))) print("Average rot error {:}".format(np.average(np.array(e_rot)))) print("Number of fails {:}".format(fails)) - - assert_array_less(e_pos, 1e-4 * np.ones(n_tests)) + if fails: + raise SystemExit(f"{fails}/{n_tests} IK solves exceeded tolerance") if __name__ == "__main__": diff --git a/experiments/solver_benchmark.py b/experiments/solver_benchmark.py index 26b8908..830d30c 100644 --- a/experiments/solver_benchmark.py +++ b/experiments/solver_benchmark.py @@ -3,12 +3,12 @@ Compares wall time and pose-recovery accuracy across: - Riemannian (RTR) with spectral init from the partial-Gram block - Riemannian (RTR) with the legacy bound-smoothing + MDS init +- Riemannian (RTR) with zero-configuration init - BFGS — scipy.optimize.minimize, method='BFGS' - L-BFGS-B — scipy.optimize.minimize, method='L-BFGS-B' with anchor pinning -Wall time includes any per-pose initialization the solver performs (e.g. -bound-smoothing for the legacy init or spectral eigendecomposition for the -new one). +Wall time is measured around each solve() call, which internally includes +goal-graph construction, per-pose initialization, optimization, and decoding. Usage: python experiments/solver_benchmark.py [--robots ur10,kuka,...] [--n-poses N] @@ -22,13 +22,8 @@ import numpy as np +from graphik.solvers import RiemannianSolver, ScipySolver from graphik.utils.utils import table_environment -from graphik.utils import ( - distance_matrix_from_graph, - adjacency_matrix_from_graph, - bound_smoothing, - graph_from_pos, -) from graphik.utils.roboturdf import ( load_ur10, load_kuka, @@ -36,8 +31,6 @@ load_schunk_lwa4d, load_schunk_lwa4p, ) -from graphik.solvers.riemannian_solver import RiemannianSolver -from graphik.solvers.nonlinear_solver import NonlinearSolver ROBOTS = { @@ -64,66 +57,32 @@ def pose_error(T_goal: np.ndarray, T_got: np.ndarray) -> tuple[float, float]: return pos, float(np.arccos(cos_theta)) -# Each runner takes the problem and returns the recovered Y (n_nodes × dim). -# Wall time is measured around each runner so pre-processing (init) is included. - -def run_rtr_spectral(graph, D_goal, omega, G_partial): - solver = RiemannianSolver(graph, jit=False, init="spectral") - return solver.solve(D_goal, omega, use_limits=True)["x"] - - -def run_rtr_bsmooth(graph, D_goal, omega, G_partial): - lb, ub = bound_smoothing(G_partial) - solver = RiemannianSolver(graph, jit=False, init="bsmooth") - return solver.solve(D_goal, omega, use_limits=True, bounds=(lb, ub))["x"] - - -def run_bfgs(graph, D_goal, omega, G_partial): - lb, ub = bound_smoothing(G_partial) - solver = NonlinearSolver(graph) - # gtol=1e-8 to match RTR's stopping tolerance. scipy's BFGS does not - # accept ftol, so we don't pass one. - return solver.solve( - D_goal, omega, use_limits=True, bounds=(lb, ub), method="BFGS", - options={"gtol": 1e-8}, - )["x"] - - -def run_lbfgsb(graph, D_goal, omega, G_partial): - lb, ub = bound_smoothing(G_partial) - solver = NonlinearSolver(graph) - return solver.solve( - D_goal, omega, use_limits=True, bounds=(lb, ub), method="L-BFGS-B", - options={"gtol": 1e-8, "ftol": 0}, - )["x"] - - +# Each config maps a label to a solver factory. Solvers are constructed once +# per robot and reused across poses. CONFIGS = [ - ("rtr-spectral", run_rtr_spectral), - ("rtr-bsmooth", run_rtr_bsmooth), - ("bfgs", run_bfgs), - ("l-bfgs-b", run_lbfgsb), + ("rtr-spectral", lambda g: RiemannianSolver(g, init="spectral", precon="gn")), + ("rtr-bsmooth", lambda g: RiemannianSolver(g, init="bsmooth", precon="gn")), + ("rtr-zero", lambda g: RiemannianSolver(g, init="zero", precon="gn")), + ("bfgs", lambda g: ScipySolver(g, method="BFGS", options={"gtol": 1e-8})), + ( + "l-bfgs-b", + lambda g: ScipySolver( + g, method="L-BFGS-B", options={"gtol": 1e-8, "ftol": 0} + ), + ), ] -def _eval(robot, graph, T_goal, ee, runner) -> Result: - G_partial = graph.from_pose(T_goal) - D_goal = distance_matrix_from_graph(G_partial) - omega = adjacency_matrix_from_graph(G_partial) +def _eval(robot, solver, T_goal, ee) -> Result: t0 = time.perf_counter() try: - Y_sol = runner(graph, D_goal, omega, G_partial) + res = solver.solve(T_goal) except Exception: return Result(time.perf_counter() - t0, np.nan, np.nan, feasible=False) wall = time.perf_counter() - t0 - G_sol = graph_from_pos(Y_sol, graph.node_ids) - try: - q_sol = graph.joint_variables(G_sol, {ee: T_goal}) - except Exception: - return Result(wall, np.nan, np.nan, feasible=False) - T_got = robot.pose(q_sol, ee) + T_got = robot.pose(res.q, ee) pos_err, rot_err = pose_error(T_goal, np.asarray(T_got)) - feasible = pos_err < 1e-2 and rot_err < 1e-2 + feasible = res.feasible and pos_err < 1e-2 and rot_err < 1e-2 return Result(wall, pos_err, rot_err, feasible) @@ -137,10 +96,11 @@ def run_robot(robot_name, n_poses, obstacles, seed): poses = [np.asarray(robot.pose(robot.random_configuration(), ee)) for _ in range(n_poses)] + solvers = {label: factory(graph) for label, factory in CONFIGS} results: dict[str, list[Result]] = {label: [] for label, _ in CONFIGS} for T_goal in poses: - for label, runner in CONFIGS: - results[label].append(_eval(robot, graph, T_goal, ee, runner)) + for label, _ in CONFIGS: + results[label].append(_eval(robot, solvers[label], T_goal, ee)) return results diff --git a/graphik/solvers/__init__.py b/graphik/solvers/__init__.py index e69de29..71485de 100644 --- a/graphik/solvers/__init__.py +++ b/graphik/solvers/__init__.py @@ -0,0 +1,12 @@ +from graphik.solvers.base import IKResult, IKSolver +from graphik.solvers.joint_angle import JointAngleSolver +from graphik.solvers.riemannian import RiemannianSolver +from graphik.solvers.scipy_solver import ScipySolver + +__all__ = [ + "IKResult", + "IKSolver", + "JointAngleSolver", + "RiemannianSolver", + "ScipySolver", +] diff --git a/graphik/solvers/base.py b/graphik/solvers/base.py new file mode 100644 index 0000000..68c142a --- /dev/null +++ b/graphik/solvers/base.py @@ -0,0 +1,57 @@ +"""Shared contract for the local (non-SDP) IK solvers. + +Every solver constructs with a ProblemGraph and exposes +``solve(T_goal, **per_problem) -> IKResult``. ``T_goal`` is a single +homogeneous transform for the primary end effector, or a +``{node: transform}`` dict accepted by ``ProblemGraph.from_pose``. +""" +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass + +import numpy as np + +from graphik.graphs.graph import ProblemGraph + +LIMIT_TOL = 1e-6 + + +@dataclass +class IKResult: + q: dict + cost: float + iterations: int + time: float + status: str + limit_violations: list + Y: np.ndarray | None = None + + @property + def feasible(self) -> bool: + return not self.limit_violations + + +class IKSolver(ABC): + """Problem-level IK solver: construct once per graph, solve per goal.""" + + def __init__(self, graph: ProblemGraph): + self.graph = graph + self.robot = graph.robot + self.dim = graph.dim + self.N = graph.number_of_nodes() + + def goals_from(self, T_goal) -> dict: + """Normalize ``T_goal`` to a ``{node: transform}`` mapping.""" + if isinstance(T_goal, dict): + return dict(T_goal) + return {self.robot.end_effectors[0]: T_goal} + + def check_limits(self, q: dict) -> list: + return self.graph.check_distance_limits( + self.graph.realization(q), tol=LIMIT_TOL + ) + + @abstractmethod + def solve(self, T_goal, **kwargs) -> IKResult: + ... diff --git a/graphik/solvers/loss.py b/graphik/solvers/costs.py similarity index 65% rename from graphik/solvers/loss.py rename to graphik/solvers/costs.py index bb70f03..d3d65d1 100644 --- a/graphik/solvers/loss.py +++ b/graphik/solvers/costs.py @@ -1,4 +1,4 @@ -"""Single-source EDM loss kernels for the non-SDP IK solvers. +"""Single-source cost kernels for the non-SDP IK solvers. NumPy dense backend: ``_dense_equality`` and ``_dense_limits`` build an S matrix once per Y and reuse it across cost/egrad/ehvp via ``memoize_last``. @@ -17,12 +17,15 @@ from typing import Any import numpy as np +from pymlg.numpy import SE2, SE3, SO2, SO3 from graphik.utils import ( distance_matrix_from_gram, distance_matrix_from_pos, memoize_last, ) +from graphik.utils.constants import LOWER, POS, ROOT +from graphik.utils.utils import list_to_variable_dict def _dense_equality(D_goal, omega, cache=True): @@ -178,3 +181,90 @@ def hessp(Y_flat, Z_flat): return ehvp(Y, Z_flat.reshape(-1, dim)).ravel() return cost_and_grad, hessp + + +def pose_cost(robot, point: str, T_goal): + """(cost, grad) closure for the squared SE(n) log pose error at ``point``. + + q is read positionally along the root-to-``point`` kinematic chain; the + returned gradient is padded to the robot's full joint count by + ``robot.jacobian``, so per-goal gradients simply sum. + """ + joints = robot.kinematic_map[ROOT][point][1:] + n = len(joints) + if robot.dim == 3: + log = SE3.Log + inverse = SE3.inverse + adjoint = SE3.adjoint + inv_left_jacobian = SE3.left_jacobian_inv + else: + log = SE2.Log + inverse = SE2.inverse + adjoint = SE2.adjoint + inv_left_jacobian = SE2.left_jacobian_inv + + def cost(q): + q_dict = {joints[idx]: q[idx] for idx in range(n)} + T = robot.pose(q_dict, point) + T_inv = inverse(T) + J = robot.jacobian(q_dict, [point]) + e = log(T_inv @ T_goal).ravel() + J_e = inv_left_jacobian(e) + J[point] = J_e @ adjoint(T_inv) @ J[point] + jac = -2 * J[point].T @ e + return e.T @ e, jac + + return cost + + +def obstacle_constraints(robot, graph, pairs: list): + """SLSQP inequality function: ||c_obs - p_node||^2 - r^2 >= 0 per pair.""" + def obstacle_constraint(q): + q_dict = list_to_variable_dict(q) + T_all = robot.get_all_poses(q_dict) + + constr = [] + for robot_node, obs_node in pairs: + p = T_all[robot_node][:-1, -1] + r = graph[robot_node][obs_node][LOWER] + c = graph.nodes[obs_node][POS] + constr += [(c - p).T @ (c - p) - r ** 2] + return np.asarray(constr) + + return obstacle_constraint + + +def obstacle_constraint_gradient(robot, graph, pairs: list): + """Jacobian of ``obstacle_constraints`` w.r.t. q, one row per pair.""" + if robot.dim == 3: + dim = 3 + ZZ = np.zeros([6, 6]) + ZZ[:3, 3:] = np.eye(3) + ZZ[3:, :3] = np.eye(3) + wedge = SO3.wedge + inverse_se = SE3.inverse + else: + dim = 2 + ZZ = np.zeros([4, 4]) + ZZ[:2, 2:] = np.eye(2) + ZZ[2:, :2] = np.eye(2) + wedge = SO2.wedge + inverse_se = SE2.inverse + + def obstacle_gradient(q): + q_dict = list_to_variable_dict(q) + T_all = robot.get_all_poses(q_dict) + J_all = robot.jacobian(q_dict, list(q_dict.keys())) + + jac = [] + for robot_node, obs_node in pairs: + T_node = T_all[robot_node] + R = T_node[:dim, :dim] + t_inv = inverse_se(T_node)[:dim, -1] + ZZ[:dim, :dim] = R @ wedge(t_inv) @ R.T + p = T_node[:dim, -1] + c = graph.nodes[obs_node][POS] + jac += [-2 * (c - p).T @ (ZZ @ J_all[robot_node])[:dim, :]] + return np.vstack(jac) + + return obstacle_gradient diff --git a/graphik/solvers/distance_solver.py b/graphik/solvers/distance_solver.py new file mode 100644 index 0000000..a0aad22 --- /dev/null +++ b/graphik/solvers/distance_solver.py @@ -0,0 +1,113 @@ +"""Shared problem-level plumbing for the EDM-based solvers.""" +from __future__ import annotations + +import time +from abc import abstractmethod +from dataclasses import dataclass +from typing import Optional, Tuple + +import networkx as nx +import numpy as np + +from graphik.graphs.graph import ProblemGraph +from graphik.solvers.base import IKResult, IKSolver +from graphik.solvers.initializations import ( + INIT_STRATEGIES, + bsmooth_init, + spectral_init, + zero_init, +) +from graphik.utils import graph_from_pos +from graphik.utils.dgp import adjacency_matrix_from_graph, distance_matrix_from_graph + + +@dataclass +class DistanceProblem: + """Per-solve state handed to ``_minimize``.""" + + G: nx.DiGraph + D_goal: np.ndarray + omega: np.ndarray + psi_L: Optional[np.ndarray] + psi_U: Optional[np.ndarray] + Y0: np.ndarray + + +@dataclass +class MinimizeInfo: + """What ``_minimize`` reports back for ``IKResult``.""" + + cost: float + iterations: int + status: str + + +class DistanceSolver(IKSolver): + def __init__( + self, + graph: ProblemGraph, + *, + init: str = "bsmooth", + use_limits: bool = True, + cache: bool = True, + ): + super().__init__(graph) + if init not in INIT_STRATEGIES: + raise ValueError( + "init must be one of " + + ", ".join(repr(s) for s in INIT_STRATEGIES) + ) + self.init = init + self.use_limits = use_limits + self.cache = cache + + def generate_initialization(self, G, D_goal, omega) -> np.ndarray: + if self.init == "spectral": + return spectral_init(D_goal, omega, self.dim) + if self.init == "zero": + return zero_init(self.graph) + return bsmooth_init(G, omega, self.dim) + + @abstractmethod + def _minimize(self, problem: DistanceProblem) -> Tuple[np.ndarray, MinimizeInfo]: + ... + + def solve(self, T_goal, *, Y_init=None) -> IKResult: + t0 = time.perf_counter() + goals = self.goals_from(T_goal) + G = self.graph.from_pose(goals) + D_goal = distance_matrix_from_graph(G) + omega = adjacency_matrix_from_graph(G) + + psi_L = psi_U = None + if self.use_limits: + psi_L, psi_U = self.graph.distance_bound_matrices() + + Y0 = ( + Y_init + if Y_init is not None + else self.generate_initialization(G, D_goal, omega) + ) + + Y, info = self._minimize( + DistanceProblem( + G=G, + D_goal=D_goal, + omega=omega, + psi_L=psi_L, + psi_U=psi_U, + Y0=Y0, + ) + ) + + G_sol = graph_from_pos(Y, self.graph.node_ids) + q = self.graph.joint_variables(G_sol, goals) + return IKResult( + q=q, + cost=info.cost, + iterations=info.iterations, + time=time.perf_counter() - t0, + status=info.status, + limit_violations=self.check_limits(q), + Y=Y, + ) diff --git a/graphik/solvers/initializations.py b/graphik/solvers/initializations.py new file mode 100644 index 0000000..8841777 --- /dev/null +++ b/graphik/solvers/initializations.py @@ -0,0 +1,49 @@ +"""Shared Y-initialization strategies for distance-geometric solvers. + +Each returns an ``(N, dim)`` point-cloud factor ordered consistently with the +problem graph's node order. ``zero_init`` is explicitly ordered by +``graph.node_ids``; the others inherit ordering from ``D_goal`` / ``omega``. +""" +from __future__ import annotations + +import numpy as np + +from graphik.utils import MDS, POS, gram_from_distance_matrix, linear_projection +from graphik.utils.dgp import bound_smoothing + +INIT_STRATEGIES = ("spectral", "bsmooth", "zero") + + +def spectral_init(D_goal, omega, dim): + """Partial-Gram spectral initialization with an eigenvalue floor.""" + G = gram_from_distance_matrix(omega * D_goal) + eigvals, eigvecs = np.linalg.eigh(G) + top = np.argsort(-eigvals)[:dim] + lam = np.maximum(eigvals[top], 1e-12) + return eigvecs[:, top] * np.sqrt(lam) + + +def _sample_bounds_init(lb, ub, omega, dim): + """Sample an EDM at 0.9 of (lb, ub), MDS, then project to ``dim``.""" + lb_sqrt = np.sqrt(lb) + ub_sqrt = np.sqrt(ub) + D_rand = (lb_sqrt + 0.9 * (ub_sqrt - lb_sqrt)) ** 2 + X_rand = MDS(gram_from_distance_matrix(D_rand), eps=1e-8) + return linear_projection(X_rand, omega, dim) + + +def bsmooth_init(G, omega, dim): + """Bound smoothing plus the legacy sampled-EDM initializer.""" + lb, ub = bound_smoothing(G) + return _sample_bounds_init(lb, ub, omega, dim) + + +def zero_init(graph): + """Use the graph realization at ``robot.zero_configuration()``.""" + G_zero = graph.realization(graph.robot.zero_configuration()) + return np.stack( + [ + np.asarray(G_zero.nodes[node][POS], dtype=float) + for node in graph.node_ids + ] + ) diff --git a/graphik/solvers/joint_angle.py b/graphik/solvers/joint_angle.py new file mode 100644 index 0000000..abcc333 --- /dev/null +++ b/graphik/solvers/joint_angle.py @@ -0,0 +1,80 @@ +"""Joint-space IK solver using SLSQP on SE(n) log pose error.""" +from __future__ import annotations + +import time + +import networkx as nx +import numpy as np +from scipy.optimize import minimize + +from graphik.graphs.graph import ProblemGraph +from graphik.solvers import costs +from graphik.solvers.base import IKResult, IKSolver +from graphik.utils.constants import BELOW, BOUNDED, OBSTACLE, ROBOT, ROOT, TYPE + + +class JointAngleSolver(IKSolver): + def __init__(self, graph: ProblemGraph): + super().__init__(graph) + self.n = self.robot.n + self.joint_order = [joint for joint in self.robot.joint_ids if joint != ROOT] + + typ = nx.get_node_attributes(self.graph, name=TYPE) + pairs = [] + for u, v, data in self.graph.edges(data=True): + if BELOW in data.get(BOUNDED, []): + if ROBOT in typ[u] and OBSTACLE in typ[v] and u != ROOT: + pairs.append((u, v)) + + self.constraints = [] + if pairs: + self.constraints = [ + { + "type": "ineq", + "fun": costs.obstacle_constraints(self.robot, self.graph, pairs), + "jac": costs.obstacle_constraint_gradient( + self.robot, self.graph, pairs + ), + } + ] + + def solve(self, T_goal, *, q_init=None) -> IKResult: + t0 = time.perf_counter() + goals = self.goals_from(T_goal) + if q_init is None: + q_init = self.robot.zero_configuration() + + goal_costs = [ + costs.pose_cost(self.robot, node, goal) for node, goal in goals.items() + ] + + def cost_and_grad(q): + f = 0.0 + grad = np.zeros(self.n) + for cg in goal_costs: + f_i, g_i = cg(q) + f += f_i + grad += g_i + return f, grad + + res = minimize( + cost_and_grad, + np.asarray([q_init[joint] for joint in self.joint_order]), + jac=True, + constraints=self.constraints, + method="SLSQP", + options={"ftol": 1e-7}, + ) + q = { + joint: value + for joint, value in zip(self.joint_order, res.x) + } + return IKResult( + q=q, + cost=float(res.fun), + iterations=int(res.nit), + time=time.perf_counter() - t0, + status=str(res.message), + limit_violations=self.check_limits(q), + Y=None, + ) diff --git a/graphik/solvers/joint_angle_solver.py b/graphik/solvers/joint_angle_solver.py deleted file mode 100644 index 03b1cd3..0000000 --- a/graphik/solvers/joint_angle_solver.py +++ /dev/null @@ -1,135 +0,0 @@ -import numpy as np -import networkx as nx -from typing import Dict, Any -from scipy.optimize import minimize -from pymlg.numpy import SE3, SO3, SE2, SO2 -from graphik.utils.constants import * -from graphik.graphs.graph import ProblemGraph -from graphik.utils.utils import list_to_variable_dict - - -class LocalSolver: - def __init__(self, robot_graph: ProblemGraph, params: Dict["str", Any]): - self.graph = robot_graph - self.robot = robot_graph.robot - self.k_map = self.robot.kinematic_map[ROOT] # get map to all nodes from root - self.n = self.robot.n - self.dim = self.graph.dim - - # create obstacle constraints - typ = nx.get_node_attributes(self.graph, name=TYPE) - pairs = [] - for u, v, data in self.graph.edges(data=True): - if BELOW in data.get(BOUNDED, []): - if ROBOT in typ[u] and OBSTACLE in typ[v] and u != ROOT: - pairs += [(u, v)] - self.m = len(pairs) - self.g = [] - if len(pairs) > 0: - fun = self.gen_obstacle_constraints(pairs) - jac = self.gen_obstacle_constraint_gradient(pairs) - self.g = [{"type": "ineq", "fun": fun, "jac": jac}] - - def gen_obstacle_constraints(self, pairs: list): - def obstacle_constraint(q): - q_dict = list_to_variable_dict(q) - T_all = self.robot.get_all_poses(q_dict) - - constr = [] - for robot_node, obs_node in pairs: - p = T_all[robot_node][:-1, -1] - r = self.graph[robot_node][obs_node][LOWER] - c = self.graph.nodes[obs_node][POS] - constr += [(c - p).T @ (c - p) - r ** 2] - return np.asarray(constr) - - return obstacle_constraint - - def gen_obstacle_constraint_gradient(self, pairs: list): - if self.dim == 3: - dim = 3 - ZZ = np.zeros([6, 6]) - ZZ[:3, 3:] = np.eye(3) - ZZ[3:, :3] = np.eye(3) - wedge = SO3.wedge - inverse_se = SE3.inverse - else: - dim = 2 - ZZ = np.zeros([4, 4]) - ZZ[:2, 2:] = np.eye(2) - ZZ[2:, :2] = np.eye(2) - wedge = SO2.wedge - inverse_se = SE2.inverse - - def obstacle_gradient(q): - q_dict = list_to_variable_dict(q) - T_all = self.robot.get_all_poses(q_dict) - J_all = self.robot.jacobian(q_dict, list(q_dict.keys())) - - jac = [] - for robot_node, obs_node in pairs: - T_node = T_all[robot_node] - R = T_node[:dim, :dim] - t_inv = inverse_se(T_node)[:dim, -1] - ZZ[:dim, :dim] = R @ wedge(t_inv) @ R.T - p = T_node[:dim, -1] - c = self.graph.nodes[obs_node][POS] - jac += [-2 * (c - p).T @ (ZZ @ J_all[robot_node])[:dim, :]] - return np.vstack(jac) - - return obstacle_gradient - - def gen_cost_and_grad_ee(self, point: str, T_goal: np.ndarray): - joints = self.k_map[point][1:] - n = len(joints) - if self.dim == 3: - log = SE3.Log - inverse = SE3.inverse - adjoint = SE3.adjoint - inv_left_jacobian = SE3.left_jacobian_inv - else: - log = SE2.Log - inverse = SE2.inverse - adjoint = SE2.adjoint - inv_left_jacobian = SE2.left_jacobian_inv - - def cost(q): - q_dict = {joints[idx]: q[idx] for idx in range(n)} - T = self.robot.pose(q_dict, point) - T_inv = inverse(T) - J = self.robot.jacobian(q_dict, [point]) - e = log(T_inv @ T_goal).ravel() - J_e = inv_left_jacobian(e) - J[point] = J_e @ adjoint(T_inv) @ J[point] - jac = -2 * J[point].T @ e - return e.T @ e, jac - - return cost - - def solve(self, goals: dict, q0: dict): - # Each per-goal cost reads the leading entries of q along its - # root-to-goal chain (q0 must be ordered along the kinematic chain) - # and returns a full-width gradient (robot.jacobian pads with zero - # columns past the chain), so goals simply sum. - goal_costs = [ - self.gen_cost_and_grad_ee(node, goal) for node, goal in goals.items() - ] - - def cost_and_grad(q): - f = 0.0 - grad = np.zeros(self.robot.n) - for cg in goal_costs: - f_i, g_i = cg(q) - f += f_i - grad += g_i - return f, grad - - res = minimize( - cost_and_grad, - np.asarray(list(q0.values())), - jac=True, - constraints=self.g, - method="SLSQP", - options={"ftol": 1e-7}, - ) - return res diff --git a/graphik/solvers/nonlinear_solver.py b/graphik/solvers/nonlinear_solver.py deleted file mode 100644 index 15677c6..0000000 --- a/graphik/solvers/nonlinear_solver.py +++ /dev/null @@ -1,122 +0,0 @@ -#!/usr/bin/env python3 -import time -import numpy as np - -from graphik.utils import ( - MDS, - linear_projection, - gram_from_distance_matrix, -) -from graphik.graphs.graph import ProblemGraph -from scipy.optimize import Bounds, minimize -from graphik.utils.constants import POS -from graphik.solvers import loss - -# scipy.optimize.minimize methods that consume hessp. Passing the kwarg to -# other methods just provokes a UserWarning per call. -_HESSP_METHODS = frozenset({ - "Newton-CG", "trust-ncg", "trust-krylov", "trust-constr", -}) - - -class NonlinearSolver: - def __init__(self, graph: ProblemGraph, *args, **kwargs): - """Distance-based IK solver wrapping ``scipy.optimize.minimize``. - - Cost / gradient / HVP come from ``graphik.solvers.loss.for_minimize`` - (dense NumPy backend). - """ - for removed in ("cost_type", "jit"): - if removed in kwargs: - raise TypeError( - f"{removed} was removed: the numba/AOT backend is gone " - "and the dense NumPy loss backend is the only backend." - ) - for key in kwargs: - setattr(self, key, kwargs[key]) - self.graph = graph - self.dim = graph.dim - self.N = graph.number_of_nodes() - - def generate_initialization(self, bounds, dim, omega): - """Sample an EDM within the supplied (lb, ub) bounds, then MDS + project.""" - lb = np.sqrt(bounds[0]) - ub = np.sqrt(bounds[1]) - D_rand = (lb + 0.9 * (ub - lb)) ** 2 - X_rand = MDS(gram_from_distance_matrix(D_rand), eps=1e-8) - return linear_projection(X_rand, omega, dim) - - def create_cost(self, D_goal, omega): - return loss.for_minimize(D_goal, omega, dim=self.dim) - - def create_cost_limits(self, D_goal, omega, psi_L, psi_U): - return loss.for_minimize( - D_goal, omega, psi_L=psi_L, psi_U=psi_U, dim=self.dim, - ) - - def position_constraints(self): - """Pin POS-tagged nodes to their stored goal positions. - - Used by L-BFGS-B to enforce known positions (end-effector goal, base - anchor, axis-frame nodes) as hard equality constraints rather than as - soft residuals. Reads positions from the graph; do not pass a random - initialization here. - """ - ub_ = np.ones((self.N, self.dim)) * np.inf - lb_ = -np.ones((self.N, self.dim)) * np.inf - for idx, (node, data) in enumerate(self.graph.nodes(data=True)): - if POS in data: - ub_[idx] = data[POS] - lb_[idx] = data[POS] - return Bounds(lb=lb_.flatten(), ub=ub_.flatten()) - - def solve( - self, - D_goal, - omega, - use_limits=False, - bounds=None, - Y_init=None, - output_log=True, - method='BFGS', - options=None, - ): - """Run scipy.optimize.minimize on the EDM loss.""" - if use_limits: - psi_L, psi_U = self.graph.distance_bound_matrices() - cost_and_grad, hessp = self.create_cost_limits(D_goal, omega, psi_L, psi_U) - else: - cost_and_grad, hessp = self.create_cost(D_goal, omega) - - if Y_init is None: - Y_init = self.generate_initialization(bounds, self.dim, omega) - Yi = np.ascontiguousarray(Y_init.flatten()) - - bnds = None - if method == 'L-BFGS-B': - bnds = self.position_constraints() - defaults = {"ftol": 1e-16, "gtol": 1e-16, "iprint": -1} - elif method == 'BFGS': - defaults = {"xrtol": 0.75e-6, "gtol": 0.25e-6, "norm": np.inf} - else: - defaults = {} - options = {**defaults, **(options or {})} - - minimize_kwargs = dict( - jac=True, method=method, bounds=bnds, options=options, - ) - if hessp is not None and method in _HESSP_METHODS: - minimize_kwargs["hessp"] = hessp - - start_time = time.time() - sol = minimize(cost_and_grad, Yi, **minimize_kwargs) - end_time = time.time() - - if output_log: - return { - "x": sol.x.reshape(omega.shape[0], self.dim), - "time": end_time - start_time, - "iterations": sol.nit, - "f(x)": sol.fun, - } - return sol.x diff --git a/graphik/solvers/riemannian.py b/graphik/solvers/riemannian.py new file mode 100644 index 0000000..a0ed093 --- /dev/null +++ b/graphik/solvers/riemannian.py @@ -0,0 +1,171 @@ +"""Riemannian trust-region EDM solver and tCG preconditioners.""" +from __future__ import annotations + +import numpy as np + +from graphik.graphs.graph import ProblemGraph +from graphik.solvers import costs, rtr +from graphik.solvers.distance_solver import ( + DistanceProblem, + DistanceSolver, + MinimizeInfo, +) +from graphik.utils import distance_matrix_from_pos, memoize_last +from graphik.utils.manifolds.fixed_rank_psd_sym import PSDFixedRank + + +def _gn_weights(omega, psi_L, psi_U): + """Active-edge weight matrix for the Gauss-Newton model: equality edges + plus whichever bound edges are active at the current distances.""" + if psi_L is None: + return lambda D: omega + diff = psi_L != psi_U + LL = diff * (psi_L > 0) + UU = diff * (psi_U > 0) + + def weights(D): + return omega + LL * (D < psi_L) + UU * (D > psi_U) + + return weights + + +def _gn_blocks(Y, W): + """Per-edge Gauss-Newton blocks of the EDM loss at Y. + + The cost is sum over (undirected, active) edges of + (w_ij * (d_goal^2 - ||y_i - y_j||^2))^2, whose GN Hessian has + graph-Laplacian block structure with edge blocks 8 w_ij u u^T, + u = y_i - y_j. Returns the (N, N, d, d) edge-block tensor E. + """ + Y_diff = Y[:, None, :] - Y[None, :, :] + return (8.0 * W)[:, :, None, None] * ( + Y_diff[..., :, None] * Y_diff[..., None, :] + ) + + +def make_gn_preconditioner(manifold, omega, psi_L=None, psi_U=None, floor_rel=1e-3): + """Inverse eigenvalue-floored Gauss-Newton Hessian preconditioner.""" + weights = _gn_weights(omega, psi_L, psi_U) + + @memoize_last + def factor(Y): + N, d = Y.shape + D = distance_matrix_from_pos(Y) + E = _gn_blocks(Y, weights(D)) + B = -E + idx = np.arange(N) + B[idx, idx] = E.sum(axis=1) + H = B.transpose(0, 2, 1, 3).reshape(N * d, N * d) + lam, V = np.linalg.eigh(H) + lam = np.maximum(lam, max(floor_rel * lam[-1], 1e-12)) + return V, lam + + def precondition(Y, r): + V, lam = factor(Y) + z = (V @ ((V.T @ r.ravel()) / lam)).reshape(r.shape) + return manifold.projection(Y, z) + + return precondition + + +def make_jacobi_preconditioner( + manifold, omega, psi_L=None, psi_U=None, floor_rel=1e-3 +): + """Block-Jacobi variant of ``make_gn_preconditioner``.""" + weights = _gn_weights(omega, psi_L, psi_U) + + @memoize_last + def factor(Y): + D = distance_matrix_from_pos(Y) + E = _gn_blocks(Y, weights(D)) + Bd = E.sum(axis=1) + lam, V = np.linalg.eigh(Bd) + lam = np.maximum(lam, np.maximum(floor_rel * lam[:, -1:], 1e-12)) + return V, lam + + def precondition(Y, r): + V, lam = factor(Y) + rv = np.einsum("nij,ni->nj", V, r) + z = np.einsum("nij,nj->ni", V, rv / lam) + return manifold.projection(Y, z) + + return precondition + + +_RTR_PASSTHROUGH = ( + "rho_prime", + "rho_regularization", + "Delta_bar", + "Delta0", + "mininner", + "maxinner", +) + + +class RiemannianSolver(DistanceSolver): + """EDM solver on the PSD fixed-rank manifold via in-house RTR.""" + + def __init__( + self, + graph: ProblemGraph, + *, + init: str = "bsmooth", + use_limits: bool = True, + cache: bool = True, + precon=None, + rtr_params: dict | None = None, + ): + super().__init__(graph, init=init, use_limits=use_limits, cache=cache) + if not (precon in (None, "gn", "jacobi") or callable(precon)): + raise ValueError("precon must be None, 'gn', 'jacobi', or callable") + self.precon = precon + self.rtr_params = dict(rtr_params or {}) + + def _minimize(self, problem: DistanceProblem): + manifold = PSDFixedRank(self.N, self.dim, cache_projection=self.cache) + cost, egrad, ehvp = costs.for_riemannian( + problem.D_goal, + problem.omega, + psi_L=problem.psi_L, + psi_U=problem.psi_U, + cache=self.cache, + ) + + if self.precon is None: + preconditioner = None + elif self.precon == "gn": + preconditioner = make_gn_preconditioner( + manifold, problem.omega, psi_L=problem.psi_L, psi_U=problem.psi_U + ) + elif self.precon == "jacobi": + preconditioner = make_jacobi_preconditioner( + manifold, problem.omega, psi_L=problem.psi_L, psi_U=problem.psi_U + ) + else: + preconditioner = self.precon + + def rgrad(Y): + return manifold.projection(Y, egrad(Y)) + + def rhess(Y, U): + return manifold.projection(Y, ehvp(Y, U)) + + rtr_kwargs = dict( + max_iterations=self.rtr_params.get("maxiter", 3000), + min_gradient_norm=self.rtr_params.get("mingradnorm", 1e-8), + theta=self.rtr_params.get("theta", 1.0), + kappa=self.rtr_params.get("kappa", 0.1), + preconditioner=preconditioner, + ) + for key in _RTR_PASSTHROUGH: + if key in self.rtr_params: + rtr_kwargs[key] = self.rtr_params[key] + + res = rtr.trust_regions( + manifold, cost, rgrad, rhess, problem.Y0, **rtr_kwargs + ) + return res.point, MinimizeInfo( + cost=res.cost, + iterations=res.iterations, + status=res.stopping_criterion, + ) diff --git a/graphik/solvers/riemannian_solver.py b/graphik/solvers/riemannian_solver.py deleted file mode 100644 index fb15650..0000000 --- a/graphik/solvers/riemannian_solver.py +++ /dev/null @@ -1,320 +0,0 @@ -#!/usr/bin/env python3 -from graphik.utils.dgp import ( - adjacency_matrix_from_graph, bound_smoothing, - distance_matrix_from_graph, graph_from_pos, -) - -import numpy as np -from graphik.utils import ( - MDS, - linear_projection, - gram_from_distance_matrix, - distance_matrix_from_pos, - memoize_last, - POS, -) -from graphik.utils.manifolds.fixed_rank_psd_sym import PSDFixedRank -from graphik.graphs.graph import ProblemGraph -from graphik.solvers import rtr, loss - - -RIEMANNIAN_INIT_STRATEGIES = ("spectral", "bsmooth", "zero") - - -def _gn_weights(omega, psi_L, psi_U): - """Active-edge weight matrix for the Gauss-Newton model: equality edges - plus whichever bound edges are active at the current distances.""" - if psi_L is None: - return lambda D: omega - diff = psi_L != psi_U - LL = diff * (psi_L > 0) - UU = diff * (psi_U > 0) - - def weights(D): - return omega + LL * (D < psi_L) + UU * (D > psi_U) - - return weights - - -def _gn_blocks(Y, W): - """Per-edge Gauss-Newton blocks of the EDM loss at Y. - - The cost is sum over (undirected, active) edges of - (w_ij * (d_goal² - ||y_i - y_j||²))², whose GN Hessian has graph-Laplacian - block structure with edge blocks 8 w_ij u u^T, u = y_i - y_j. Returns the - (N, N, d, d) edge-block tensor E (E[i, i] = 0). - """ - Y_diff = Y[:, None, :] - Y[None, :, :] - return (8.0 * W)[:, :, None, None] * ( - Y_diff[..., :, None] * Y_diff[..., None, :] - ) - - -def make_gn_preconditioner(manifold, omega, psi_L=None, psi_U=None, - floor_rel=1e-3): - """Preconditioner for RTR's truncated CG: the inverse of the (eigenvalue- - floored) Gauss-Newton Hessian of the EDM loss. - - The GN matrix is N*d x N*d — small for IK-sized graphs — so one eigh per - outer iteration (memoized per Y) buys a near-exact Newton preconditioner. - The flooring keeps it positive definite: the GN matrix is PSD with a - nullspace containing the translation directions. Output is projected back - to the tangent (horizontal) space. - """ - weights = _gn_weights(omega, psi_L, psi_U) - - @memoize_last - def factor(Y): - N, d = Y.shape - D = distance_matrix_from_pos(Y) - E = _gn_blocks(Y, weights(D)) - B = -E - idx = np.arange(N) - B[idx, idx] = E.sum(axis=1) - H = B.transpose(0, 2, 1, 3).reshape(N * d, N * d) - lam, V = np.linalg.eigh(H) - lam = np.maximum(lam, max(floor_rel * lam[-1], 1e-12)) - return V, lam - - def precondition(Y, r): - V, lam = factor(Y) - z = (V @ ((V.T @ r.ravel()) / lam)).reshape(r.shape) - return manifold.projection(Y, z) - - return precondition - - -def make_jacobi_preconditioner(manifold, omega, psi_L=None, psi_U=None, - floor_rel=1e-3): - """Block-Jacobi variant of `make_gn_preconditioner`: keeps only the N - diagonal d x d blocks of the GN Hessian. Cheaper apply, weaker model.""" - weights = _gn_weights(omega, psi_L, psi_U) - - @memoize_last - def factor(Y): - D = distance_matrix_from_pos(Y) - E = _gn_blocks(Y, weights(D)) - Bd = E.sum(axis=1) # (N, d, d) diagonal blocks - lam, V = np.linalg.eigh(Bd) # batched - lam = np.maximum(lam, np.maximum(floor_rel * lam[:, -1:], 1e-12)) - return V, lam - - def precondition(Y, r): - V, lam = factor(Y) - rv = np.einsum("nij,ni->nj", V, r) - z = np.einsum("nij,nj->ni", V, rv / lam) - return manifold.projection(Y, z) - - return precondition - - -class RiemannianSolver: - def __init__( - self, - graph: ProblemGraph, - cache=True, - init="bsmooth", - *args, - **kwargs, - ): - for removed in ("cost_type", "jit"): - if removed in kwargs: - raise TypeError( - f"{removed} was removed: the numba/AOT backend is gone " - "and the dense NumPy loss backend is the only backend." - ) - self.params = {} - for key in kwargs: - setattr(self, key, kwargs[key]) - self.graph = graph - self.dim = graph.dim - self.N = graph.number_of_nodes() - # Single switch that enables both per-Y memoization caches: - # the solver's cost-state builder and the manifold's projection - # operator. Both are perf-only — same algorithmic trajectory. - self.cache = cache - if init not in RIEMANNIAN_INIT_STRATEGIES: - raise ValueError( - "init must be one of " - + ", ".join(repr(strategy) for strategy in RIEMANNIAN_INIT_STRATEGIES) - ) - self.init = init - - def zero_configuration_initialization(self): - """Return an RTR point-cloud factor from the robot zero configuration.""" - G_zero = self.graph.realization(self.graph.robot.zero_configuration()) - return np.stack( - [ - np.asarray(G_zero.nodes[node][POS], dtype=float) - for node in self.graph.node_ids - ] - ) - - def generate_initialization(self, D_goal, omega, bounds=None): - """Build a starting point for RTR. Dispatches on `self.init`: - - - 'spectral': Smith–Cai–Tasissa style. Center the partially- - observed distance matrix (zeros for unknowns), eigendecompose, take - the top-d eigenvectors weighted by sqrt(eigvals). One N×N symmetric - eigendecomposition. - - 'bsmooth' (default): legacy. Triangle-inequality smoothing of distance bounds - (networkx all-pairs Bellman-Ford), sample an EDM at 0.9 of (lb, ub), - classical MDS, then project to dim along the known-edge structure. - Requires a ``bounds=(lb, ub)`` tuple (e.g. from ``bound_smoothing``). - - 'zero': realize the graph at ``robot.zero_configuration()`` and use - the resulting node positions directly, ordered by ``graph.node_ids``. - """ - if self.init == "spectral": - n = D_goal.shape[0] - D_partial = omega * D_goal - J = np.eye(n) - np.full((n, n), 1.0 / n) - G = -0.5 * J @ D_partial @ J - eigvals, eigvecs = np.linalg.eigh(G) # ascending - top = np.argsort(-eigvals)[:self.dim] - # Floor eigenvalues to keep Y full-rank: PSDFixedRank requires - # rank == dim, and degenerate problems can yield fewer than dim - # positive eigvals here, which would otherwise produce zero - # columns and break the manifold projection. - lam = np.maximum(eigvals[top], 1e-12) - return eigvecs[:, top] * np.sqrt(lam) - - if self.init == "zero": - return self.zero_configuration_initialization() - - if bounds is None: - raise ValueError("init='bsmooth' requires bounds=(lb, ub)") - lb, ub = bounds - lb_sqrt = np.sqrt(lb) - ub_sqrt = np.sqrt(ub) - D_rand = (lb_sqrt + 0.9 * (ub_sqrt - lb_sqrt)) ** 2 - X_rand = MDS(gram_from_distance_matrix(D_rand), eps=1e-8) - return linear_projection(X_rand, omega, self.dim) - - def create_cost(self, D_goal, omega): - return loss.for_riemannian(D_goal, omega, cache=self.cache) - - def create_cost_limits(self, D_goal, omega, psi_L, psi_U): - return loss.for_riemannian( - D_goal, omega, psi_L=psi_L, psi_U=psi_U, cache=self.cache, - ) - - def solve( - self, - D_goal, - omega, - use_limits=False, - bounds=None, - Y_init=None, - method=None, - output_log=True, - precon=None, - ): - manifold = PSDFixedRank( - self.N, self.dim, - cache_projection=self.cache, - ) - - psi_L = psi_U = None - if not use_limits: - cost, egrad, ehess = self.create_cost(D_goal, omega) - else: - psi_L, psi_U = self.graph.distance_bound_matrices() - cost, egrad, ehess = self.create_cost_limits( - D_goal, omega, psi_L, psi_U - ) - - if precon is None: - preconditioner = None - elif precon == "gn": - preconditioner = make_gn_preconditioner( - manifold, omega, psi_L=psi_L, psi_U=psi_U - ) - elif precon == "jacobi": - preconditioner = make_jacobi_preconditioner( - manifold, omega, psi_L=psi_L, psi_U=psi_U - ) - elif callable(precon): - preconditioner = precon - else: - raise ValueError("precon must be None, 'gn', 'jacobi', or callable") - - if Y_init is None: - Y_init = self.generate_initialization(D_goal, omega, bounds=bounds) - - def rgrad(Y): - return manifold.projection(Y, egrad(Y)) - - def rhess(Y, U): - return manifold.projection(Y, ehess(Y, U)) - - rtr_kwargs = dict( - max_iterations=self.params.get("maxiter", 3000), - min_gradient_norm=self.params.get("mingradnorm", 1e-8), - theta=self.params.get("theta", 1.0), - kappa=self.params.get("kappa", 0.1), - preconditioner=preconditioner, - ) - for k in ("rho_prime", "rho_regularization", - "Delta_bar", "Delta0", "mininner", "maxinner"): - if k in self.params: - rtr_kwargs[k] = self.params[k] - res = rtr.trust_regions(manifold, cost, rgrad, rhess, Y_init, **rtr_kwargs) - - if output_log: - return { - "x": res.point, - "f(x)": res.cost, - "iterations": res.iterations, - "stopping_criterion": res.stopping_criterion, - "time": res.time, - "gradnorm": res.gradient_norm, - } - return res.point - -def solve_with_riemannian( - graph, - T_goal, - cache=True, - precon=None, - init="bsmooth", -): - """One-shot wrapper: build the IK problem from ``T_goal``, solve via - ``RiemannianSolver``, decode joint angles from the recovered point cloud. - - Parameters - ---------- - cache : bool, default True - Memoize per-Y cost state and projection ops across cost/grad/HVP - calls within a single iteration. Pure perf, identical trajectory. - precon : None | "gn" | "jacobi" | callable, default None - tCG preconditioner. "gn" (inverse Gauss-Newton Hessian, one small - eigh per outer iteration) cuts inner HVPs ~10x and wall time 2-5x - at equal-or-better success rate on the UR10 benchmark — see - experiments/rtr_preconditioner_study.py. Off by default because it - changes the optimization trajectory (baselines would shift). - init : {"spectral", "bsmooth", "zero"}, default "bsmooth" - Initialization strategy forwarded to ``RiemannianSolver``. - - Uses the bsmooth initialization (the ``RiemannianSolver`` default), so the - triangle-inequality bounds from ``bound_smoothing(G)`` are computed here - and forwarded as the ``bounds`` argument to ``solve`` when needed. - """ - G = graph.from_pose(T_goal) - solver = RiemannianSolver(graph, cache=cache, init=init) - D_goal = distance_matrix_from_graph(G) - omega = adjacency_matrix_from_graph(G) - bounds = None - if init == "bsmooth": - lb, ub = bound_smoothing(G) - bounds = (lb, ub) - sol_info = solver.solve( - D_goal, omega, use_limits=True, bounds=bounds, precon=precon - ) - G_sol = graph_from_pos(sol_info["x"], graph.node_ids) - q_sol = graph.joint_variables(G_sol, {f"p{graph.robot.n}": T_goal}) - - broken_limits = graph.check_distance_limits(graph.realization(q_sol), tol=1e-6) - if len(broken_limits) > 0: - return None, None - else: - return q_sol, sol_info["x"] diff --git a/graphik/solvers/scipy_solver.py b/graphik/solvers/scipy_solver.py new file mode 100644 index 0000000..6ebaaec --- /dev/null +++ b/graphik/solvers/scipy_solver.py @@ -0,0 +1,87 @@ +"""EDM solver backed by scipy.optimize.minimize.""" +from __future__ import annotations + +import numpy as np +from scipy.optimize import Bounds, minimize + +from graphik.graphs.graph import ProblemGraph +from graphik.solvers import costs +from graphik.solvers.distance_solver import ( + DistanceProblem, + DistanceSolver, + MinimizeInfo, +) +from graphik.utils.constants import POS + +_HESSP_METHODS = frozenset( + {"Newton-CG", "trust-ncg", "trust-krylov", "trust-constr"} +) + +_METHOD_DEFAULTS = { + "L-BFGS-B": {"ftol": 1e-16, "gtol": 1e-16}, + "BFGS": {"xrtol": 0.75e-6, "gtol": 0.25e-6, "norm": np.inf}, +} + + +class ScipySolver(DistanceSolver): + def __init__( + self, + graph: ProblemGraph, + *, + init: str = "bsmooth", + use_limits: bool = True, + cache: bool = True, + method: str = "BFGS", + options: dict | None = None, + ): + super().__init__(graph, init=init, use_limits=use_limits, cache=cache) + self.method = method + self.options = dict(options or {}) + + def position_constraints(self, G) -> Bounds: + """Pin POS-tagged nodes of ``G`` to their stored positions. + + For refactor parity this is called with the construction graph, where + only base anchors carry POS. The follow-up fix switches to the + per-solve goal graph. + """ + N = G.number_of_nodes() + ub = np.full((N, self.dim), np.inf) + lb = np.full((N, self.dim), -np.inf) + for idx, (node, data) in enumerate(G.nodes(data=True)): + if POS in data: + ub[idx] = data[POS] + lb[idx] = data[POS] + return Bounds(lb=lb.flatten(), ub=ub.flatten()) + + def _minimize(self, problem: DistanceProblem): + cost_and_grad, hessp = costs.for_minimize( + problem.D_goal, + problem.omega, + dim=self.dim, + psi_L=problem.psi_L, + psi_U=problem.psi_U, + cache=self.cache, + ) + + bounds = None + if self.method == "L-BFGS-B": + bounds = self.position_constraints(self.graph) + options = {**_METHOD_DEFAULTS.get(self.method, {}), **self.options} + + minimize_kwargs = dict( + jac=True, method=self.method, bounds=bounds, options=options + ) + if hessp is not None and self.method in _HESSP_METHODS: + minimize_kwargs["hessp"] = hessp + + sol = minimize( + cost_and_grad, + np.ascontiguousarray(problem.Y0.flatten()), + **minimize_kwargs, + ) + return sol.x.reshape(-1, self.dim), MinimizeInfo( + cost=float(sol.fun), + iterations=int(sol.nit), + status=str(sol.message), + ) diff --git a/tests/baselines/runner.py b/tests/baselines/runner.py index ddf68cd..c53c426 100644 --- a/tests/baselines/runner.py +++ b/tests/baselines/runner.py @@ -1,24 +1,18 @@ """Run a baseline case and measure pose error. -`run_case(case)` is the only public entrypoint; it returns a dict that goes +``run_case(case)`` is the only public entrypoint; it returns a dict that goes straight into the baseline JSON. All randomness is reseeded per-case so cases are independent regardless of order. """ from __future__ import annotations + import time + import numpy as np +from graphik.solvers import RiemannianSolver, ScipySolver from graphik.utils.roboturdf import load_schunk_lwa4d, load_ur10 from graphik.utils.utils import table_environment -from graphik.utils.dgp import ( - adjacency_matrix_from_graph, - bound_smoothing, - distance_matrix_from_graph, - graph_from_pos, -) -from graphik.solvers.riemannian_solver import solve_with_riemannian -from graphik.solvers.nonlinear_solver import NonlinearSolver - from tests.baselines.cases import Case ROBOT_LOADERS = { @@ -39,23 +33,13 @@ def _setup(case: Case): def _solve_riemannian(graph, T_goal): - q_sol, _ = solve_with_riemannian(graph, T_goal) - return q_sol + res = RiemannianSolver(graph).solve(T_goal) + return res.q if res.feasible else None def _solve_nonlinear(graph, T_goal, method: str): - G = graph.from_pose(T_goal) - D_goal = distance_matrix_from_graph(G) - omega = adjacency_matrix_from_graph(G) - lb, ub = bound_smoothing(G) - solver = NonlinearSolver(graph) - sol = solver.solve(D_goal, omega, use_limits=True, bounds=(lb, ub), method=method) - G_sol = graph_from_pos(sol["x"], graph.node_ids) - q_sol = graph.joint_variables(G_sol, {f"p{graph.robot.n}": T_goal}) - if len(graph.check_distance_limits(graph.realization(q_sol), tol=1e-6)) > 0: - return None - return q_sol - + res = ScipySolver(graph, method=method).solve(T_goal) + return res.q if res.feasible else None def _measure(robot, T_goal, q_sol) -> dict: @@ -66,10 +50,13 @@ def _measure(robot, T_goal, q_sol) -> dict: R_rel = T_goal[:3, :3] @ T_sol[:3, :3].T cos_theta = (np.trace(R_rel) - 1.0) / 2.0 rot_err = float(np.arccos(np.clip(cos_theta, -1.0, 1.0))) - # q_sol stored as a list sorted by joint id for forensic comparison only; - # the gate compares pose error, not joint values (multiple IK solutions exist). q_serialised = [(k, float(v)) for k, v in sorted(q_sol.items())] - return {"solved": True, "trans_err": t_err, "rot_err": rot_err, "q_sol": q_serialised} + return { + "solved": True, + "trans_err": t_err, + "rot_err": rot_err, + "q_sol": q_serialised, + } def run_case(case: Case) -> dict: diff --git a/tests/helpers.py b/tests/helpers.py new file mode 100644 index 0000000..07b7c7e --- /dev/null +++ b/tests/helpers.py @@ -0,0 +1,20 @@ +"""Shared builders for solver tests.""" +import numpy as np + +from graphik.graphs import ProblemGraph +from graphik.robots import Robot +from graphik.utils.utils import list_to_variable_dict + + +def planar_chain(n): + """n-joint planar chain with unit links and +/-pi joint limits.""" + params = { + "link_lengths": list_to_variable_dict(np.ones(n)), + "theta": list_to_variable_dict(np.zeros(n)), + "joint_limits_upper": list_to_variable_dict(np.pi * np.ones(n)), + "joint_limits_lower": list_to_variable_dict(-np.pi * np.ones(n)), + "num_joints": n, + "dim": 2, + } + robot = Robot(params) + return robot, ProblemGraph(robot) diff --git a/tests/test_loss.py b/tests/test_costs.py similarity index 67% rename from tests/test_loss.py rename to tests/test_costs.py index 95abd2b..8c60df4 100644 --- a/tests/test_loss.py +++ b/tests/test_costs.py @@ -1,4 +1,4 @@ -"""Finite-difference parity tests for graphik.solvers.loss. +"""Finite-difference parity tests for graphik.solvers.costs. The dense backend is checked against an FD approximation of its own gradient and HVP. egrad equals the analytical gradient of cost; ehvp @@ -11,7 +11,7 @@ import numpy as np -from graphik.solvers import loss +from graphik.solvers import costs from graphik.utils.dgp import distance_matrix_from_pos @@ -46,7 +46,7 @@ def _make_problem(n=6, d=3, seed=0): class TestDenseEquality(unittest.TestCase): def setUp(self): self.D_goal, self.omega, self.n, self.d, self.rng = _make_problem(seed=0) - self.cost, self.egrad, self.ehvp = loss._dense_equality( + self.cost, self.egrad, self.ehvp = costs._dense_equality( self.D_goal, self.omega, cache=True ) @@ -85,7 +85,7 @@ def setUp(self): psi_L[pair_L] = psi_L[pair_L[::-1]] = 5.0 * self.D_goal[pair_L] psi_U[pair_U] = psi_U[pair_U[::-1]] = 0.2 * self.D_goal[pair_U] self.psi_L, self.psi_U = psi_L, psi_U - self.cost, self.egrad, self.ehvp = loss._dense_limits( + self.cost, self.egrad, self.ehvp = costs._dense_limits( self.D_goal, self.omega, self.psi_L, self.psi_U, cache=True ) @@ -137,8 +137,8 @@ def setUp(self): self.D_goal, self.omega, self.n, self.d, self.rng = _make_problem(seed=4) def test_dense_dispatch_equality(self): - cost, egrad, _ = loss.for_riemannian(self.D_goal, self.omega) - ref_cost, ref_egrad, _ = loss._dense_equality(self.D_goal, self.omega, cache=True) + cost, egrad, _ = costs.for_riemannian(self.D_goal, self.omega) + ref_cost, ref_egrad, _ = costs._dense_equality(self.D_goal, self.omega, cache=True) Y = self.rng.standard_normal((self.n, self.d)) np.testing.assert_allclose(cost(Y), ref_cost(Y), atol=1e-12, rtol=0) np.testing.assert_allclose(egrad(Y), ref_egrad(Y), atol=1e-12, rtol=0) @@ -150,10 +150,10 @@ def test_dense_dispatch_limits(self): pair_U = (1, self.n - 2) psi_L[pair_L] = psi_L[pair_L[::-1]] = 5.0 * self.D_goal[pair_L] psi_U[pair_U] = psi_U[pair_U[::-1]] = 0.2 * self.D_goal[pair_U] - cost, _, _ = loss.for_riemannian( + cost, _, _ = costs.for_riemannian( self.D_goal, self.omega, psi_L=psi_L, psi_U=psi_U ) - ref_cost, _, _ = loss._dense_limits( + ref_cost, _, _ = costs._dense_limits( self.D_goal, self.omega, psi_L, psi_U, cache=True ) Y = self.rng.standard_normal((self.n, self.d)) @@ -165,7 +165,7 @@ def setUp(self): self.D_goal, self.omega, self.n, self.d, self.rng = _make_problem(seed=5) def test_cost_and_grad_shapes(self): - cost_and_grad, hessp = loss.for_minimize( + cost_and_grad, hessp = costs.for_minimize( self.D_goal, self.omega, dim=self.d ) Y_flat = self.rng.standard_normal(self.n * self.d) @@ -178,8 +178,8 @@ def test_cost_and_grad_shapes(self): def test_cost_and_grad_matches_riemannian(self): # for_minimize must return the same numbers as for_riemannian, just flattened. - cost, egrad, _ = loss.for_riemannian(self.D_goal, self.omega) - cost_and_grad, _ = loss.for_minimize( + cost, egrad, _ = costs.for_riemannian(self.D_goal, self.omega) + cost_and_grad, _ = costs.for_minimize( self.D_goal, self.omega, dim=self.d ) Y = self.rng.standard_normal((self.n, self.d)) @@ -196,7 +196,7 @@ def test_limits_dispatch_shapes(self): pair_U = (1, self.n - 2) psi_L[pair_L] = psi_L[pair_L[::-1]] = 5.0 * self.D_goal[pair_L] psi_U[pair_U] = psi_U[pair_U[::-1]] = 0.2 * self.D_goal[pair_U] - cost_and_grad, hessp = loss.for_minimize( + cost_and_grad, hessp = costs.for_minimize( self.D_goal, self.omega, dim=self.d, psi_L=psi_L, psi_U=psi_U, ) @@ -209,5 +209,88 @@ def test_limits_dispatch_shapes(self): self.assertEqual(hv.shape, (self.n * self.d,)) +class TestPoseCost(unittest.TestCase): + """FD parity for the SE(n) log pose loss.""" + + def test_cost_and_grad_matches_finite_differences_2d(self): + from graphik.graphs import ProblemGraph + from graphik.robots import Robot + from graphik.utils.utils import list_to_variable_dict + + n = 4 + params = { + "link_lengths": list_to_variable_dict(np.ones(n)), + "theta": list_to_variable_dict(np.zeros(n)), + "joint_limits_upper": list_to_variable_dict(np.pi * np.ones(n)), + "joint_limits_lower": list_to_variable_dict(-np.pi * np.ones(n)), + "num_joints": n, + "dim": 2, + } + robot = Robot(params) + ProblemGraph(robot) + point = f"p{n}" + q_goal = np.array([0.4, -0.3, 0.5, 0.2]) + T_goal = robot.pose(list_to_variable_dict(q_goal), point) + cost_and_grad = costs.pose_cost(robot, point, T_goal) + + q = np.array([-0.2, 0.6, -0.1, 0.3]) + _, grad = cost_and_grad(q) + + eps = 1e-6 + fd = np.zeros(n) + for i in range(n): + dq = np.zeros(n) + dq[i] = eps + f_plus, _ = cost_and_grad(q + dq) + f_minus, _ = cost_and_grad(q - dq) + fd[i] = (f_plus - f_minus) / (2 * eps) + + np.testing.assert_allclose(grad, fd, rtol=1e-4, atol=1e-6) + + +class TestForMinimizeUR10(unittest.TestCase): + """UR10-scale FD checks for the flat for_minimize interface.""" + + @classmethod + def setUpClass(cls): + from graphik.utils.dgp import adjacency_matrix_from_graph, pos_from_graph + from graphik.utils.roboturdf import load_ur10 + + robot, graph = load_ur10() + rng = np.random.default_rng(0) + joint_names = list(robot.random_configuration().keys()) + q = {name: float(rng.uniform(-np.pi, np.pi)) for name in joint_names} + G = graph.realization(q) + cls.omega = adjacency_matrix_from_graph(G) + cls.D_goal = distance_matrix_from_pos(pos_from_graph(G)) + cls.dim = robot.dim + cls.n_points = cls.omega.shape[0] + cls.rng = np.random.default_rng(1) + + def test_gradient_matches_fd(self): + cost_and_grad, _ = costs.for_minimize(self.D_goal, self.omega, dim=self.dim) + for trial in range(5): + Y = self.rng.standard_normal((self.n_points, self.dim)).ravel() + _, g = cost_and_grad(Y) + g_fd = _fd_gradient(Y, lambda Yf: cost_and_grad(Yf)[0]) + np.testing.assert_allclose( + g, g_fd, atol=1e-6, + err_msg=f"trial {trial}: analytical gradient != FD", + ) + + def test_hvp_matches_fd_of_grad(self): + cost_and_grad, hessp = costs.for_minimize(self.D_goal, self.omega, dim=self.dim) + grad_only = lambda Y: cost_and_grad(Y)[1] + for trial in range(5): + Y = self.rng.standard_normal((self.n_points, self.dim)).ravel() + w = self.rng.standard_normal(Y.size) + hv = hessp(Y, w) + hv_fd = _fd_jvp(Y, grad_only, w) + np.testing.assert_allclose( + hv, hv_fd, atol=1e-3, + err_msg=f"trial {trial}: analytical HVP != FD", + ) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_distance_solver.py b/tests/test_distance_solver.py new file mode 100644 index 0000000..f0021b6 --- /dev/null +++ b/tests/test_distance_solver.py @@ -0,0 +1,66 @@ +"""Pipeline tests for the DistanceSolver mid-layer.""" +import numpy as np +import pytest + +from graphik.solvers.distance_solver import ( + DistanceProblem, + DistanceSolver, + MinimizeInfo, +) +from graphik.solvers.initializations import zero_init +from tests.helpers import planar_chain + + +class _IdentityMinimizer(DistanceSolver): + """Returns Y0 untouched; no optimizer is involved.""" + + def _minimize(self, problem: DistanceProblem): + self.last_problem = problem + return problem.Y0, MinimizeInfo(cost=0.0, iterations=1, status="identity") + + +def test_invalid_init_rejected_at_construction(): + robot, graph = planar_chain(4) + with pytest.raises(ValueError): + _IdentityMinimizer(graph, init="bogus") + + +def test_pipeline_decodes_exact_realization(): + robot, graph = planar_chain(4) + solver = _IdentityMinimizer(graph) + q_zero = robot.zero_configuration() + T_goal = np.asarray(robot.pose(q_zero, f"p{robot.n}")) + Y0 = zero_init(graph) + + result = solver.solve(T_goal, Y_init=Y0) + + assert result.Y is Y0 + assert result.status == "identity" + assert result.iterations == 1 + assert result.time > 0 + assert set(result.q) == set(q_zero) + np.testing.assert_allclose(np.array(list(result.q.values())), 0.0, atol=1e-8) + assert result.feasible + + +def test_problem_state_is_complete(): + robot, graph = planar_chain(4) + solver = _IdentityMinimizer(graph) + T_goal = np.asarray(robot.pose(robot.zero_configuration(), f"p{robot.n}")) + solver.solve(T_goal, Y_init=zero_init(graph)) + + p = solver.last_problem + N = graph.number_of_nodes() + assert p.D_goal.shape == (N, N) + assert p.omega.shape == (N, N) + assert p.psi_L is not None and p.psi_U is not None + assert p.Y0.shape == (N, graph.dim) + + +def test_use_limits_false_skips_bound_matrices(): + robot, graph = planar_chain(4) + solver = _IdentityMinimizer(graph, use_limits=False) + T_goal = np.asarray(robot.pose(robot.zero_configuration(), f"p{robot.n}")) + solver.solve(T_goal, Y_init=zero_init(graph)) + assert solver.last_problem.psi_L is None + assert solver.last_problem.psi_U is None diff --git a/tests/test_initializations.py b/tests/test_initializations.py new file mode 100644 index 0000000..3079bbb --- /dev/null +++ b/tests/test_initializations.py @@ -0,0 +1,49 @@ +"""Behavior tests for the shared Y-initialization strategies.""" +import numpy as np + +from graphik.solvers.initializations import ( + INIT_STRATEGIES, + bsmooth_init, + spectral_init, + zero_init, +) +from graphik.utils import POS +from graphik.utils.dgp import adjacency_matrix_from_graph, distance_matrix_from_graph +from graphik.utils.roboturdf import load_ur10 + + +def _ur10_problem(seed=0): + np.random.seed(seed) + robot, graph = load_ur10() + T_goal = robot.pose(robot.random_configuration(), f"p{robot.n}") + G = graph.from_pose(T_goal) + return graph, G, distance_matrix_from_graph(G), adjacency_matrix_from_graph(G) + + +def test_strategy_names(): + assert INIT_STRATEGIES == ("spectral", "bsmooth", "zero") + + +def test_spectral_shape_and_rank(): + graph, G, D_goal, omega = _ur10_problem() + Y = spectral_init(D_goal, omega, graph.dim) + assert Y.shape == (D_goal.shape[0], graph.dim) + assert np.isfinite(Y).all() + assert np.linalg.matrix_rank(Y) == graph.dim + + +def test_bsmooth_shape(seed=0): + graph, G, D_goal, omega = _ur10_problem(seed) + Y = bsmooth_init(G, omega, graph.dim) + assert Y.shape == (D_goal.shape[0], graph.dim) + assert np.isfinite(Y).all() + + +def test_zero_init_matches_zero_realization(): + graph, G, D_goal, omega = _ur10_problem() + Y = zero_init(graph) + G_zero = graph.realization(graph.robot.zero_configuration()) + expected = np.stack( + [np.asarray(G_zero.nodes[node][POS], dtype=float) for node in graph.node_ids] + ) + np.testing.assert_array_equal(Y, expected) diff --git a/tests/test_joint_angle_solver.py b/tests/test_joint_angle_solver.py index b5ec325..bf70c0a 100644 --- a/tests/test_joint_angle_solver.py +++ b/tests/test_joint_angle_solver.py @@ -1,80 +1,90 @@ -"""Behavior tests for LocalSolver (joint-space SLSQP solver).""" +"""Behavior tests for JointAngleSolver (joint-space SLSQP solver).""" +from types import SimpleNamespace + import numpy as np -import pytest -from graphik.robots import Robot -from graphik.graphs import ProblemGraph -from graphik.solvers.joint_angle_solver import LocalSolver +import graphik.solvers.joint_angle as joint_angle_mod +from graphik.solvers import IKResult, JointAngleSolver from graphik.utils.utils import list_to_variable_dict - - -def planar_chain(n): - params = { - "link_lengths": list_to_variable_dict(np.ones(n)), - "theta": list_to_variable_dict(np.zeros(n)), - "joint_limits_upper": list_to_variable_dict(np.pi * np.ones(n)), - "joint_limits_lower": list_to_variable_dict(-np.pi * np.ones(n)), - "num_joints": n, - "dim": 2, - } - robot = Robot(params) - return robot, ProblemGraph(robot) - - -class TestGradientConsistency2d: - def test_cost_and_grad_matches_finite_differences(self): - """The 2D branch must use the true SE2 inverse left Jacobian; a - stubbed identity makes the returned gradient inconsistent with the - returned cost away from the goal.""" - n = 4 - robot, graph = planar_chain(n) - solver = LocalSolver(graph, {}) - point = f"p{n}" - - q_goal = np.array([0.4, -0.3, 0.5, 0.2]) - T_goal = robot.pose(list_to_variable_dict(q_goal), point) - cost_and_grad = solver.gen_cost_and_grad_ee(point, T_goal) - - q = np.array([-0.2, 0.6, -0.1, 0.3]) - _, grad = cost_and_grad(q) - - eps = 1e-6 - fd = np.zeros(n) - for i in range(n): - dq = np.zeros(n) - dq[i] = eps - f_plus, _ = cost_and_grad(q + dq) - f_minus, _ = cost_and_grad(q - dq) - fd[i] = (f_plus - f_minus) / (2 * eps) - - np.testing.assert_allclose(grad, fd, rtol=1e-4, atol=1e-6) +from tests.helpers import planar_chain class TestMultiGoalSolve: def test_solve_optimizes_all_goals_not_just_the_last(self): - """Two consistent goals (intermediate point and end effector, both - from the same configuration). If solve() drops every goal but the - last dict entry, the end-effector goal here is silently ignored.""" + """Two consistent goals through the shared solve(T_goal) signature.""" n = 6 robot, graph = planar_chain(n) - solver = LocalSolver(graph, {}) + solver = JointAngleSolver(graph) q_goal = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6]) q_goal_dict = list_to_variable_dict(q_goal) T_ee = robot.pose(q_goal_dict, f"p{n}") T_mid = robot.pose(q_goal_dict, "p3") - # End-effector goal first, intermediate goal last: the historical - # bug kept only the last entry. goals = {f"p{n}": T_ee, "p3": T_mid} - # Start with the first three joints already at the goal, the rest - # off: the p3 cost alone is zero here, so a solver that ignores the - # p{n} goal terminates immediately. q0 = np.concatenate([q_goal[:3], np.zeros(3)]) - res = solver.solve(goals, list_to_variable_dict(q0)) + result = solver.solve(goals, q_init=list_to_variable_dict(q0)) - q_sol = list_to_variable_dict(res.x) - err_ee = np.linalg.norm(robot.pose(q_sol, f"p{n}") - T_ee) - err_mid = np.linalg.norm(robot.pose(q_sol, "p3") - T_mid) + assert isinstance(result, IKResult) + err_ee = np.linalg.norm(robot.pose(result.q, f"p{n}") - T_ee) + err_mid = np.linalg.norm(robot.pose(result.q, "p3") - T_mid) assert err_mid < 1e-3 assert err_ee < 1e-3 + + +class TestSharedContract: + def test_q_init_defaults_to_zero_configuration(self): + n = 4 + robot, graph = planar_chain(n) + solver = JointAngleSolver(graph) + q_goal = list_to_variable_dict(np.array([0.3, -0.2, 0.4, 0.1])) + T_goal = np.asarray(robot.pose(q_goal, f"p{n}")) + + result = solver.solve(T_goal) + + T_sol = np.asarray(robot.pose(result.q, f"p{n}")) + assert np.linalg.norm(T_sol - T_goal) < 1e-3 + + def test_result_fields(self): + n = 4 + robot, graph = planar_chain(n) + solver = JointAngleSolver(graph) + T_goal = np.asarray( + robot.pose( + list_to_variable_dict(np.array([0.4, -0.3, 0.5, 0.2])), + f"p{n}", + ) + ) + result = solver.solve(T_goal) + assert result.Y is None + assert isinstance(result.cost, float) + assert isinstance(result.status, str) + assert result.iterations >= 1 + assert result.time > 0 + assert isinstance(result.limit_violations, list) + + def test_q_init_is_ordered_by_joint_name_not_dict_insertion_order(self, monkeypatch): + n = 4 + robot, graph = planar_chain(n) + solver = JointAngleSolver(graph) + q_goal = list_to_variable_dict(np.array([0.2, -0.1, 0.3, -0.2])) + T_goal = np.asarray(robot.pose(q_goal, f"p{n}")) + q_init = list_to_variable_dict(np.array([0.4, -0.3, 0.2, -0.1])) + shuffled = {joint: q_init[joint] for joint in reversed(list(q_init))} + captured = {} + + def fake_minimize(fun, x0, **kwargs): + captured["x0"] = np.asarray(x0) + return SimpleNamespace( + x=np.asarray(x0), fun=0.0, nit=0, message="fake" + ) + + monkeypatch.setattr(joint_angle_mod, "minimize", fake_minimize) + + result = solver.solve(T_goal, q_init=shuffled) + + expected = np.array([q_init[joint] for joint in solver.joint_order]) + np.testing.assert_array_equal(captured["x0"], expected) + assert result.q == { + joint: value for joint, value in zip(solver.joint_order, expected) + } diff --git a/tests/test_nonlinear_solver_gradients.py b/tests/test_nonlinear_solver_gradients.py deleted file mode 100644 index eb8ade6..0000000 --- a/tests/test_nonlinear_solver_gradients.py +++ /dev/null @@ -1,81 +0,0 @@ -"""Finite-difference checks for the analytical gradient and HVP exposed by -``NonlinearSolver.create_cost``. - -After the loss-consolidation refactor there is a single dense backend; the -prior ``loop`` / ``sparse`` / ``dense`` test variants collapse to one. -""" -from __future__ import annotations - -import unittest - -import numpy as np - -from graphik.solvers.nonlinear_solver import NonlinearSolver -from graphik.utils.dgp import ( - adjacency_matrix_from_graph, - distance_matrix_from_pos, - pos_from_graph, -) -from graphik.utils.roboturdf import load_ur10 - - -def _numerical_gradient(Y, f, eps=1e-6): - g = np.zeros_like(Y) - perturb = np.zeros_like(Y) - for i in range(Y.size): - perturb[i] = eps - loss_plus, _ = f(Y + perturb) - loss_minus, _ = f(Y - perturb) - g[i] = (loss_plus - loss_minus) / (2 * eps) - perturb[i] = 0 - return g - - -def _numerical_hvp(Y, grad_fn, w, eps=1e-5): - return (grad_fn(Y + eps * w) - grad_fn(Y - eps * w)) / (2 * eps) - - -class NonlinearSolverGradientTests(unittest.TestCase): - @classmethod - def setUpClass(cls): - robot, graph = load_ur10() - rng = np.random.default_rng(0) - joint_names = list(robot.random_configuration().keys()) - q = {name: float(rng.uniform(-np.pi, np.pi)) for name in joint_names} - G = graph.realization(q) - cls.graph = graph - cls.omega = adjacency_matrix_from_graph(G) - cls.D_goal = distance_matrix_from_pos(pos_from_graph(G)) - cls.dim = robot.dim - cls.n_points = cls.omega.shape[0] - cls.rng = np.random.default_rng(1) - - def test_gradient_matches_fd(self): - solver = NonlinearSolver(self.graph) - cost_and_grad, _ = solver.create_cost(self.D_goal, self.omega) - for trial in range(5): - Y = self.rng.standard_normal((self.n_points, self.dim)).ravel() - _, g = cost_and_grad(Y) - g_fd = _numerical_gradient(Y, cost_and_grad) - np.testing.assert_allclose( - g, g_fd, atol=1e-6, - err_msg=f"trial {trial}: analytical gradient != FD", - ) - - def test_hvp_matches_fd_of_grad(self): - solver = NonlinearSolver(self.graph) - cost_and_grad, hessp = solver.create_cost(self.D_goal, self.omega) - grad_only = lambda Y: cost_and_grad(Y)[1] - for trial in range(5): - Y = self.rng.standard_normal((self.n_points, self.dim)).ravel() - w = self.rng.standard_normal(Y.size) - hv = hessp(Y, w) - hv_fd = _numerical_hvp(Y, grad_only, w) - np.testing.assert_allclose( - hv, hv_fd, atol=1e-3, - err_msg=f"trial {trial}: analytical HVP != FD", - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_preconditioner.py b/tests/test_preconditioner.py index c9e7315..4f58e58 100644 --- a/tests/test_preconditioner.py +++ b/tests/test_preconditioner.py @@ -2,17 +2,11 @@ import numpy as np import pytest -from graphik.solvers.riemannian_solver import ( +from graphik.solvers.riemannian import ( RiemannianSolver, make_gn_preconditioner, make_jacobi_preconditioner, ) -from graphik.utils.dgp import ( - adjacency_matrix_from_graph, - bound_smoothing, - distance_matrix_from_graph, - graph_from_pos, -) from graphik.utils.manifolds.fixed_rank_psd_sym import PSDFixedRank from graphik.utils.roboturdf import load_ur10 @@ -24,7 +18,6 @@ def _setup(seed=3, N=10, d=3): rng = np.random.default_rng(seed) manifold = PSDFixedRank(N, d) Y = rng.standard_normal((N, d)) - # A connected random edge mask (symmetric, zero diagonal) omega = (rng.random((N, N)) < 0.5).astype(float) omega = np.triu(omega, 1) omega += omega.T @@ -37,7 +30,6 @@ def test_output_is_tangent(factory): precon = factory(manifold, omega) r = manifold.projection(Y, rng.standard_normal(Y.shape)) z = precon(Y, r) - # Horizontal-space condition for the PSDFixedRank quotient: Y^T Z symmetric asym = Y.T @ z - z.T @ Y np.testing.assert_allclose(asym, 0.0, atol=1e-9) @@ -48,11 +40,9 @@ def test_symmetric_positive_definite_on_tangent_space(factory): precon = factory(manifold, omega) u = manifold.projection(Y, rng.standard_normal(Y.shape)) v = manifold.projection(Y, rng.standard_normal(Y.shape)) - # Symmetry: == np.testing.assert_allclose( np.sum(u * precon(Y, v)), np.sum(precon(Y, u) * v), rtol=1e-9 ) - # Positive definiteness along tangent directions assert np.sum(u * precon(Y, u)) > 0 assert np.sum(v * precon(Y, v)) > 0 @@ -61,24 +51,17 @@ def test_gn_solve_matches_unpreconditioned_quality(): np.random.seed(5) robot, graph = load_ur10() q = robot.random_configuration() - T_goal = robot.pose(q, f"p{robot.n}") + T_goal = np.asarray(robot.pose(q, f"p{robot.n}")) - G = graph.from_pose(T_goal) - D_goal = distance_matrix_from_graph(G) - omega = adjacency_matrix_from_graph(G) - bounds = bound_smoothing(G) + solver = RiemannianSolver(graph, precon="gn") + result = solver.solve(T_goal) - solver = RiemannianSolver(graph) - out = solver.solve(D_goal, omega, use_limits=True, bounds=bounds, precon="gn") - G_sol = graph_from_pos(out["x"], graph.node_ids) - q_sol = graph.joint_variables(G_sol, {f"p{robot.n}": T_goal}) - T_sol = robot.pose(q_sol, f"p{robot.n}") + T_sol = np.asarray(robot.pose(result.q, f"p{robot.n}")) assert np.linalg.norm(T_sol[:3, 3] - T_goal[:3, 3]) < 1e-2 assert np.linalg.norm(T_sol[:3, :3] - T_goal[:3, :3]) < 1e-2 def test_unknown_precon_raises(): robot, graph = load_ur10() - solver = RiemannianSolver(graph) with pytest.raises(ValueError): - solver.solve(np.eye(4), np.zeros((4, 4)), precon="bogus") + RiemannianSolver(graph, precon="bogus") diff --git a/tests/test_pymlg_conventions.py b/tests/test_pymlg_conventions.py index 77107b3..7150c12 100644 --- a/tests/test_pymlg_conventions.py +++ b/tests/test_pymlg_conventions.py @@ -107,14 +107,13 @@ def test_so2_exp_log_round_trip(): def test_se3_left_jacobian_inv_at_origin_is_identity(): - """left_jacobian_inv(0) = I_6. Used by joint_angle_solver.""" + """left_jacobian_inv(0) = I_6. Used by costs.pose_cost.""" J = SE3.left_jacobian_inv(np.zeros(6)) np.testing.assert_allclose(J, np.eye(6), atol=1e-12) def test_se2_left_jacobian_inv_exists(): - """LocalSolver.gen_cost_and_grad_ee uses SE2.left_jacobian_inv when - dim==2; its 2D gradient is wrong without it.""" + """costs.pose_cost uses SE2.left_jacobian_inv in 2D.""" assert hasattr(SE2, "left_jacobian_inv") J = SE2.left_jacobian_inv(np.zeros(3)) np.testing.assert_allclose(J, np.eye(3), atol=1e-12) diff --git a/tests/test_removed_kwargs.py b/tests/test_removed_kwargs.py index e909121..d733ee7 100644 --- a/tests/test_removed_kwargs.py +++ b/tests/test_removed_kwargs.py @@ -1,31 +1,25 @@ -"""Removed-kwarg guards on the solver constructors. - -Both solver __init__s setattr unknown kwargs onto self, so without an -explicit guard a stale ``jit=True`` from a pre-removal caller would be -silently swallowed instead of erroring. -""" +"""Stale kwargs from removed APIs must raise TypeError.""" import unittest -from graphik.solvers.nonlinear_solver import NonlinearSolver -from graphik.solvers.riemannian_solver import RiemannianSolver +from graphik.solvers import JointAngleSolver, RiemannianSolver, ScipySolver class TestRemovedKwargs(unittest.TestCase): - def test_riemannian_rejects_jit(self): + def test_riemannian_rejects_stale_kwargs(self): with self.assertRaisesRegex(TypeError, "jit"): RiemannianSolver(None, jit=True) - - def test_riemannian_rejects_cost_type(self): with self.assertRaisesRegex(TypeError, "cost_type"): RiemannianSolver(None, cost_type="dense") - def test_nonlinear_rejects_jit(self): + def test_scipy_rejects_stale_kwargs(self): with self.assertRaisesRegex(TypeError, "jit"): - NonlinearSolver(None, jit=True) - - def test_nonlinear_rejects_cost_type(self): + ScipySolver(None, jit=True) with self.assertRaisesRegex(TypeError, "cost_type"): - NonlinearSolver(None, cost_type="dense") + ScipySolver(None, cost_type="dense") + + def test_joint_angle_rejects_old_params_positional(self): + with self.assertRaises(TypeError): + JointAngleSolver(None, {}) if __name__ == "__main__": diff --git a/tests/test_scipy_solver.py b/tests/test_scipy_solver.py new file mode 100644 index 0000000..7c1e433 --- /dev/null +++ b/tests/test_scipy_solver.py @@ -0,0 +1,36 @@ +"""Behavior tests for ScipySolver.""" +import numpy as np + +from graphik.solvers.scipy_solver import ScipySolver +from graphik.utils.constants import POS +from tests.helpers import planar_chain + + +def test_bfgs_reaches_reachable_goal(): + np.random.seed(42) + robot, graph = planar_chain(6) + solver = ScipySolver(graph, method="BFGS") + q_goal = robot.random_configuration() + T_goal = np.asarray(robot.pose(q_goal, f"p{robot.n}")) + + result = solver.solve(T_goal) + + d = graph.dim + T_sol = np.asarray(robot.pose(result.q, f"p{robot.n}")) + assert np.linalg.norm(T_goal[:d, d] - T_sol[:d, d]) < 1e-2 + assert np.linalg.norm(T_goal[:d, :d] - T_sol[:d, :d]) < 1e-2 + assert result.Y.shape == (graph.number_of_nodes(), d) + + +def test_position_constraints_pin_exactly_the_pos_tagged_nodes(): + robot, graph = planar_chain(4) + solver = ScipySolver(graph, method="L-BFGS-B") + bnds = solver.position_constraints(graph) + lb = bnds.lb.reshape(-1, graph.dim) + ub = bnds.ub.reshape(-1, graph.dim) + for idx, (node, data) in enumerate(graph.nodes(data=True)): + if POS in data: + np.testing.assert_array_equal(lb[idx], data[POS]) + np.testing.assert_array_equal(ub[idx], data[POS]) + else: + assert np.all(lb[idx] == -np.inf) and np.all(ub[idx] == np.inf) diff --git a/tests/test_solver_contract.py b/tests/test_solver_contract.py new file mode 100644 index 0000000..b65abd5 --- /dev/null +++ b/tests/test_solver_contract.py @@ -0,0 +1,112 @@ +"""Shared-contract tests and cross-solver solve(T_goal) matrix.""" +import numpy as np +import pytest + +from graphik.solvers import ( + IKResult, + IKSolver, + JointAngleSolver, + RiemannianSolver, + ScipySolver, +) +from graphik.solvers.distance_solver import DistanceSolver +from tests.helpers import planar_chain + + +class _StubSolver(IKSolver): + """Minimal concrete IKSolver for base-class helpers.""" + + def solve(self, T_goal, **kwargs): + return IKResult( + q={}, + cost=0.0, + iterations=0, + time=0.0, + status="stub", + limit_violations=[], + ) + + +def test_ik_solver_is_abstract(): + robot, graph = planar_chain(4) + with pytest.raises(TypeError): + IKSolver(graph) + + +def test_feasible_iff_no_limit_violations(): + kw = dict(q={}, cost=0.0, iterations=0, time=0.0, status="") + assert IKResult(limit_violations=[], **kw).feasible + assert not IKResult(limit_violations=[{"edge": ("p1", "p2")}], **kw).feasible + + +def test_goals_from_wraps_single_transform_with_primary_ee(): + robot, graph = planar_chain(4) + solver = _StubSolver(graph) + T = np.eye(3) + goals = solver.goals_from(T) + assert goals == {robot.end_effectors[0]: T} + assert robot.end_effectors[0] == f"p{robot.n}" + + +def test_goals_from_copies_dict_input(): + robot, graph = planar_chain(4) + solver = _StubSolver(graph) + original = {"p2": np.eye(3)} + goals = solver.goals_from(original) + assert goals == original + assert goals is not original + + +def test_check_limits_empty_at_interior_configuration(): + robot, graph = planar_chain(4) + solver = _StubSolver(graph) + q = {f"p{i}": 0.1 for i in range(1, 5)} + assert solver.check_limits(q) == [] + + +SOLVER_FACTORIES = [ + pytest.param(lambda g: RiemannianSolver(g), id="riemannian"), + pytest.param(lambda g: ScipySolver(g, method="BFGS"), id="scipy-bfgs"), + pytest.param(lambda g: ScipySolver(g, method="L-BFGS-B"), id="scipy-lbfgsb"), + pytest.param(lambda g: JointAngleSolver(g), id="joint-angle"), +] + + +@pytest.mark.parametrize("factory", SOLVER_FACTORIES) +def test_solve_returns_well_formed_result(factory): + np.random.seed(42) + robot, graph = planar_chain(6) + solver = factory(graph) + q_goal = robot.random_configuration() + T_goal = np.asarray(robot.pose(q_goal, f"p{robot.n}")) + + result = solver.solve(T_goal) + + assert isinstance(result, IKResult) + assert set(result.q.keys()) == set(q_goal.keys()) + assert isinstance(result.cost, float) + assert result.iterations >= 0 + assert result.time > 0 + assert isinstance(result.status, str) + assert isinstance(result.limit_violations, list) + assert result.feasible == (len(result.limit_violations) == 0) + if isinstance(solver, DistanceSolver): + assert result.Y.shape == (graph.number_of_nodes(), graph.dim) + else: + assert result.Y is None + + +@pytest.mark.parametrize("factory", SOLVER_FACTORIES) +def test_solve_reaches_reachable_goal(factory): + np.random.seed(42) + robot, graph = planar_chain(6) + solver = factory(graph) + q_goal = robot.random_configuration() + T_goal = np.asarray(robot.pose(q_goal, f"p{robot.n}")) + + result = solver.solve(T_goal) + + d = graph.dim + T_sol = np.asarray(robot.pose(result.q, f"p{robot.n}")) + assert np.linalg.norm(T_goal[:d, d] - T_sol[:d, d]) < 1e-2 + assert np.linalg.norm(T_goal[:d, :d] - T_sol[:d, :d]) < 1e-2 From 59ffbfcc6dc3e0090de5da29a5dda5eef389e048 Mon Sep 17 00:00:00 2001 From: Filip Maric Date: Thu, 11 Jun 2026 08:55:01 +0200 Subject: [PATCH 8/9] fix(solvers): pin L-BFGS-B goal nodes Build L-BFGS-B position Bounds from the per-solve goal graph so POS-tagged end-effector goal nodes are hard-pinned instead of only represented by soft EDM residuals. Regenerate the solver-consolidation baseline for the intentional L-BFGS-B behavior change. --- graphik/solvers/scipy_solver.py | 7 +- ...6-06-11-solver-consolidation-baseline.json | 244 ++++++++++++++++++ tests/test_scipy_solver.py | 37 ++- 3 files changed, 275 insertions(+), 13 deletions(-) create mode 100644 tests/baselines/2026-06-11-solver-consolidation-baseline.json diff --git a/graphik/solvers/scipy_solver.py b/graphik/solvers/scipy_solver.py index 6ebaaec..289a285 100644 --- a/graphik/solvers/scipy_solver.py +++ b/graphik/solvers/scipy_solver.py @@ -41,9 +41,8 @@ def __init__( def position_constraints(self, G) -> Bounds: """Pin POS-tagged nodes of ``G`` to their stored positions. - For refactor parity this is called with the construction graph, where - only base anchors carry POS. The follow-up fix switches to the - per-solve goal graph. + Called with the per-solve goal graph, so the end-effector goal nodes + (POS-tagged there) are pinned alongside the base anchors. """ N = G.number_of_nodes() ub = np.full((N, self.dim), np.inf) @@ -66,7 +65,7 @@ def _minimize(self, problem: DistanceProblem): bounds = None if self.method == "L-BFGS-B": - bounds = self.position_constraints(self.graph) + bounds = self.position_constraints(problem.G) options = {**_METHOD_DEFAULTS.get(self.method, {}), **self.options} minimize_kwargs = dict( diff --git a/tests/baselines/2026-06-11-solver-consolidation-baseline.json b/tests/baselines/2026-06-11-solver-consolidation-baseline.json new file mode 100644 index 0000000..95bc016 --- /dev/null +++ b/tests/baselines/2026-06-11-solver-consolidation-baseline.json @@ -0,0 +1,244 @@ +{ + "baseline_label": "solver-consolidation-lbfgsb-ee-pinning", + "cases": [ + { + "elapsed_sec": 0.7572524168062955, + "name": "riemann_schunk_free", + "obstacles": false, + "q_sol": [ + [ + "p1", + 1.4801681993055937 + ], + [ + "p2", + -2.4223480695934665 + ], + [ + "p3", + 3.1270835954174565 + ], + [ + "p4", + 0.6297268332254121 + ], + [ + "p5", + 0.0005973837141309488 + ], + [ + "p6", + -1.65042783433981 + ], + [ + "p7", + -2.4960495602032347 + ] + ], + "robot": "schunk_lwa4d", + "rot_err": 0.0010536527563540677, + "seed": 42, + "solved": true, + "solver": "riemannian", + "trans_err": 0.0009624828265227867 + }, + { + "elapsed_sec": 9.900694834068418, + "name": "riemann_schunk_obstacles", + "obstacles": true, + "q_sol": [ + [ + "p1", + -2.4644641300393495 + ], + [ + "p2", + 2.9558558768020684 + ], + [ + "p3", + -1.8984299735892958 + ], + [ + "p4", + 0.6235219059259767 + ], + [ + "p5", + 2.566408466058777 + ], + [ + "p6", + -2.241225078262413 + ], + [ + "p7", + -2.3126562618440003 + ] + ], + "robot": "schunk_lwa4d", + "rot_err": 0.00040182225863074885, + "seed": 42, + "solved": true, + "solver": "riemannian", + "trans_err": 0.0003658774543899778 + }, + { + "elapsed_sec": 0.09440699988044798, + "name": "riemann_ur10_free", + "obstacles": false, + "q_sol": [ + [ + "p1", + -0.7895761971977215 + ], + [ + "p2", + -2.053988169563212 + ], + [ + "p3", + -1.4574615597920593 + ], + [ + "p4", + 2.137215019393664 + ], + [ + "p5", + -2.1606410541657977 + ], + [ + "p6", + -2.161949643420102 + ] + ], + "robot": "ur10", + "rot_err": 0.0006731579956268483, + "seed": 42, + "solved": true, + "solver": "riemannian", + "trans_err": 0.0006140372098513563 + }, + { + "elapsed_sec": 3.7709082909859717, + "name": "riemann_ur10_obstacles", + "obstacles": true, + "q_sol": [ + [ + "p1", + 1.3890639892329724 + ], + [ + "p2", + 1.7501975941676926 + ], + [ + "p3", + 1.9623548898985819 + ], + [ + "p4", + -2.31653963950726 + ], + [ + "p5", + -1.9308093061820248 + ], + [ + "p6", + 0.8073654422649152 + ] + ], + "robot": "ur10", + "rot_err": 0.00044422340080434023, + "seed": 42, + "solved": true, + "solver": "riemannian", + "trans_err": 0.0005816867501727357 + }, + { + "elapsed_sec": 0.05102654197253287, + "name": "nonlin_bfgs_schunk", + "obstacles": false, + "q_sol": [ + [ + "p1", + 1.485784198178859 + ], + [ + "p2", + -2.4228730990880067 + ], + [ + "p3", + 3.1343388916609203 + ], + [ + "p4", + 0.62676186166186 + ], + [ + "p5", + -0.01119074952525843 + ], + [ + "p6", + -1.6529289831947986 + ], + [ + "p7", + -2.4997576056800574 + ] + ], + "robot": "schunk_lwa4d", + "rot_err": 0.0007748905568186698, + "seed": 42, + "solved": true, + "solver": "nonlinear_bfgs", + "trans_err": 0.0006874708806062909 + }, + { + "elapsed_sec": 0.3198007920291275, + "name": "nonlin_lbfgsb_schunk", + "obstacles": false, + "q_sol": [ + [ + "p1", + 1.4791452125463092 + ], + [ + "p2", + -2.4361518035920247 + ], + [ + "p3", + -0.016663408559282843 + ], + [ + "p4", + -0.5956465056665671 + ], + [ + "p5", + -3.1412664752464847 + ], + [ + "p6", + -1.6730792057969357 + ], + [ + "p7", + -2.4948141433341355 + ] + ], + "robot": "schunk_lwa4d", + "rot_err": 0.0028847869691253003, + "seed": 42, + "solved": true, + "solver": "nonlinear_lbfgsb", + "trans_err": 0.002839600381299503 + } + ], + "schema_version": 1 +} \ No newline at end of file diff --git a/tests/test_scipy_solver.py b/tests/test_scipy_solver.py index 7c1e433..eeed55c 100644 --- a/tests/test_scipy_solver.py +++ b/tests/test_scipy_solver.py @@ -22,15 +22,34 @@ def test_bfgs_reaches_reachable_goal(): assert result.Y.shape == (graph.number_of_nodes(), d) -def test_position_constraints_pin_exactly_the_pos_tagged_nodes(): +def test_position_constraints_pin_goal_nodes(): + np.random.seed(42) robot, graph = planar_chain(4) solver = ScipySolver(graph, method="L-BFGS-B") - bnds = solver.position_constraints(graph) + T_goal = np.asarray(robot.pose(robot.random_configuration(), f"p{robot.n}")) + G = graph.from_pose(T_goal) + bnds = solver.position_constraints(G) lb = bnds.lb.reshape(-1, graph.dim) - ub = bnds.ub.reshape(-1, graph.dim) - for idx, (node, data) in enumerate(graph.nodes(data=True)): - if POS in data: - np.testing.assert_array_equal(lb[idx], data[POS]) - np.testing.assert_array_equal(ub[idx], data[POS]) - else: - assert np.all(lb[idx] == -np.inf) and np.all(ub[idx] == np.inf) + + pinned = { + node: idx for idx, (node, data) in enumerate(G.nodes(data=True)) if POS in data + } + assert f"p{robot.n}" in pinned + ee_row = lb[pinned[f"p{robot.n}"]] + np.testing.assert_allclose(ee_row, T_goal[:2, 2], atol=1e-12) + for idx in range(lb.shape[0]): + if idx not in pinned.values(): + assert np.all(lb[idx] == -np.inf) + + +def test_lbfgsb_solution_has_pinned_ee(): + np.random.seed(42) + robot, graph = planar_chain(6) + solver = ScipySolver(graph, method="L-BFGS-B") + q_goal = robot.random_configuration() + T_goal = np.asarray(robot.pose(q_goal, f"p{robot.n}")) + + result = solver.solve(T_goal) + + node_index = list(graph.node_ids).index(f"p{robot.n}") + np.testing.assert_allclose(result.Y[node_index], T_goal[:2, 2], atol=1e-12) From 1e4b27e0a80742462064f0ef398cd78eba9654d9 Mon Sep 17 00:00:00 2001 From: Filip Maric Date: Thu, 11 Jun 2026 16:19:57 +0200 Subject: [PATCH 9/9] refactor: dedupe limit geometry, consolidate loaders, prune dead code Simplification pass over the non-SDP codebase (sdp_snl, convex_iteration, and sdp_formulations are deliberately untouched): - graphs: extract _revolute_edge_bounds, the revolute reachable-distance core previously duplicated between _root_angle_limits_3d and _set_limits_3d; collapse check_distance_limits to one classification + two-sided check; O(1) node index in distance_bound_matrices; replace the star import with explicit imports; cached_property for node lists. - robots: drop the no-op kinematic_map property pair and the dead limited_joints property; cached_property for joint_ids/end_effectors/ T_base. - roboturdf: collapse the five identical load_* functions into load_urdf_robot + URDF_FILES (old names kept as thin wrappers); remove unused members (find_joint_by_name, find_joints_with_parent_link, joint_limits, n_urdf_joints, scene, urdf_ind_to_q); make get_parents and find_end_effector_joints side-effect free. - delete robot_visualization.py (unused; plot_3d_chain_manipulator called the nonexistent Robot.get_pose) and rtr's unused stopping-criterion label helpers. - move table_environment to graphik/utils/environments.py; give graphik.utils explicit re-exports with __all__. - packaging: migrate setup.py to pyproject.toml; export the public API (solvers, ProblemGraph, Robot) from the top-level graphik package. Verified behavior-preserving: full suite passes (123) and graph construction artifacts (distance/bound matrices, BOUNDED tags, limited_joints, violation reports) are bit-identical to the pre-refactor tree for UR10, KUKA, and a planar chain. Co-Authored-By: Claude Opus 4.8 --- README.md | 2 +- experiments/cidgik_example.py | 2 +- experiments/riemannian_example.py | 2 +- experiments/solver_benchmark.py | 2 +- graphik/__init__.py | 24 +++ graphik/graphs/graph.py | 286 +++++++++++++-------------- graphik/robots/robot.py | 74 +++---- graphik/solvers/joint_angle.py | 5 +- graphik/solvers/rtr.py | 14 -- graphik/utils/__init__.py | 99 +++++++++- graphik/utils/dgp.py | 2 +- graphik/utils/environments.py | 21 ++ graphik/utils/robot_visualization.py | 274 ------------------------- graphik/utils/roboturdf.py | 199 +++++-------------- graphik/utils/utils.py | 15 -- pyproject.toml | 37 ++++ setup.py | 30 --- tests/baselines/runner.py | 2 +- 18 files changed, 401 insertions(+), 689 deletions(-) create mode 100644 graphik/utils/environments.py delete mode 100644 graphik/utils/robot_visualization.py create mode 100644 pyproject.toml delete mode 100644 setup.py diff --git a/README.md b/README.md index c050343..a62dadd 100644 --- a/README.md +++ b/README.md @@ -58,7 +58,7 @@ GraphIK's interface between robot models and IK solvers is the [`ProblemGraph`]( If you are considering an environment with spherical obstacles, you can include constraints that prevent collisions. In this example, we will use a set of spheres that approximate a table: ```python -from graphik.utils.utils import table_environment +from graphik.utils.environments import table_environment obstacles = table_environment() # This loop is not needed if you are not using obstacle avoidance constraints for idx, obs in enumerate(obstacles): diff --git a/experiments/cidgik_example.py b/experiments/cidgik_example.py index 0b847a9..ed031d4 100644 --- a/experiments/cidgik_example.py +++ b/experiments/cidgik_example.py @@ -1,4 +1,4 @@ -from graphik.utils.utils import table_environment +from graphik.utils.environments import table_environment from graphik.solvers.convex_iteration import solve_with_cidgik # Multiple robot models to try out, or you can implement your own diff --git a/experiments/riemannian_example.py b/experiments/riemannian_example.py index 953fcce..aa762bc 100644 --- a/experiments/riemannian_example.py +++ b/experiments/riemannian_example.py @@ -1,6 +1,6 @@ from time import perf_counter -from graphik.utils.utils import table_environment +from graphik.utils.environments import table_environment from graphik.solvers import RiemannianSolver # Multiple robot models to try out, or you can implement your own diff --git a/experiments/solver_benchmark.py b/experiments/solver_benchmark.py index 830d30c..d8ce77d 100644 --- a/experiments/solver_benchmark.py +++ b/experiments/solver_benchmark.py @@ -23,7 +23,7 @@ import numpy as np from graphik.solvers import RiemannianSolver, ScipySolver -from graphik.utils.utils import table_environment +from graphik.utils.environments import table_environment from graphik.utils.roboturdf import ( load_ur10, load_kuka, diff --git a/graphik/__init__.py b/graphik/__init__.py index e69de29..84f1185 100644 --- a/graphik/__init__.py +++ b/graphik/__init__.py @@ -0,0 +1,24 @@ +"""graphIK: distance-geometric inverse kinematics. + +Top-level convenience exports; submodules (``graphik.utils``, +``graphik.solvers.sdp_snl``, ...) remain importable directly. +""" +from graphik.graphs import ProblemGraph +from graphik.robots import Robot +from graphik.solvers import ( + IKResult, + IKSolver, + JointAngleSolver, + RiemannianSolver, + ScipySolver, +) + +__all__ = [ + "IKResult", + "IKSolver", + "JointAngleSolver", + "ProblemGraph", + "RiemannianSolver", + "Robot", + "ScipySolver", +] diff --git a/graphik/graphs/graph.py b/graphik/graphs/graph.py index ab0300e..5c57e58 100644 --- a/graphik/graphs/graph.py +++ b/graphik/graphs/graph.py @@ -1,4 +1,5 @@ """Unified ProblemGraph class parameterized by dim in {2, 3}.""" +from functools import cached_property from math import sqrt, atan2 from typing import Any, Dict, List, Optional, Tuple, Union @@ -10,7 +11,35 @@ from pymlg.numpy import SE3 from graphik.robots.robot import Robot -from graphik.utils import * # noqa: F401,F403 (existing convention; brings constants and helpers into scope: ROOT, POS, TYPE, BASE, ROBOT, END_EFFECTOR, OBSTACLE, BOUNDED, BELOW, ABOVE, LOWER, UPPER, DIST, MAIN_PREFIX, AUX_PREFIX, TRANSFORM, trans_axis, rot_axis, skew, normalize, wraptopi, R2, best_fit_transform, level2_descendants, max_min_distance_revolute, distance_matrix_from_graph, adjacency_matrix_from_graph, graph_complete_edges) +from graphik.utils.constants import ( + ABOVE, + AUX_PREFIX, + BASE, + BELOW, + BOUNDED, + DIST, + END_EFFECTOR, + LOWER, + MAIN_PREFIX, + OBSTACLE, + POS, + ROBOT, + ROOT, + TYPE, + UPPER, +) +from graphik.utils.dgp import ( + adjacency_matrix_from_graph, + distance_matrix_from_graph, + graph_complete_edges, +) +from graphik.utils.geometry import ( + best_fit_transform, + max_min_distance_revolute, + skew, +) +from graphik.utils.kinematics import R2, rot_axis, trans_axis +from graphik.utils.utils import level2_descendants, normalize, wraptopi def _bounded_tags(limit) -> List[str]: @@ -45,29 +74,17 @@ def __init__(self, robot: Robot, params: Optional[Dict] = None): # ------------------------------------------------------------------ # Properties (unchanged from the previous ProblemGraph base class) # ------------------------------------------------------------------ - @property + @cached_property def base_nodes(self) -> List[str]: - try: - return self._base_nodes - except AttributeError: - self._base_nodes = [n for n, d in self.nodes(data=TYPE) if BASE in d] - return self._base_nodes + return [n for n, d in self.nodes(data=TYPE) if BASE in d] - @property + @cached_property def structure_nodes(self) -> List[str]: - try: - return self._structure_nodes - except AttributeError: - self._structure_nodes = [n for n, d in self.nodes(data=TYPE) if ROBOT in d] - return self._structure_nodes + return [n for n, d in self.nodes(data=TYPE) if ROBOT in d] - @property + @cached_property def end_effector_nodes(self) -> List[str]: - try: - return self._end_effector_nodes - except AttributeError: - self._end_effector_nodes = [n for n, d in self.nodes(data=TYPE) if END_EFFECTOR in d] - return self._end_effector_nodes + return [n for n, d in self.nodes(data=TYPE) if END_EFFECTOR in d] @property def base(self) -> nx.DiGraph: @@ -157,54 +174,45 @@ def clear_obstacles(self): obstacles = [n for n, t in node_types.items() if t == [OBSTACLE]] self.remove_nodes_from(obstacles) - def check_distance_limits(self, G: nx.DiGraph, tol=1e-10) -> List[Dict[str, List[Any]]]: + def check_distance_limits(self, G: nx.DiGraph, tol=1e-10) -> List[Dict[str, Any]]: typ = nx.get_node_attributes(self, name=TYPE) broken_limits = [] for u, v, data in self.edges(data=True): bounds = data.get(BOUNDED, []) - if BELOW in bounds or ABOVE in bounds: - if G[u][v][DIST] < data[LOWER] - tol: - bl = {} - if (ROBOT in typ[u] and OBSTACLE in typ[v]) or ( - OBSTACLE in typ[u] and ROBOT in typ[v] - ): - bl["edge"] = (u, v) - bl["value"] = G[u][v][DIST] - data[LOWER] - bl["type"] = OBSTACLE - bl["side"] = LOWER - broken_limits += [bl] - if ROBOT in typ[u] and ROBOT in typ[v]: - bl["edge"] = (u, v) - bl["value"] = G[u][v][DIST] - data[LOWER] - bl["type"] = "joint" - bl["side"] = LOWER - broken_limits += [bl] - if G[u][v][DIST] > data[UPPER] + tol: - bl = {} - if (ROBOT in typ[u] and OBSTACLE in typ[v]) or ( - OBSTACLE in typ[u] and ROBOT in typ[v] - ): - bl["edge"] = (u, v) - bl["value"] = G[u][v][DIST] - data[UPPER] - bl["type"] = OBSTACLE - bl["side"] = UPPER - broken_limits += [bl] - if ROBOT in typ[u] and ROBOT in typ[v]: - bl["edge"] = (u, v) - bl["value"] = G[u][v][DIST] - data[UPPER] - bl["type"] = "joint" - bl["side"] = UPPER - broken_limits += [bl] + if BELOW not in bounds and ABOVE not in bounds: + continue + if (ROBOT in typ[u] and OBSTACLE in typ[v]) or ( + OBSTACLE in typ[u] and ROBOT in typ[v] + ): + kind = OBSTACLE + elif ROBOT in typ[u] and ROBOT in typ[v]: + kind = "joint" + else: + continue + dist = G[u][v][DIST] + for side, violated in ( + (LOWER, dist < data[LOWER] - tol), + (UPPER, dist > data[UPPER] + tol), + ): + if violated: + broken_limits.append( + { + "edge": (u, v), + "value": dist - data[side], + "type": kind, + "side": side, + } + ) return broken_limits def distance_bound_matrices(self) -> Tuple[ArrayLike, ArrayLike]: n_nodes = self.number_of_nodes() L = np.zeros([n_nodes, n_nodes]) U = np.zeros([n_nodes, n_nodes]) + index = {node: idx for idx, node in enumerate(self.nodes())} for e1, e2, data in self.edges(data=True): if BOUNDED in data: - udx = self.node_ids.index(e1) - vdx = self.node_ids.index(e2) + udx, vdx = index[e1], index[e2] bounds = data[BOUNDED] if BELOW in bounds: L[udx, vdx] = data[LOWER] ** 2 @@ -362,6 +370,16 @@ def set_limits(self): self._set_limits_3d() # --- 2D bodies (verbatim from the old ProblemGraphPlanar) --- + @staticmethod + def _two_link_bounds(l1: float, l2: float, angle_limit: float) -> Tuple[float, float]: + """Max/min reach of a 2-link planar chain whose hinge ranges over + ``[-angle_limit, angle_limit]``. The upper bound is full extension + ``l1 + l2``; the lower bound is the law of cosines at the extreme + angle (interior angle ``pi - angle_limit``).""" + upper = l1 + l2 + lower = sqrt(l1 ** 2 + l2 ** 2 - 2 * l1 * l2 * cos(pi - angle_limit)) + return upper, lower + def _root_angle_limits_2d(self): ax = "x" S = self.structure @@ -372,11 +390,10 @@ def _root_angle_limits_2d(self): lb = self.robot.lb[node] ub = self.robot.ub[node] lim = max(abs(ub), abs(lb)) + upper, lower = self._two_link_bounds(l1, l2, lim) self.add_edge(ax, node) - self[ax][node][UPPER] = l1 + l2 - self[ax][node][LOWER] = sqrt( - l1 ** 2 + l2 ** 2 - 2 * l1 * l2 * cos(pi - lim) - ) + self[ax][node][UPPER] = upper + self[ax][node][LOWER] = lower self[ax][node][BOUNDED] = [BELOW] def _set_limits_2d(self): @@ -392,64 +409,69 @@ def _set_limits_2d(self): lb = self.robot.lb[ids[2]] ub = self.robot.ub[ids[2]] lim = max(abs(ub), abs(lb)) + upper, lower = self._two_link_bounds(l1, l2, lim) self.add_edge(u, v) - self[u][v][UPPER] = l1 + l2 - self[u][v][LOWER] = sqrt( - l1 ** 2 + l2 ** 2 - 2 * l1 * l2 * cos(pi - lim) - ) + self[u][v][UPPER] = upper + self[u][v][LOWER] = lower self[u][v][BOUNDED] = [BELOW] # --- 3D bodies (verbatim from the old ProblemGraphRevolute) --- + @staticmethod + def _revolute_edge_bounds(T0, T1, T2, upper_limit): + """Distance bounds between the fixed point at ``T0`` and the point at + ``T2`` as it revolves around the z-axis of joint frame ``T1``. + + Returns ``(d_max, d_min, limit)`` where ``limit`` tags which bound is + attained at the zero configuration (``BELOW``/``ABOVE``), ``False`` + when the circle degenerates, and ``None`` when neither bound is + attained. When a bound is attained, it is tightened to the distance at + the joint's ``upper_limit`` angle (joint limits are symmetric here). + """ + N = T1[:3, 2] + C = T1[:3, 3] + (N.dot(T2[:3, 3] - T1[:3, 3])) * N + r = la.norm(T2[:3, 3] - C) + P = T0[:3, 3] + d_max, d_min = max_min_distance_revolute(r, P, C, N) + d = la.norm(T2[:3, 3] - T0[:3, 3]) + + if d_max == d_min: + limit = False + elif d == d_max: + limit = BELOW + elif d == d_min: + limit = ABOVE + else: + limit = None + + if limit: + T_rel = SE3.inverse(T1) @ T2 + d_limit = la.norm( + (T1 @ rot_axis(upper_limit, "z") @ T_rel)[:3, 3] - T0[:3, 3] + ) + if limit == ABOVE: + d_max = d_limit + else: + d_min = d_limit + return d_max, d_min, limit + def _root_angle_limits_3d(self): - axis_length = self.axis_length robot = self.robot - upper_limits = self.robot.ub limited_joints = self.limited_joints T1 = robot.nodes[ROOT]["T0"] - base_names = ["x", "y"] - names = ["p1", "q1"] - T_axis = trans_axis(axis_length, "z") + T_axis = trans_axis(self.axis_length, "z") - for base_node in base_names: - for node in names: + for base_node in ["x", "y"]: + for node in ["p1", "q1"]: T0 = np.eye(4) T0[:3, 3] = self.nodes[base_node][POS] - if node[0] == "p": - T2 = robot.nodes["p1"]["T0"] - else: - T2 = robot.nodes["p1"]["T0"] @ T_axis - - N = T1[:3, 2] - C = T1[:3, 3] + (N.dot(T2[:3, 3] - T1[:3, 3])) * N - r = np.linalg.norm(T2[:3, 3] - C) - P = T0[:3, 3] - d_max, d_min = max_min_distance_revolute(r, P, C, N) - d = np.linalg.norm(T2[:3, 3] - T0[:3, 3]) - - if d_max == d_min: - limit = False - elif d == d_max: - limit = BELOW - elif d == d_min: - limit = ABOVE - else: - limit = None + T2 = robot.nodes["p1"]["T0"] + if node[0] == AUX_PREFIX: + T2 = T2 @ T_axis + d_max, d_min, limit = self._revolute_edge_bounds( + T0, T1, T2, robot.ub["p1"] + ) if limit: - if node[0] == "p": - T_rel = SE3.inverse(T1) @ robot.nodes["p1"]["T0"] - else: - T_rel = SE3.inverse(T1) @ (robot.nodes["p1"]["T0"] @ T_axis) - - d_limit = la.norm( - (T1 @ rot_axis(upper_limits["p1"], "z") @ T_rel)[:3, 3] - - T0[:3, 3] - ) - - if limit == ABOVE: - d_max = d_limit - else: - d_min = d_limit limited_joints += ["p1"] self.add_edge(base_node, node) @@ -464,56 +486,31 @@ def _set_limits_3d(self): robot = self.robot kinematic_map = self.robot.kinematic_map T_axis = trans_axis(self.axis_length, "z") - end_effectors = self.robot.end_effectors - upper_limits = self.robot.ub limited_joints = [] - for ee in end_effectors: + for ee in self.robot.end_effectors: k_map = kinematic_map[ROOT][ee] for idx in range(2, len(k_map)): cur, prev = k_map[idx], k_map[idx - 2] names = [ - (MAIN_PREFIX + str(prev[1:]), MAIN_PREFIX + str(cur[1:])), - (MAIN_PREFIX + str(prev[1:]), AUX_PREFIX + str(cur[1:])), - (AUX_PREFIX + str(prev[1:]), MAIN_PREFIX + str(cur[1:])), - (AUX_PREFIX + str(prev[1:]), AUX_PREFIX + str(cur[1:])), + (prefix_u + str(prev[1:]), prefix_v + str(cur[1:])) + for prefix_u in (MAIN_PREFIX, AUX_PREFIX) + for prefix_v in (MAIN_PREFIX, AUX_PREFIX) ] for ids in names: path = kinematic_map[prev][cur] - T0, T1, T2 = [ - robot.nodes[path[0]]["T0"], - robot.nodes[path[1]]["T0"], - robot.nodes[path[2]]["T0"], - ] + T0 = robot.nodes[path[0]]["T0"] + T1 = robot.nodes[path[1]]["T0"] + T2 = robot.nodes[path[2]]["T0"] if AUX_PREFIX in ids[0]: T0 = T0 @ T_axis if AUX_PREFIX in ids[1]: T2 = T2 @ T_axis - N = T1[:3, 2] - C = T1[:3, 3] + (N.dot(T2[:3, 3] - T1[:3, 3])) * N - r = la.norm(T2[:3, 3] - C) - P = T0[:3, 3] - d_max, d_min = max_min_distance_revolute(r, P, C, N) - - d = la.norm(T2[:3, 3] - T0[:3, 3]) - if d_max == d_min: - limit = False - elif d == d_max: - limit = BELOW - elif d == d_min: - limit = ABOVE - else: - limit = None - + d_max, d_min, limit = self._revolute_edge_bounds( + T0, T1, T2, robot.ub[cur] + ) if limit: - rot_limit = rot_axis(upper_limits[cur], "z") - T_rel = SE3.inverse(T1) @ T2 - d_limit = la.norm((T1 @ rot_limit @ T_rel)[:3, 3] - T0[:3, 3]) - if limit == ABOVE: - d_max = d_limit - else: - d_min = d_limit limited_joints += [cur] self.add_edge(ids[0], ids[1]) @@ -600,24 +597,23 @@ def joint_variables( k_map = kinematic_map[ROOT][ee] for idx in range(1, len(k_map)): cur, aux_cur = k_map[idx], f"q{k_map[idx][1:]}" - pred, aux_pred = (k_map[idx - 1], f"q{k_map[idx-1][1:]}") + pred = k_map[idx - 1] T_prev = T[pred] T_prev_0 = self.robot.nodes[pred]["T0"] T_0 = self.robot.nodes[cur]["T0"] - T_0_q = self.robot.nodes[cur]["T0"] @ trans_axis(axis_length, "z") T_rel = SE3.inverse(T_prev_0) @ T_0 - ps_0 = (SE3.inverse(T_prev_0) @ T_0)[:3, 3] - qs_0 = (SE3.inverse(T_prev_0) @ T_0_q)[:3, 3] + # Rest-pose direction of the auxiliary (q) node in the + # predecessor frame. T_0_q = T_0 @ trans_axis, so the + # inverse is already captured in T_rel. + qs_0 = (T_rel @ trans_axis(axis_length, "z"))[:3, 3] - p = B_inv[:3, :3] @ G.nodes[cur][POS] + B_inv[:3, 3] qnorm = G.nodes[cur][POS] + ( G.nodes[aux_cur][POS] - G.nodes[cur][POS] ) / la.norm(G.nodes[aux_cur][POS] - G.nodes[cur][POS]) q = B_inv[:3, :3] @ qnorm + B_inv[:3, 3] T_prev_inv = SE3.inverse(T_prev) - ps = T_prev_inv[:3, :3] @ (p - T_prev[:3, 3]) qs = T_prev_inv[:3, :3] @ (q - T_prev[:3, 3]) theta[cur] = arctan2( diff --git a/graphik/robots/robot.py b/graphik/robots/robot.py index 38fad23..cf739c4 100644 --- a/graphik/robots/robot.py +++ b/graphik/robots/robot.py @@ -1,7 +1,8 @@ #!/usr/bin/env python3 """Unified Robot class parameterized by dim in {2, 3}.""" +from functools import cached_property from math import pi -from typing import Any, Dict, List, Tuple, Union +from typing import Any, Dict, List, Union import networkx as nx import numpy as np @@ -72,43 +73,17 @@ def __init__(self, params: Dict): # ------------------------------------------------------------------ # Properties (unchanged behaviour from the previous Robot base class) # ------------------------------------------------------------------ - @property - def kinematic_map(self) -> dict: - return self._kinematic_map - - @kinematic_map.setter - def kinematic_map(self, kinematic_map: dict): - self._kinematic_map = kinematic_map - - @property + @cached_property def joint_ids(self) -> List[str]: - try: - return self._joint_ids - except AttributeError: - self._joint_ids = list(self.kinematic_map.keys()) - return self._joint_ids + return list(self.kinematic_map.keys()) - @property + @cached_property def end_effectors(self) -> List: - if not hasattr(self, "_end_effectors"): - self._end_effectors = [x for x in self.nodes() if self.out_degree(x) == 0] - return self._end_effectors + return [x for x in self.nodes() if self.out_degree(x) == 0] - @property + @cached_property def T_base(self) -> np.ndarray: - try: - return self._T_base - except AttributeError: - self._T_base = self.nodes[ROOT]["T0"] - return self._T_base - - @property - def limited_joints(self) -> List[str]: - return self._limited_joints - - @limited_joints.setter - def limited_joints(self, lim: List[str]): - self._limited_joints = lim + return self.nodes[ROOT]["T0"] @property def ub(self) -> Dict[str, Any]: @@ -155,28 +130,31 @@ def get_all_poses(self, joint_angles: Dict[str, Any]) -> Dict[str, np.ndarray]: # Dim-aware kinematics # ------------------------------------------------------------------ def _screw_axis_from_T0(self, T0: np.ndarray) -> np.ndarray: - """Return the screw axis (twist coords) for a revolute joint with frame T0.""" + """Return the screw axis (twist coords) for a revolute joint with frame T0. + + The twist is ``[omega, v]`` with ``v = -omega × q``. A planar joint + always rotates about +z, so the SE(2) twist ``[omega_z, v_x, v_y]`` + reduces to the closed form ``[1, q_y, -q_x]``. + """ if self.dim == 2: - omega = np.array([0, 0, 1]) - q = np.hstack((T0[:2, 2], 0)) - return np.hstack((np.cross(-omega, q), omega))[[5, 0, 1]] + qx, qy = T0[0, 2], T0[1, 2] + return np.array([1.0, qy, -qx]) else: omega = T0[:3, 2] q = T0[:3, 3] return np.hstack((omega, np.cross(-omega, q))) def set_geometric_attributes(self): - for ee in self.end_effectors: - k_map = self.kinematic_map[ROOT][ee] - for idx in range(len(k_map)): - cur = k_map[idx] - self.nodes[cur]["S"] = self._screw_axis_from_T0(self.nodes[cur]["T0"]) - if idx != 0: - pred = k_map[idx - 1] - self[pred][cur][TRANSFORM] = ( - self._SE.inverse(self.nodes[pred]["T0"]) - @ self.nodes[cur]["T0"] - ) + # The screw axis S is a per-node property and the relative TRANSFORM is + # a per-edge property; neither depends on which end-effector path + # reaches it, so iterate over nodes/edges directly rather than walking + # (and recomputing along) every kinematic path. + for node in self.nodes: + self.nodes[node]["S"] = self._screw_axis_from_T0(self.nodes[node]["T0"]) + for pred, cur in self.edges: + self[pred][cur][TRANSFORM] = ( + self._SE.inverse(self.nodes[pred]["T0"]) @ self.nodes[cur]["T0"] + ) def from_params(self) -> Dict[str, np.ndarray]: """Construct T_zero from link lengths. 2D-only today (no fk_tree_3d helper).""" diff --git a/graphik/solvers/joint_angle.py b/graphik/solvers/joint_angle.py index abcc333..e99f3e2 100644 --- a/graphik/solvers/joint_angle.py +++ b/graphik/solvers/joint_angle.py @@ -65,10 +65,7 @@ def cost_and_grad(q): method="SLSQP", options={"ftol": 1e-7}, ) - q = { - joint: value - for joint, value in zip(self.joint_order, res.x) - } + q = dict(zip(self.joint_order, res.x)) return IKResult( q=q, cost=float(res.fun), diff --git a/graphik/solvers/rtr.py b/graphik/solvers/rtr.py index 86c0731..bd90721 100644 --- a/graphik/solvers/rtr.py +++ b/graphik/solvers/rtr.py @@ -23,20 +23,6 @@ class StopReason(IntEnum): MODEL_INCREASED = 5 -_TCG_LABEL = { - StopReason.NEGATIVE_CURVATURE: "negative curvature", - StopReason.EXCEEDED_TR: "exceeded trust region", - StopReason.REACHED_TARGET_LINEAR: "reached target residual-kappa (linear)", - StopReason.REACHED_TARGET_SUPERLINEAR: "reached target residual-theta (superlinear)", - StopReason.MAX_INNER_ITER: "maximum inner iterations", - StopReason.MODEL_INCREASED: "model increased", -} - - -def stopping_criterion_label(reason: StopReason) -> str: - return _TCG_LABEL[reason] - - @dataclass class RTRResult: point: np.ndarray diff --git a/graphik/utils/__init__.py b/graphik/utils/__init__.py index 4acf8b9..04ae394 100644 --- a/graphik/utils/__init__.py +++ b/graphik/utils/__init__.py @@ -1,5 +1,94 @@ -from graphik.utils.dgp import * -from graphik.utils.geometry import * -from graphik.utils.utils import * -from graphik.utils.kinematics import * -from graphik.utils.constants import * +"""Public surface of graphik.utils. + +Everything re-exported here may be imported as ``from graphik.utils import X``. +Demo/benchmark fixtures (``graphik.utils.environments``) and the URDF loader +(``graphik.utils.roboturdf``) are deliberately not re-exported; import them +from their modules directly. +""" +from graphik.utils.constants import ( + ABOVE, + AUX_PREFIX, + BASE, + BELOW, + BOUNDED, + DIST, + END_EFFECTOR, + FEASIBLE, + INFEASIBLE, + LOWER, + MAIN_PREFIX, + OBSTACLE, + POS, + RADIUS, + ROBOT, + ROOT, + SOLVER_ERROR, + TRANSFORM, + TYPE, + UNDEFINED, + UPPER, +) +from graphik.utils.dgp import ( + MDS, + adjacency_matrix_from_graph, + bound_smoothing, + distance_matrix_from_gram, + distance_matrix_from_graph, + distance_matrix_from_pos, + gram_from_distance_matrix, + graph_complete_edges, + graph_from_pos, + graph_from_pos_dict, + linear_projection, + normalize_positions, + pos_from_graph, +) +from graphik.utils.geometry import ( + best_fit_transform, + max_min_distance_revolute, + skew, +) +from graphik.utils.kinematics import ( + R2, + Rx, + Ry, + Rz, + dh_to_se2, + dh_to_se3, + fk_3d, + fk_tree_2d, + modified_dh_to_se3, + modified_fk_3d, + rot_axis, + trans_axis, +) +from graphik.utils.utils import ( + flatten, + level2_descendants, + list_to_variable_dict, + memoize_last, + normalize, + wraptopi, +) + +__all__ = [ + # constants + "ABOVE", "AUX_PREFIX", "BASE", "BELOW", "BOUNDED", "DIST", "END_EFFECTOR", + "FEASIBLE", "INFEASIBLE", "LOWER", "MAIN_PREFIX", "OBSTACLE", "POS", + "RADIUS", "ROBOT", "ROOT", "SOLVER_ERROR", "TRANSFORM", "TYPE", + "UNDEFINED", "UPPER", + # dgp + "MDS", "adjacency_matrix_from_graph", "bound_smoothing", + "distance_matrix_from_gram", "distance_matrix_from_graph", + "distance_matrix_from_pos", "gram_from_distance_matrix", + "graph_complete_edges", "graph_from_pos", "graph_from_pos_dict", + "linear_projection", "normalize_positions", "pos_from_graph", + # geometry + "best_fit_transform", "max_min_distance_revolute", "skew", + # kinematics + "R2", "Rx", "Ry", "Rz", "dh_to_se2", "dh_to_se3", "fk_3d", "fk_tree_2d", + "modified_dh_to_se3", "modified_fk_3d", "rot_axis", "trans_axis", + # utils + "flatten", "level2_descendants", "list_to_variable_dict", "memoize_last", + "normalize", "wraptopi", +] diff --git a/graphik/utils/dgp.py b/graphik/utils/dgp.py index f058a44..89f7b01 100644 --- a/graphik/utils/dgp.py +++ b/graphik/utils/dgp.py @@ -3,7 +3,7 @@ import numpy as np import networkx as nx from numpy.typing import NDArray -from graphik.utils.constants import * +from graphik.utils.constants import DIST, LOWER, POS, UPPER def gram_from_distance_matrix(D: NDArray) -> NDArray: diff --git a/graphik/utils/environments.py b/graphik/utils/environments.py new file mode 100644 index 0000000..296e9b4 --- /dev/null +++ b/graphik/utils/environments.py @@ -0,0 +1,21 @@ +"""Obstacle environments for benchmarks and examples.""" +import numpy as np + + +def table_environment(height=0.9, width=0.8, n_height=9, n_width=8, obs_inflation=2.): + """Sphere-approximated table: a tabletop grid plus four legs. + + Returns a list of ``(center, radius)`` pairs suitable for + ``ProblemGraph.add_spherical_obstacle``. + """ + # Make the tabletop + radius = 0.5*height/n_height + tabletop_obs = [(np.asarray([2*(i+0.5)*radius, 2*(j+0.5)*radius, height+radius]), obs_inflation*radius) + for i in range(-n_width//2, n_width//2) for j in range(-n_width//2, n_width//2)] + + leg1 = [(np.asarray([-width/2+radius, -width/2+radius, (2*i+1)*radius]), obs_inflation*radius) for i in range(0, n_height)] + leg2 = [(np.asarray([-width/2+radius, width/2-radius, (2*i+1)*radius]), obs_inflation*radius) for i in range(0, n_height)] + leg3 = [(np.asarray([ width/2-radius, -width/2+radius, (2*i+1)*radius]), obs_inflation*radius) for i in range(0, n_height)] + leg4 = [(np.asarray([ width/2-radius, width/2-radius, (2*i+1)*radius]), obs_inflation*radius) for i in range(0, n_height)] + + return tabletop_obs + leg1 + leg2 + leg3 + leg4 diff --git a/graphik/utils/robot_visualization.py b/graphik/utils/robot_visualization.py deleted file mode 100644 index 836ab47..0000000 --- a/graphik/utils/robot_visualization.py +++ /dev/null @@ -1,274 +0,0 @@ -import numpy as np -from matplotlib import pyplot as plt - -# from matplotlib.cbook import get_sample_data -from matplotlib.offsetbox import OffsetImage, AnnotationBbox -from matplotlib import rc, cm -from graphik.robots import Robot -import matplotlib as mpl -from mpl_toolkits.mplot3d import Axes3D - -# mpl.rcParams['legend.fontsize'] = 10 - -# LaTeX text rendering is opt-in: enabling it at import time breaks the -# module for users without a TeX install. Call this before plotting if -# you want publication-style labels. -def use_latex_text(): - rc("font", **{"family": "serif", "serif": ["Computer Modern"]}) - rc("text", usetex=True) - - -def plot_heatmap( - X, Y, Z, fig=None, ax=None, title=None, xlabel=None, ylabel=None, cmap=cm.coolwarm -): - """ - Taken from an answer at https://stackoverflow.com/questions/33282368/plotting-a-2d-heatmap-with-matplotlib by user - Erasmus Cedernaes. - :param X: matrix coordinates (same size as Z) from np.meshgrid - :param Y: matrix coordinates (same size as Z) from np.meshgrid - :param Z: matrix coordinates (same size as Z) from np.meshgrid - :param fig: figure handle - :param ax: axes handle - :param title: figure title - :param xlabel: x-axis label - :param ylabel: y-axis label - :param cmap: colormap to use - :return: - """ - if ax is None: - fig, ax = plt.subplots() - z = Z[:-1, :-1] - z_min, z_max = z.min(), z.max() - c = ax.pcolormesh(X, Y, Z, cmap=cmap, vmin=z_min, vmax=z_max) - if title is not None: - ax.set_title(title) - if xlabel is not None: - plt.xlabel(xlabel) - if ylabel is not None: - plt.ylabel(ylabel) - # set the limits of the plot to the limits of the data - ax.axis([X.min(), X.max(), Y.min(), Y.max()]) - fig.colorbar(c, ax=ax) - return fig, ax - - -def plot_surface( - X, Y, Z, fig=None, ax=None, title=None, xlabel=None, ylabel=None, cmap=cm.coolwarm -): - if ax is None: - fig = plt.figure() - ax = fig.gca(projection="3d") - - # Plot the surface. - surf = ax.plot_surface(X, Y, Z, cmap=cmap, linewidth=0, antialiased=False) - # Customize the z axis. - # ax.set_zlim(-1.01, 1.01) - # ax.zaxis.set_major_locator(LinearLocator(10)) - # ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f')) - # Add a color bar which maps values to colors. - fig.colorbar(surf, shrink=0.5, aspect=5) - return fig, ax - - -def get_image(path, zoom=1): - return OffsetImage(plt.imread(path), zoom=zoom) - - -def plot_image(path, x, y, ax, zoom=1): - ab = AnnotationBbox(get_image(path, zoom=zoom), (x, y), frameon=False) - ax.add_artist(ab) - - -def get_tree_paths(parent_nodes): - n_nodes = len(parent_nodes) + 1 - heads = [[0]] - for idx, parent_node in enumerate(parent_nodes): - for head in heads: - if head[-1] == parent_node: - head.append(idx + 1) - break - else: - heads.append([parent_node, idx + 1]) - return heads - - -def plot_planar_manipulator_robot( - robot: Robot, - config, - fig_handle=None, - ax_handle=None, - arm_color=(0.72, 0.53, 0.04, 1), - joint_color=(0.4, 0.4, 0.4, 1), - joint_size=20, - joint_thickness=16.0, - joint_edge_width=3.5, - alpha=1.0, -): - - if fig_handle is None: - fig_handle, ax_handle = plt.subplots() - - joint_x = [] - joint_y = [] - - for i in range(robot.n + 1): - joint_x += [robot.pose(config, "p" + str(i))[:-1, -1][0]] - joint_y += [robot.pose(config, "p" + str(i))[:-1, -1][1]] - joint_x = np.array(joint_x) - joint_y = np.array(joint_y) - - for joint in robot: - i1 = int(joint[1:]) - for child in robot.successors(joint): - i2 = int(child[1:]) - - ax_handle.plot( - [joint_x[i1], joint_x[i2]], - [joint_y[i1], joint_y[i2]], - "o-", - markersize=joint_size, - markerfacecolor=joint_color, - markeredgewidth=joint_edge_width, - markeredgecolor=(0, 0, 0, 1), - color=arm_color, - linewidth=joint_thickness, - alpha=alpha, - ) - - # plt.xlim([joint_x.min() - 0.125, joint_x.max() + 0.125]) - # plt.ylim([joint_y.min() - 0.125, joint_y.max() + 0.125]) - ax_handle.set_aspect("equal", adjustable="box") - - return True - - -def plot_planar_manipulator( - joint_x, - joint_y, - parent_nodes=None, - fig_handle=None, - ax_handle=None, - arm_color=(0.72, 0.53, 0.04, 1), - joint_color=(0.4, 0.4, 0.4, 1), - joint_size=20, - joint_thickness=16.0, - joint_edge_width=3.5, - alpha=1.0, -): - if fig_handle is None: - fig_handle, ax_handle = plt.subplots() - # fig_handle = plt.figure() - # ax_handle =fig_handle.gca() - - if parent_nodes is not None: - heads = get_tree_paths(parent_nodes) - for head in heads: - print("Head: {:}".format(head)) - plt.plot( - joint_x[head], - joint_y[head], - "o-", - markersize=joint_size, - markerfacecolor=joint_color, - markeredgewidth=joint_edge_width, - markeredgecolor=(0, 0, 0, 1), - color=arm_color, - linewidth=joint_thickness, - alpha=alpha, - ) - else: - plt.plot( - joint_x, - joint_y, - "o-", - markersize=joint_size, - markerfacecolor=joint_color, - markeredgewidth=joint_edge_width, - markeredgecolor=(0, 0, 0, 1), - color=arm_color, - linewidth=joint_thickness, - alpha=alpha, - ) - plt.xlim([joint_x.min() - 0.125, joint_x.max() + 0.125]) - plt.ylim([joint_y.min() - 0.125, joint_y.max() + 0.125]) - ax_handle.set_aspect("equal", adjustable="box") - return fig_handle, ax_handle - - -def add_fixed_point_to_graphic(fig, ax, joint_location=(0.0, 0.0), orientation=0): - x, y = (joint_location[0], joint_location[1] - 0.042) - plot_image("assets/fixed_point.png", x, y, ax, zoom=0.13) - - -def plot_3d_chain_manipulator(robot: Robot, input_angles, fig=None, ax=None): - """ - TODO: add tree capability (same as 2D) - :param robot: - :param input_angles: - :param title: - :return: - """ - pose_list = [] - for joint in input_angles: - pose_list.append(robot.get_pose(input_angles, joint)) - x = [] - y = [] - z = [] - # For frame axes - u1 = [] - v1 = [] - w1 = [] - u2 = [] - v2 = [] - w2 = [] - u3 = [] - v3 = [] - w3 = [] - if fig is None: - fig = plt.figure() - if ax is None: - ax = fig.gca(projection="3d") - # ax.set_aspect('equal', adjustable='box') - ax.autoscale(False) - for pose in pose_list: - R = pose[:3, :3] - t = pose[:3, 3] - print("t: {:}".format(t)) - print("type of t: {:}".format(type(t))) - x.append(t[0]) - y.append(t[1]) - z.append(t[2]) - u1.append(R[0, 0]) - v1.append(R[1, 0]) - w1.append(R[2, 0]) - u2.append(R[0, 1]) - v2.append(R[1, 1]) - w2.append(R[2, 1]) - u3.append(R[0, 2]) - v3.append(R[1, 2]) - w3.append(R[2, 2]) - ax.quiver(x, y, z, u1, v1, w1, length=0.1, normalize=True, color="r") - ax.quiver(x, y, z, u2, v2, w2, length=0.1, normalize=True, color="g") - ax.quiver(x, y, z, u3, v3, w3, length=0.1, normalize=True, color="b") - ax.plot(x, y, z) - - return fig, ax - - -if __name__ == "__main__": - # # Test manipulator plotting - # joint_y = np.array([0., 1., 1., 2., 2.]) - # joint_x = np.array([0., 0., 1., 0., 1.]) - # parent_joints = [0, 1, 1, 3] - # plot_planar_manipulator(joint_x, joint_y, parent_nodes=parent_joints) - # # add_fixed_point_to_graphic(fig, ax, joint_location=(x_joint[0], y_joint[0])) - # plt.show() - - # Test plot_heatmap - x = np.linspace(0.0, 1.0, 50) - y = x - X, Y = np.meshgrid(x, y) - Z = x ** 2 + y * x - 2 * y + 2.0 - Z = X ** 2 + Y * X - 2 * Y + 2.0 - plot_heatmap(X, Y, Z, title="Test Computer Modern Font") - plt.show() diff --git a/graphik/utils/roboturdf.py b/graphik/utils/roboturdf.py index 3600bcb..5be7a54 100644 --- a/graphik/utils/roboturdf.py +++ b/graphik/utils/roboturdf.py @@ -35,29 +35,14 @@ def __init__(self, fname): self._joints = list(self.urdf.robot.joints) self._links = list(self.urdf.robot.links) - self.urdf_ind_to_q, self.q_to_urdf_ind = self.joint_map() + self.q_to_urdf_ind = { + f"p{q_ind}": self._joints.index(joint) + for q_ind, joint in enumerate(self.urdf.actuated_joints, start=1) + } self.n_q_joints = len(self.q_to_urdf_ind) - self.n_urdf_joints = len(self.urdf_ind_to_q) self.ee_joints = self.find_end_effector_joints() self.T_zero = self.extract_T_zero_from_URDF(frame="joint") - self.scene = None - - self.parents = None - - def joint_map(self): - urdf_ind_to_q = {} - q_to_urdf_ind = {} - q_ind = 1 - label = "p{0}" - - for joint in self.urdf.actuated_joints: - j = self._joints.index(joint) - urdf_ind_to_q[j] = label.format(q_ind) - q_to_urdf_ind[label.format(q_ind)] = j - q_ind += 1 - - return urdf_ind_to_q, q_to_urdf_ind def find_first_joint(self): """ @@ -81,14 +66,6 @@ def find_actuated_joints_with_parent_link(self, link): return parent_joints - def find_joints_with_parent_link(self, link): - parent_joints = [] - for joint in self._joints: - if joint.parent == link.name: - parent_joints.append(joint) - - return parent_joints - def find_joints_actuated_child_joints(self, joint): child_link = self.find_link_by_name(joint.child) children_joints = self.find_actuated_joints_with_parent_link(child_link) @@ -104,24 +81,18 @@ def find_joints_child_joints_from_list(self, joint, joints): return children_joints def get_parents(self, joints): + """Map each joint's graphik label to the labels of its child joints.""" base_joint = self.find_first_joint() - if not (base_joint in joints): + if base_joint not in joints: raise ValueError("Base joint not in joints") - label_base = "p{0}" - # parents = {'p0': [label_base.format(joints.index(base_joint))]} parents = {} - for joint in joints: children = self.find_joints_child_joints_from_list(joint, joints) - if children == []: - child_labels = [] - else: - child_labels = [label_base.format(joints.index(cj)) for cj in children] - parent_label = label_base.format(joints.index(joint)) - parents[parent_label] = child_labels - - self.parents = parents + parents[f"p{joints.index(joint)}"] = [ + f"p{joints.index(cj)}" for cj in children + ] + return parents def actuated_joint_index(self, joint): try: @@ -135,12 +106,6 @@ def find_link_by_name(self, name): return link return None - def find_joint_by_name(self, name): - for joint in self._joints: - if joint.name == name: - return joint - return None - def extract_T_zero_from_URDF(self, q=None, frame="joint"): """ T is located at the joint's origin, the rotation such that @@ -192,8 +157,6 @@ def find_end_effector_joints(self): if child_joints == []: ee_joints.append(joint) - self.ee_joints = ee_joints - return ee_joints def map_to_urdf_ind(self, q): @@ -210,48 +173,13 @@ def map_to_urdf_ind(self, q): return urdf_q - def joint_limits(self): - ub = {} - lb = {} - - for joint in self.urdf.actuated_joints: - ubi = np.clip(joint.limit.upper, -np.pi, np.pi) - lbi = np.clip(joint.limit.lower, -np.pi, np.pi) - label = "p{0}" - joint_label = label.format(self.actuated_joint_index(joint)) - ub[joint_label] = ubi - lb[joint_label] = lbi - - return ub, lb - - def get_graphik_labels(self, joints): - """ - Assigned joint labels according to the p{i}, where i is the joint - index. The first joint has label p0 - Parameters - ---------- - joints : list - list of URDF Joint objects - - Returns - ------- - labels : list - list of the labels - """ - n = len(joints) - - label = "p{0}" - labels = [label.format(i) for i in range(n)] - return labels - def make_Revolute3d(self, ub, lb, randomized_links = False, randomize_percentage = 0.4): # if all the child lists have len 1, then chain, otherwise tree params = {} # assign parents joints = list(self.T_zero.keys()) - self.get_parents(joints) - params["parents"] = self.parents + params["parents"] = self.get_parents(joints) T_list = list(self.T_zero.values()) if randomized_links: @@ -267,8 +195,8 @@ def make_Revolute3d(self, ub, lb, randomized_links = False, randomize_percentage T_mod[idx + 1] = T_mod[idx] @ T_delta T_list = T_mod - # Assign Transforms - T_labels = self.get_graphik_labels(joints) + # Assign Transforms, labelled p0..p{n} in joint order + T_labels = [f"p{idx}" for idx in range(len(joints))] T_zero = dict(zip(T_labels, T_list)) T0 = T_zero["p0"] for key, val in T_zero.items(): @@ -311,79 +239,54 @@ def get_T_from_joint_axis(axis: np.ndarray): return T -def load_schunk_lwa4p(limits=None, randomized_links = False, randomize_percentage = 0.4): - fname = graphik.__path__[0] + "/robots/urdfs/lwa4p.urdf" - urdf_robot = RobotURDF(fname) - n = urdf_robot.n_q_joints - if limits is None: - ub = np.ones(n) * np.pi - lb = -ub - else: - lb = limits[0] - ub = limits[1] - robot = urdf_robot.make_Revolute3d(ub, lb, randomized_links, randomize_percentage) # make the Revolute class from a URDF - graph = ProblemGraph(robot) - return robot, graph +# Bundled URDFs, keyed by the short names accepted by load_urdf_robot. +URDF_FILES = { + "lwa4p": "lwa4p.urdf", + "lwa4d": "lwa4d.urdf", + "kuka": "kuka_iiwr.urdf", + "panda": "panda_arm.urdf", + "ur10": "ur10_mod.urdf", +} -def load_schunk_lwa4d(limits=None, randomized_links = False, randomize_percentage = 0.4): - fname = graphik.__path__[0] + "/robots/urdfs/lwa4d.urdf" +def load_urdf_robot(name, limits=None, randomized_links=False, randomize_percentage=0.4): + """Load a bundled URDF by short name (see ``URDF_FILES``) or an explicit + path, returning the ``(Robot, ProblemGraph)`` pair. + + ``limits`` is an optional ``(lb, ub)`` pair of per-joint arrays; the + default is symmetric +/-pi limits on every joint. + """ + fname = name + if name in URDF_FILES: + fname = graphik.__path__[0] + "/robots/urdfs/" + URDF_FILES[name] urdf_robot = RobotURDF(fname) - n = urdf_robot.n_q_joints if limits is None: - ub = np.ones(n) * np.pi + ub = np.ones(urdf_robot.n_q_joints) * np.pi lb = -ub else: - lb = limits[0] - ub = limits[1] - robot = urdf_robot.make_Revolute3d(ub, lb, randomized_links, randomize_percentage) # make the Revolute class from a URDF - graph = ProblemGraph(robot) - return robot, graph + lb, ub = limits + robot = urdf_robot.make_Revolute3d(ub, lb, randomized_links, randomize_percentage) + return robot, ProblemGraph(robot) -def load_kuka(limits=None, randomized_links = False, randomize_percentage = 0.4): - fname = graphik.__path__[0] + "/robots/urdfs/kuka_iiwr.urdf" - urdf_robot = RobotURDF(fname) - n = urdf_robot.n_q_joints - if limits is None: - ub = np.ones(n) * np.pi - lb = -ub - else: - lb = limits[0] - ub = limits[1] - robot = urdf_robot.make_Revolute3d(ub, lb, randomized_links, randomize_percentage) # make the Revolute class from a URDF - graph = ProblemGraph(robot) - return robot, graph +def load_schunk_lwa4p(limits=None, randomized_links=False, randomize_percentage=0.4): + return load_urdf_robot("lwa4p", limits, randomized_links, randomize_percentage) -def load_panda(limits=None, randomized_links = False, randomize_percentage = 0.4): - fname = graphik.__path__[0] + "/robots/urdfs/panda_arm.urdf" - urdf_robot = RobotURDF(fname) - n = urdf_robot.n_q_joints - if limits is None: - ub = np.ones(n) * np.pi - lb = -ub - else: - lb = limits[0] - ub = limits[1] - robot = urdf_robot.make_Revolute3d(ub, lb, randomized_links, randomize_percentage) # make the Revolute class from a URDF - # print(robot.structure.nodes()) - graph = ProblemGraph(robot) - return robot, graph -def load_ur10(limits=None, randomized_links = False, randomize_percentage = 0.4): - fname = graphik.__path__[0] + "/robots/urdfs/ur10_mod.urdf" - urdf_robot = RobotURDF(fname) - n = urdf_robot.n_q_joints - if limits is None: - ub = np.ones(n) * np.pi - lb = -ub - else: - lb = limits[0] - ub = limits[1] - robot = urdf_robot.make_Revolute3d(ub, lb, randomized_links, randomize_percentage) # make the Revolute class from a URDF - # print(robot.structure.nodes()) - graph = ProblemGraph(robot) - return robot, graph +def load_schunk_lwa4d(limits=None, randomized_links=False, randomize_percentage=0.4): + return load_urdf_robot("lwa4d", limits, randomized_links, randomize_percentage) + + +def load_kuka(limits=None, randomized_links=False, randomize_percentage=0.4): + return load_urdf_robot("kuka", limits, randomized_links, randomize_percentage) + + +def load_panda(limits=None, randomized_links=False, randomize_percentage=0.4): + return load_urdf_robot("panda", limits, randomized_links, randomize_percentage) + + +def load_ur10(limits=None, randomized_links=False, randomize_percentage=0.4): + return load_urdf_robot("ur10", limits, randomized_links, randomize_percentage) def load_truncated_ur10(n: int, limits=None): diff --git a/graphik/utils/utils.py b/graphik/utils/utils.py index 4b22b2e..2810f28 100644 --- a/graphik/utils/utils.py +++ b/graphik/utils/utils.py @@ -56,18 +56,3 @@ def normalize(v): if norm == 0: return v return v / norm - - -def table_environment(height=0.9, width=0.8, n_height=9, n_width=8, obs_inflation=2.): - - # Make the tabletop - radius = 0.5*height/n_height - tabletop_obs = [(np.asarray([2*(i+0.5)*radius, 2*(j+0.5)*radius, height+radius]), obs_inflation*radius) - for i in range(-n_width//2, n_width//2) for j in range(-n_width//2, n_width//2)] - - leg1 = [(np.asarray([-width/2+radius, -width/2+radius, (2*i+1)*radius]), obs_inflation*radius) for i in range(0, n_height)] - leg2 = [(np.asarray([-width/2+radius, width/2-radius, (2*i+1)*radius]), obs_inflation*radius) for i in range(0, n_height)] - leg3 = [(np.asarray([ width/2-radius, -width/2+radius, (2*i+1)*radius]), obs_inflation*radius) for i in range(0, n_height)] - leg4 = [(np.asarray([ width/2-radius, width/2-radius, (2*i+1)*radius]), obs_inflation*radius) for i in range(0, n_height)] - - return tabletop_obs + leg1 + leg2 + leg3 + leg4 diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..e45ab11 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,37 @@ +[build-system] +requires = ["setuptools>=64"] +build-backend = "setuptools.build_meta" + +[project] +name = "graphIK" +version = "0.2.0" +description = "Graph-based kinematics library" +authors = [ + { name = "Filip Maric", email = "filip.maric@robotics.utias.utoronto.ca" }, + { name = "Matthew Giamou", email = "matthew.giamou@robotics.utias.utoronto.ca" }, +] +license = { text = "MIT" } +requires-python = ">=3.11" +dependencies = [ + "numpy >= 2", + "scipy >= 1.17", + "matplotlib >= 3.1", + "cvxpy >= 1.8", + "networkx >= 3.6", + "yourdfpy >= 0.0.60", + "trimesh >= 4.6.1", + "pymlg @ git+https://github.com/decargroup/pymlg@8e6dc5ea61327ddfc2c8c1d16f276ae829f22db8#egg=pymlg", +] + +[project.optional-dependencies] +dev = ["pytest"] + +[project.urls] +Homepage = "https://github.com/utiasSTARS/GraphIK" + +[tool.setuptools.packages.find] +include = ["graphik", "graphik.*"] +exclude = ["tests*", "experiments*"] + +[tool.setuptools.package-data] +graphik = ["robots/urdfs/*", "robots/urdfs/meshes/**/*"] diff --git a/setup.py b/setup.py deleted file mode 100644 index d60f8d8..0000000 --- a/setup.py +++ /dev/null @@ -1,30 +0,0 @@ -from setuptools import setup, find_packages - -setup( - name="graphIK", - version="0.2.0", - description="Graph-based kinematics library", - author="Filip Maric, Matthew Giamou", - author_email="filip.maric@robotics.utias.utoronto.ca, matthew.giamou@robotics.utias.utoronto.ca", - license="MIT", - url="https://github.com/utiasSTARS/GraphIK", - packages=find_packages(include=["graphik", "graphik.*"], exclude=["tests*", "experiments*"]), - package_data={ - 'graphik': ['robots/urdfs/*', 'robots/urdfs/meshes/**/*'], - }, - include_package_data=True, - install_requires=[ - "numpy >= 2", - "scipy >= 1.17", - "matplotlib >= 3.1", - "cvxpy >= 1.8", - "networkx >= 3.6", - "yourdfpy >= 0.0.60", - "trimesh >= 4.6.1", - "pymlg @ git+https://github.com/decargroup/pymlg@8e6dc5ea61327ddfc2c8c1d16f276ae829f22db8#egg=pymlg", - ], - extras_require={ - "dev": ["pytest"], - }, - python_requires=">=3.11", -) diff --git a/tests/baselines/runner.py b/tests/baselines/runner.py index c53c426..09b6426 100644 --- a/tests/baselines/runner.py +++ b/tests/baselines/runner.py @@ -12,7 +12,7 @@ from graphik.solvers import RiemannianSolver, ScipySolver from graphik.utils.roboturdf import load_schunk_lwa4d, load_ur10 -from graphik.utils.utils import table_environment +from graphik.utils.environments import table_environment from tests.baselines.cases import Case ROBOT_LOADERS = {