From 98091e4b745a8fb9e715445ca90d172f9f62ce83 Mon Sep 17 00:00:00 2001 From: Chung Min Kim Date: Sat, 7 Feb 2026 18:06:07 -0800 Subject: [PATCH 01/11] Add J-PARSE velocity IK implementation with example and tests --- examples/14_jparse_ik.py | 143 ++++++++++++ examples/pyroki_snippets/__init__.py | 1 + examples/pyroki_snippets/_jparse.py | 271 ++++++++++++++++++++++ tests/test_jparse.py | 327 +++++++++++++++++++++++++++ 4 files changed, 742 insertions(+) create mode 100644 examples/14_jparse_ik.py create mode 100644 examples/pyroki_snippets/_jparse.py create mode 100644 tests/test_jparse.py diff --git a/examples/14_jparse_ik.py b/examples/14_jparse_ik.py new file mode 100644 index 00000000..60cd153d --- /dev/null +++ b/examples/14_jparse_ik.py @@ -0,0 +1,143 @@ +"""J-PARSE Velocity IK + +Singularity-aware velocity IK using J-PARSE with PyRoKi. + +J-PARSE (Jacobian-based Projection Algorithm for Resolving Singularities +Effectively) provides smooth behavior near singularities by decomposing +the Jacobian into non-singular and singular directions. + +Compare J-PARSE against standard pseudo-inverse and damped least squares +methods using the GUI controls. +""" + +import time + +import numpy as np +import pyroki as pk +import viser +from robot_descriptions.loaders.yourdfpy import load_robot_description +from viser.extras import ViserUrdf + +from pyroki_snippets._jparse import jparse_step + + +def main(): + """Main function for J-PARSE velocity IK example.""" + + urdf = load_robot_description("panda_description") + target_link_name = "panda_hand" + + # Create robot. + robot = pk.Robot.from_urdf(urdf) + target_link_index = robot.links.names.index(target_link_name) + + # Initial configuration (middle of joint range). + cfg = np.array((robot.joints.lower_limits + robot.joints.upper_limits) / 2.0) + + # Set up visualizer. + server = viser.ViserServer() + server.scene.add_grid("/ground", width=2, height=2) + urdf_vis = ViserUrdf(server, urdf, root_node_name="/base") + + # Target gizmo. + ik_target = server.scene.add_transform_controls( + "/ik_target", scale=0.2, position=(0.61, 0.0, 0.56), wxyz=(0, 0, 1, 0) + ) + + # --- GUI controls --- + with server.gui.add_folder("Method"): + method_handle = server.gui.add_dropdown( + "IK Method", + options=["jparse", "pinv", "dls"], + initial_value="jparse", + ) + gamma_handle = server.gui.add_slider( + "Gamma (singularity threshold)", 0.01, 0.5, 0.01, 0.1 + ) + singular_gain_handle = server.gui.add_slider( + "Singular direction gain", 0.5, 5.0, 0.1, 2.0 + ) + + with server.gui.add_folder("Control"): + position_gain_handle = server.gui.add_slider( + "Position gain", 1.0, 20.0, 0.5, 5.0 + ) + nullspace_gain_handle = server.gui.add_slider( + "Nullspace gain", 0.0, 2.0, 0.05, 0.5 + ) + max_vel_handle = server.gui.add_slider("Max joint velocity", 0.5, 5.0, 0.1, 2.0) + + with server.gui.add_folder("Visualization"): + show_ellipsoid_handle = server.gui.add_checkbox("Show ellipsoid", True) + + with server.gui.add_folder("Metrics"): + manipulability_handle = server.gui.add_number( + "Manipulability", 0.0, disabled=True + ) + condition_handle = server.gui.add_number( + "Inv. condition number", 0.0, disabled=True + ) + max_vel_display = server.gui.add_number( + "Max joint vel (rad/s)", 0.0, disabled=True + ) + pos_error_handle = server.gui.add_number( + "Position error (m)", 0.0, disabled=True + ) + singularity_handle = server.gui.add_text( + "Singularity status", "", disabled=True + ) + + # Manipulability ellipsoid visualization. + manip_ellipse = pk.viewer.ManipulabilityEllipse( + server, + robot, + root_node_name="/manipulability", + target_link_name=target_link_name, + ) + + # Control loop at ~50Hz. + while True: + start_time = time.time() + + # Take one velocity IK step. + cfg, info = jparse_step( + robot=robot, + cfg=cfg, + target_link_index=target_link_index, + target_position=np.array(ik_target.position), + method=method_handle.value, # type: ignore + gamma=gamma_handle.value, + singular_direction_gain=singular_gain_handle.value, + position_gain=position_gain_handle.value, + nullspace_gain=nullspace_gain_handle.value, + max_joint_velocity=max_vel_handle.value, + dt=0.02, + ) + + # Update metrics. + manipulability_handle.value = round(info["manipulability"], 4) + condition_handle.value = round(info["inverse_condition_number"], 4) + max_vel_display.value = round(info["max_joint_vel"], 3) + pos_error_handle.value = round(info["position_error"], 4) + + # Singularity warning. + icn = info["inverse_condition_number"] + if icn < gamma_handle.value: + singularity_handle.value = f"NEAR SINGULAR (icn={icn:.3f})" + else: + singularity_handle.value = "OK" + + # Update ellipsoid. + manip_ellipse.set_visibility(show_ellipsoid_handle.value) + manip_ellipse.update(cfg) + + # Update robot visualization. + urdf_vis.update_cfg(cfg) + + # Sleep to maintain ~50Hz. + elapsed = time.time() - start_time + time.sleep(max(0.0, 0.02 - elapsed)) + + +if __name__ == "__main__": + main() diff --git a/examples/pyroki_snippets/__init__.py b/examples/pyroki_snippets/__init__.py index d6051f93..95168631 100644 --- a/examples/pyroki_snippets/__init__.py +++ b/examples/pyroki_snippets/__init__.py @@ -9,3 +9,4 @@ from ._solve_ik_with_multiple_targets import ( solve_ik_with_multiple_targets as solve_ik_with_multiple_targets, ) +from ._jparse import jparse_step as jparse_step diff --git a/examples/pyroki_snippets/_jparse.py b/examples/pyroki_snippets/_jparse.py new file mode 100644 index 00000000..8e01b56d --- /dev/null +++ b/examples/pyroki_snippets/_jparse.py @@ -0,0 +1,271 @@ +""" +JAX implementation of the J-PARSE algorithm and velocity IK controller. + +J-PARSE (Jacobian-based Projection Algorithm for Resolving Singularities +Effectively) provides singularity-aware inverse kinematics by computing a +modified pseudo-inverse that handles singular configurations smoothly. + +Reference: https://github.com/chungmin99/jparse +""" + +from __future__ import annotations + +from typing import Literal + +import jax +import jax.numpy as jnp +import jaxlie +import numpy as np +import pyroki as pk +from jax.typing import ArrayLike + + +def compute_jacobian( + robot: pk.Robot, + cfg: ArrayLike, + target_link_index: int, + position_only: bool = True, +) -> jnp.ndarray: + """Compute geometric Jacobian via autodiff on pyroki FK. + + Args: + robot: PyRoKi robot model. + cfg: Joint configuration (actuated_count,). + target_link_index: Index of the target link. + position_only: If True, return 3xn position Jacobian. + If False, return 6xn Jacobian (via SE3 log). + + Returns: + Jacobian matrix. Shape (3, n) if position_only else (6, n). + """ + cfg = jnp.asarray(cfg) + + if position_only: + jacobian = jax.jacfwd( + lambda q: jaxlie.SE3(robot.forward_kinematics(q)).translation() + )(cfg)[target_link_index] + else: + + def get_pose_log(q: ArrayLike) -> jnp.ndarray: + poses = robot.forward_kinematics(q) + return jaxlie.SE3(poses[target_link_index]).log() + + jacobian = jax.jacfwd(get_pose_log)(cfg) + + return jacobian + + +def jparse_pseudoinverse( + jacobian: ArrayLike, + gamma: float = 0.1, + singular_direction_gain: float = 1.0, +) -> tuple[jnp.ndarray, jnp.ndarray]: + """Compute J-PARSE pseudo-inverse of a Jacobian matrix. + + The J-PARSE algorithm decomposes the Jacobian using SVD and constructs + a modified pseudo-inverse that: + 1. Clamps singular values below gamma*sigma_max (safety Jacobian) + 2. Projects commands onto non-singular directions (projection Jacobian) + 3. Provides smooth feedback in singular directions + + Formula: J_parse = J_safety^+ @ J_proj @ J_proj^+ + J_safety^+ @ Phi + + Args: + jacobian: The m x n Jacobian matrix. + gamma: Singularity threshold in (0, 1). Directions with + sigma/sigma_max < gamma are treated as singular. + singular_direction_gain: Gain applied to singular directions. + + Returns: + Tuple of (J_parse, nullspace_projector): + - J_parse: n x m J-PARSE pseudo-inverse matrix. + - nullspace_projector: n x n nullspace projection matrix. + """ + J = jnp.asarray(jacobian) + m, n = J.shape + + U, S, Vt = jnp.linalg.svd(J, full_matrices=True) + k = S.shape[0] # min(m, n) + + sigma_max = jnp.max(S) + threshold = gamma * sigma_max + + # Mask: True for non-singular directions (sigma > threshold). + non_singular = S > threshold + + # Safety singular values: clamp below threshold. + S_safety = jnp.where(non_singular, S, threshold) + + # Projection singular values: keep only non-singular directions. + S_proj = jnp.where(non_singular, S, 0.0) + + # Reconstruct matrices from SVD components. + # J_safety = U[:, :k] @ diag(S_safety) @ Vt[:k, :] + U_k = U[:, :k] + Vt_k = Vt[:k, :] + + J_safety = U_k * S_safety[None, :] @ Vt_k + J_proj = U_k * S_proj[None, :] @ Vt_k + + J_safety_pinv = jnp.linalg.pinv(J_safety) + J_proj_pinv = jnp.linalg.pinv(J_proj) + + # Singular direction feedback: Phi = sum_i phi_i * u_i @ u_i^T @ Kp + # where phi_i = sigma_i / (sigma_max * gamma) for singular directions. + phi = jnp.where(non_singular, 0.0, S / (sigma_max * gamma)) + # Phi_singular = U_k @ diag(phi) @ U_k^T @ Kp + Kp = singular_direction_gain * jnp.eye(m) + Phi_singular = (U_k * phi[None, :]) @ U_k.T @ Kp + + J_parse = J_safety_pinv @ J_proj @ J_proj_pinv + J_safety_pinv @ Phi_singular + + # Nullspace projector: N = I - J_safety^+ @ J_safety + nullspace = jnp.eye(n) - J_safety_pinv @ J_safety + + return J_parse, nullspace + + +def pinv(jacobian: ArrayLike) -> jnp.ndarray: + """Standard Moore-Penrose pseudo-inverse (for comparison).""" + return jnp.linalg.pinv(jnp.asarray(jacobian)) + + +def damped_least_squares( + jacobian: ArrayLike, + damping: float = 0.05, +) -> jnp.ndarray: + """Damped least squares (Levenberg-Marquardt) pseudo-inverse. + + Args: + jacobian: The m x n Jacobian matrix. + damping: Damping factor lambda. + + Returns: + The n x m DLS pseudo-inverse. + """ + J = jnp.asarray(jacobian) + n = J.shape[1] + return jnp.linalg.inv(J.T @ J + damping**2 * jnp.eye(n)) @ J.T + + +def manipulability_measure(jacobian: ArrayLike) -> jnp.ndarray: + """Yoshikawa's manipulability measure: sqrt(det(J @ J^T)).""" + J = jnp.asarray(jacobian) + return jnp.sqrt(jnp.linalg.det(J @ J.T)) + + +def inverse_condition_number(jacobian: ArrayLike) -> jnp.ndarray: + """Inverse condition number: sigma_min / sigma_max.""" + J = jnp.asarray(jacobian) + S = jnp.linalg.svd(J, compute_uv=False) + return jnp.min(S) / jnp.max(S) + + +def jparse_step( + robot: pk.Robot, + cfg: ArrayLike, + target_link_index: int, + target_position: ArrayLike, + *, + method: Literal["jparse", "pinv", "dls"] = "jparse", + gamma: float = 0.1, + singular_direction_gain: float = 1.0, + position_gain: float = 5.0, + nullspace_gain: float = 0.5, + max_joint_velocity: float = 2.0, + dls_damping: float = 0.05, + dt: float = 0.02, + home_cfg: ArrayLike | None = None, +) -> tuple[np.ndarray, dict]: + """Single velocity IK step using J-PARSE (or pinv/DLS). + + Computes Jacobian -> pseudo-inverse -> joint velocities -> integrate -> clamp. + + Args: + robot: PyRoKi robot model. + cfg: Current joint configuration (actuated_count,). + target_link_index: Index of the target link. + target_position: Target position (3,). + method: IK method — "jparse", "pinv", or "dls". + gamma: J-PARSE singularity threshold. + singular_direction_gain: Gain for singular directions. + position_gain: Proportional gain for position error. + nullspace_gain: Gain for nullspace motion toward home. + max_joint_velocity: Maximum joint velocity (rad/s). + dls_damping: Damping factor for DLS method. + dt: Time step for integration. + home_cfg: Home configuration for nullspace. If None, uses joint midpoints. + + Returns: + Tuple of (new_cfg, info): + - new_cfg: New joint configuration after integration (numpy). + - info: Dict with position_error, max_joint_vel, jacobian, + manipulability, inverse_condition_number. + """ + cfg = jnp.asarray(cfg) + target_position = jnp.asarray(target_position) + + # Current end-effector position. + poses = robot.forward_kinematics(cfg) + current_pos = jaxlie.SE3(poses[target_link_index]).translation() + + # Position error and desired task-space velocity. + pos_error = target_position - current_pos + pos_error_mag = float(jnp.linalg.norm(pos_error)) + v_des = position_gain * pos_error + + # Compute Jacobian. + jacobian = compute_jacobian(robot, cfg, target_link_index, position_only=True) + + # Compute pseudo-inverse and nullspace. + if method == "jparse": + J_inv, N = jparse_pseudoinverse(jacobian, gamma, singular_direction_gain) + elif method == "pinv": + J_inv = pinv(jacobian) + N = jnp.eye(jacobian.shape[1]) - J_inv @ jacobian + else: # dls + J_inv = damped_least_squares(jacobian, dls_damping) + N = jnp.eye(jacobian.shape[1]) - J_inv @ jacobian + + # Primary task joint velocities. + dq = J_inv @ v_des + + # Nullspace motion toward home configuration. + if nullspace_gain > 0: + if home_cfg is None: + lower = robot.joints.lower_limits + upper = robot.joints.upper_limits + home = (lower + upper) / 2.0 + else: + home = jnp.asarray(home_cfg) + dq_null = N @ (-nullspace_gain * (cfg - home)) + dq = dq + dq_null + + # Track raw max velocity before limiting. + max_joint_vel = float(jnp.max(jnp.abs(dq))) + + # Apply velocity limits. + scale = jnp.where( + jnp.max(jnp.abs(dq)) > max_joint_velocity, + max_joint_velocity / jnp.max(jnp.abs(dq)), + 1.0, + ) + dq = dq * scale + + # Integrate. + new_cfg = cfg + dq * dt + + # Clamp to joint limits. + lower = robot.joints.lower_limits + upper = robot.joints.upper_limits + new_cfg = jnp.clip(new_cfg, lower, upper) + + info = { + "position_error": pos_error_mag, + "max_joint_vel": max_joint_vel, + "jacobian": np.asarray(jacobian), + "manipulability": float(manipulability_measure(jacobian)), + "inverse_condition_number": float(inverse_condition_number(jacobian)), + } + + return np.asarray(new_cfg), info diff --git a/tests/test_jparse.py b/tests/test_jparse.py new file mode 100644 index 00000000..343b2cd7 --- /dev/null +++ b/tests/test_jparse.py @@ -0,0 +1,327 @@ +"""Tests for the JAX J-PARSE implementation. + +Verifies correctness of the JAX port and cross-validates against the +original numpy-based jparse_robotics package. +""" + +import sys + +import numpy as np +import pytest + +sys.path.insert(0, "examples") + +from pyroki_snippets._jparse import ( + damped_least_squares, + inverse_condition_number, + jparse_pseudoinverse, + jparse_step, + manipulability_measure, + pinv, +) + + +class TestJParsePseudoinverse: + """Test the core jparse_pseudoinverse function.""" + + def test_well_conditioned_2dof(self): + """Well-conditioned 2-DOF planar Jacobian produces valid pseudo-inverse.""" + l1, l2 = 1.0, 1.0 + theta1, theta2 = np.pi / 4, np.pi / 4 + J = np.array( + [ + [ + -l1 * np.sin(theta1) - l2 * np.sin(theta1 + theta2), + -l2 * np.sin(theta1 + theta2), + ], + [ + l1 * np.cos(theta1) + l2 * np.cos(theta1 + theta2), + l2 * np.cos(theta1 + theta2), + ], + ] + ) + + J_parse, N = jparse_pseudoinverse(J, gamma=0.1) + + assert J_parse.shape == (2, 2) + # Valid pseudo-inverse: J @ J_parse @ J ≈ J + np.testing.assert_allclose(J @ J_parse @ J, J, atol=1e-6) + # Well-conditioned: should be close to pinv + np.testing.assert_allclose(J_parse, np.linalg.pinv(J), atol=0.1) + + def test_near_singular_no_nan(self): + """Near-singular Jacobian produces finite output.""" + J = np.array([[1.0, 1.0], [1.0, 1.001]]) + J_parse, N = jparse_pseudoinverse(J, gamma=0.1) + assert np.all(np.isfinite(J_parse)) + + def test_singular_no_nan(self): + """Exactly singular Jacobian produces finite output.""" + J = np.array([[1.0, 1.0], [1.0, 1.0]]) + J_parse, N = jparse_pseudoinverse(J, gamma=0.1) + assert np.all(np.isfinite(J_parse)) + + def test_redundant_system(self): + """Redundant 2x3 system has correct shapes and valid pseudo-inverse.""" + J = np.array([[1.0, 0.0, 1.0], [0.0, 1.0, 1.0]]) + J_parse, N = jparse_pseudoinverse(J, gamma=0.1) + assert J_parse.shape == (3, 2) + np.testing.assert_allclose(J @ J_parse @ J, J, atol=1e-6) + + +class TestNullspace: + """Test nullspace projector.""" + + def test_nullspace_shape(self): + """Nullspace projector has shape (n, n).""" + J = np.array([[1.0, 0.0, 1.0], [0.0, 1.0, 1.0]]) + _, N = jparse_pseudoinverse(J, gamma=0.1) + assert N.shape == (3, 3) + + def test_nullspace_projects_to_zero(self): + """Nullspace vectors project to ~zero in task space.""" + J = np.array([[1.0, 0.0, 1.0], [0.0, 1.0, 1.0]]) + _, N = jparse_pseudoinverse(J, gamma=0.1) + null_vec = N @ np.array([1.0, 1.0, 1.0]) + task_effect = J @ null_vec + np.testing.assert_allclose(task_effect, 0, atol=0.1) + + +class TestSingularDirectionGains: + """Test singular direction gain parameters.""" + + def test_finite_output_with_gain(self): + """Finite output with specified singular direction gain.""" + J = np.array([[1.0, 1.0], [1.0, 1.001]]) + J_parse, _ = jparse_pseudoinverse(J, gamma=0.1, singular_direction_gain=2.0) + assert np.all(np.isfinite(J_parse)) + + def test_gain_affects_output(self): + """Different gains produce different results for near-singular case.""" + J = np.array([[1.0, 1.0], [1.0, 1.001]]) + J1, _ = jparse_pseudoinverse(J, gamma=0.1, singular_direction_gain=1.0) + J2, _ = jparse_pseudoinverse(J, gamma=0.1, singular_direction_gain=3.0) + assert not np.allclose(J1, J2) + + +class TestComparisonMethods: + """Test pinv and damped_least_squares comparison methods.""" + + def test_pinv_matches_numpy(self): + """pinv matches numpy's pinv.""" + J = np.array([[1.0, 0.5], [0.5, 1.0]]) + np.testing.assert_allclose(pinv(J), np.linalg.pinv(J), atol=1e-6) + + def test_dls_formula(self): + """DLS matches expected formula.""" + J = np.array([[1.0, 0.0], [0.0, 1.0]]) + d = 0.1 + J_dls = damped_least_squares(J, damping=d) + expected = np.linalg.inv(J.T @ J + d**2 * np.eye(2)) @ J.T + np.testing.assert_allclose(J_dls, expected, atol=1e-6) + + +class TestMetrics: + """Test manipulability and condition number metrics.""" + + def test_manipulability_identity(self): + """Identity Jacobian has manipulability 1.""" + J = np.eye(2) + m = manipulability_measure(J) + np.testing.assert_allclose(float(m), 1.0, atol=1e-6) + + def test_manipulability_near_singular_is_low(self): + """Near-singular has lower manipulability than identity.""" + J_good = np.eye(2) + J_bad = np.array([[1.0, 1.0], [1.0, 1.001]]) + assert float(manipulability_measure(J_bad)) < float( + manipulability_measure(J_good) + ) + + def test_inverse_condition_identity(self): + """Identity Jacobian has inverse condition number 1.""" + J = np.eye(2) + np.testing.assert_allclose(float(inverse_condition_number(J)), 1.0, atol=1e-6) + + def test_inverse_condition_range(self): + """Inverse condition number is in [0, 1].""" + rng = np.random.default_rng(42) + for _ in range(10): + J = rng.standard_normal((3, 4)) + icn = float(inverse_condition_number(J)) + assert 0 <= icn <= 1 + + +class TestEdgeCases: + """Test edge cases.""" + + def test_1x1_jacobian(self): + """1x1 Jacobian produces correct pseudo-inverse.""" + J = np.array([[2.0]]) + J_parse, _ = jparse_pseudoinverse(J, gamma=0.1) + assert J_parse.shape == (1, 1) + np.testing.assert_allclose(float(J_parse[0, 0]), 0.5, atol=1e-6) + + def test_tall_jacobian(self): + """Tall Jacobian (m > n) has correct shape.""" + J = np.array([[1.0], [0.5], [0.25]]) + J_parse, _ = jparse_pseudoinverse(J, gamma=0.1) + assert J_parse.shape == (1, 3) + + def test_very_small_gamma(self): + """Small gamma: well-conditioned case close to standard pinv.""" + J = np.array([[1.0, 0.5], [0.5, 1.0]]) + J_parse, _ = jparse_pseudoinverse(J, gamma=0.01) + np.testing.assert_allclose(J_parse, np.linalg.pinv(J), atol=0.01) + + def test_large_gamma(self): + """Large gamma still produces finite output.""" + J = np.array([[1.0, 0.5], [0.5, 1.0]]) + J_parse, _ = jparse_pseudoinverse(J, gamma=0.99) + assert np.all(np.isfinite(J_parse)) + + +class TestFeatureParity: + """Cross-validate JAX implementation against original jparse_robotics.""" + + @pytest.fixture(autouse=True) + def _skip_if_no_jparse(self): + jparse_robotics = pytest.importorskip("jparse_robotics") + self.JParseCore = jparse_robotics.JParseCore + self.np_manipulability = jparse_robotics.manipulability_measure + self.np_icn = jparse_robotics.inverse_condition_number + + @pytest.fixture() + def test_jacobians(self): + """Collection of test Jacobians.""" + rng = np.random.default_rng(42) + return { + "well_conditioned": np.array([[1.0, 0.5], [0.5, 1.0]]), + "near_singular": np.array([[1.0, 1.0], [1.0, 1.001]]), + "redundant": np.array([[1.0, 0.0, 1.0], [0.0, 1.0, 1.0]]), + "random_6x7": rng.standard_normal((6, 7)), + } + + def test_jparse_output_matches(self, test_jacobians): + """J-PARSE output matches original for multiple test cases.""" + gamma = 0.1 + core = self.JParseCore(gamma=gamma) + + for name, J in test_jacobians.items(): + J_np = core.compute(J) + J_jax, _ = jparse_pseudoinverse(J, gamma=gamma) + np.testing.assert_allclose( + np.array(J_jax), np.array(J_np), atol=1e-5, err_msg=f"Failed for {name}" + ) + + def test_manipulability_matches(self, test_jacobians): + """Manipulability measure matches original.""" + for name, J in test_jacobians.items(): + m_np = self.np_manipulability(J) + m_jax = float(manipulability_measure(J)) + # Relaxed tolerance: JAX defaults to float32 vs numpy float64. + np.testing.assert_allclose( + m_jax, m_np, rtol=0.2, atol=1e-4, err_msg=f"Failed for {name}" + ) + + def test_inverse_condition_matches(self, test_jacobians): + """Inverse condition number matches original.""" + for name, J in test_jacobians.items(): + icn_np = self.np_icn(J) + icn_jax = float(inverse_condition_number(J)) + np.testing.assert_allclose( + icn_jax, icn_np, atol=1e-6, err_msg=f"Failed for {name}" + ) + + def test_dls_matches(self, test_jacobians): + """DLS matches original.""" + core = self.JParseCore(gamma=0.1) + damping = 0.05 + for name, J in test_jacobians.items(): + J_np = core.damped_least_squares(J, damping=damping) + J_jax = damped_least_squares(J, damping=damping) + # Relaxed tolerance: JAX defaults to float32 vs numpy float64. + np.testing.assert_allclose( + np.array(J_jax), np.array(J_np), atol=1e-4, err_msg=f"Failed for {name}" + ) + + +class TestIntegration: + """Integration test with pyroki robot (no Viser).""" + + @pytest.fixture() + def panda_robot(self): + """Load Panda robot.""" + import pyroki as pk + from robot_descriptions.loaders.yourdfpy import load_robot_description + + urdf = load_robot_description("panda_description") + robot = pk.Robot.from_urdf(urdf) + target_link_index = robot.links.names.index("panda_hand") + return robot, target_link_index + + def test_position_error_decreases(self, panda_robot): + """Position error decreases over several IK steps.""" + robot, target_link_index = panda_robot + + # Start from joint midpoints. + cfg = np.array((robot.joints.lower_limits + robot.joints.upper_limits) / 2.0) + + # Target position (reachable). + target_pos = np.array([0.5, 0.1, 0.4]) + + errors = [] + for _ in range(50): + cfg, info = jparse_step( + robot=robot, + cfg=cfg, + target_link_index=target_link_index, + target_position=target_pos, + gamma=0.1, + dt=0.02, + ) + errors.append(info["position_error"]) + + # Error should decrease overall. + assert errors[-1] < errors[0] + + def test_joint_limits_respected(self, panda_robot): + """All configs remain within joint limits.""" + robot, target_link_index = panda_robot + cfg = np.array((robot.joints.lower_limits + robot.joints.upper_limits) / 2.0) + target_pos = np.array([0.5, 0.1, 0.4]) + + lower = np.array(robot.joints.lower_limits) + upper = np.array(robot.joints.upper_limits) + + for _ in range(30): + cfg, _ = jparse_step( + robot=robot, + cfg=cfg, + target_link_index=target_link_index, + target_position=target_pos, + dt=0.02, + ) + assert np.all(cfg >= lower - 1e-6) + assert np.all(cfg <= upper + 1e-6) + + def test_info_dict_keys(self, panda_robot): + """Info dict contains expected keys.""" + robot, target_link_index = panda_robot + cfg = np.array((robot.joints.lower_limits + robot.joints.upper_limits) / 2.0) + target_pos = np.array([0.5, 0.1, 0.4]) + + _, info = jparse_step( + robot=robot, + cfg=cfg, + target_link_index=target_link_index, + target_position=target_pos, + ) + expected_keys = { + "position_error", + "max_joint_vel", + "jacobian", + "manipulability", + "inverse_condition_number", + } + assert set(info.keys()) == expected_keys From 4157ad6c6a5170f9c25c4df35ca2e672cd088f3e Mon Sep 17 00:00:00 2001 From: Chung Min Kim Date: Sat, 7 Feb 2026 18:06:15 -0800 Subject: [PATCH 02/11] Add Claude Code project documentation --- CLAUDE.md | 74 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 74 insertions(+) create mode 100644 CLAUDE.md diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 00000000..c2c43911 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,74 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## Project Overview + +PyRoki is a JAX-based Python library for differentiable robot kinematics and optimization. It provides forward kinematics, collision detection, and cost/constraint factories that integrate with the `jaxls` nonlinear least-squares solver. Python 3.10+. + +## Common Commands + +### Install +```bash +sudo apt install -y libsuitesparse-dev # system dependency for scikit-sparse +pip install -e ".[dev]" +``` + +### Test +```bash +pytest # run all tests +pytest tests/test_ik_vmap.py # run single test file +pytest -k "test_name" # run specific test by name +``` + +### Lint & Format +```bash +ruff check docs/ src/ examples/ # lint +ruff format docs/ src/ examples/ # format +ruff format docs/ src/ examples/ --check # check formatting without modifying +``` + +### Type Check +```bash +pyright . +``` + +### Build Docs +```bash +cd docs && sphinx-build source build -b dirhtml +``` + +## Architecture + +### Core Pipeline + +The typical workflow is: **URDF → Robot → forward kinematics → residuals → costs → jaxls solver**. + +1. **`Robot`** (`_robot.py`): `jax_dataclasses.pytree_dataclass` holding joint/link info. Created via `Robot.from_urdf()`. The `forward_kinematics(cfg)` method returns SE(3) transforms (as wxyz-xyz arrays) for all links. Supports `jax.vmap` for batched configs. + +2. **Residuals** (`_residuals/`): Functions with signature `(vals: VarValues, robot, joint_var, ...) -> Array`. These compute optimization residuals (pose error, joint limits, collision distances, etc.). + +3. **Costs** (`costs.py`): Thin wrappers created via `Cost.factory(residual_fn)`. Costs are soft penalties (least squares); constraints use `kind="constraint_leq_zero"` for augmented Lagrangian enforcement. + +4. **Collision** (`collision/`): `RobotCollision.from_urdf()` builds collision geometry (spheres, capsules, boxes). `CollGeom` is the base class for geometry primitives. `collide()` computes signed distances between geometry pairs. + +5. **Solver**: Uses `jaxls` (external) for Levenberg-Marquardt optimization with optional augmented Lagrangian constraints. Problems are built with `jaxls.LeastSquaresProblem` and solved with `jaxls.solve()`. + +### Key Patterns + +- All core data structures are `@jdc.pytree_dataclass` (JAX-compatible pytrees via `jax_dataclasses`) +- SE(3) transforms use `jaxlie.SE3` (Lie group library) +- Typed arrays use `jaxtyping.Float[Array, "shape"]` annotations +- Variables for optimization use `jaxls.Var` subclasses; `Robot.joint_var_cls` is auto-created from URDF +- `@jdc.jit` is used instead of `@jax.jit` for methods on pytree dataclasses + +### Examples Structure + +Examples live in `examples/` and use helper modules from `examples/pyroki_snippets/` for reusable solve patterns (e.g., `solve_ik`, `solve_ik_with_collision`). Visualization uses `viser`. + +## CI + +Three workflows run on PRs to `main`: +- **pytest**: Tests across Python 3.10–3.13 +- **pyright**: Type checking across Python 3.10–3.13 +- **formatting**: `ruff check` + `ruff format --check` on `docs/`, `src/`, `examples/` From 8d7bc481ddfbbc4e00344e8be7162eeffdb2251d Mon Sep 17 00:00:00 2001 From: Chung Min Kim Date: Sat, 7 Feb 2026 22:29:59 -0800 Subject: [PATCH 03/11] Add orientation tracking support to J-PARSE velocity IK --- examples/14_jparse_ik.py | 23 ++++- examples/pyroki_snippets/_jparse.py | 140 +++++++++++++++++++++++---- tests/test_jparse.py | 145 +++++++++++++++++++++++++++- 3 files changed, 283 insertions(+), 25 deletions(-) diff --git a/examples/14_jparse_ik.py b/examples/14_jparse_ik.py index 60cd153d..3a89e294 100644 --- a/examples/14_jparse_ik.py +++ b/examples/14_jparse_ik.py @@ -54,14 +54,21 @@ def main(): gamma_handle = server.gui.add_slider( "Gamma (singularity threshold)", 0.01, 0.5, 0.01, 0.1 ) - singular_gain_handle = server.gui.add_slider( - "Singular direction gain", 0.5, 5.0, 0.1, 2.0 + singular_gain_position_handle = server.gui.add_slider( + "Singular gain (position)", 0.5, 5.0, 0.1, 2.0 + ) + singular_gain_angular_handle = server.gui.add_slider( + "Singular gain (angular)", 0.5, 5.0, 0.1, 2.0 ) with server.gui.add_folder("Control"): + track_orientation_handle = server.gui.add_checkbox("Track orientation", True) position_gain_handle = server.gui.add_slider( "Position gain", 1.0, 20.0, 0.5, 5.0 ) + orientation_gain_handle = server.gui.add_slider( + "Orientation gain", 0.1, 5.0, 0.1, 1.0 + ) nullspace_gain_handle = server.gui.add_slider( "Nullspace gain", 0.0, 2.0, 0.05, 0.5 ) @@ -83,6 +90,9 @@ def main(): pos_error_handle = server.gui.add_number( "Position error (m)", 0.0, disabled=True ) + ori_error_handle = server.gui.add_number( + "Orientation error (rad)", 0.0, disabled=True + ) singularity_handle = server.gui.add_text( "Singularity status", "", disabled=True ) @@ -100,15 +110,21 @@ def main(): start_time = time.time() # Take one velocity IK step. + target_wxyz = ( + np.array(ik_target.wxyz) if track_orientation_handle.value else None + ) cfg, info = jparse_step( robot=robot, cfg=cfg, target_link_index=target_link_index, target_position=np.array(ik_target.position), + target_wxyz=target_wxyz, method=method_handle.value, # type: ignore gamma=gamma_handle.value, - singular_direction_gain=singular_gain_handle.value, + singular_direction_gain_position=singular_gain_position_handle.value, + singular_direction_gain_angular=singular_gain_angular_handle.value, position_gain=position_gain_handle.value, + orientation_gain=orientation_gain_handle.value, nullspace_gain=nullspace_gain_handle.value, max_joint_velocity=max_vel_handle.value, dt=0.02, @@ -119,6 +135,7 @@ def main(): condition_handle.value = round(info["inverse_condition_number"], 4) max_vel_display.value = round(info["max_joint_vel"], 3) pos_error_handle.value = round(info["position_error"], 4) + ori_error_handle.value = round(info["orientation_error"], 4) # Singularity warning. icn = info["inverse_condition_number"] diff --git a/examples/pyroki_snippets/_jparse.py b/examples/pyroki_snippets/_jparse.py index 8e01b56d..7de64d8c 100644 --- a/examples/pyroki_snippets/_jparse.py +++ b/examples/pyroki_snippets/_jparse.py @@ -33,7 +33,7 @@ def compute_jacobian( cfg: Joint configuration (actuated_count,). target_link_index: Index of the target link. position_only: If True, return 3xn position Jacobian. - If False, return 6xn Jacobian (via SE3 log). + If False, return 6xn Jacobian (translation + SO3 log). Returns: Jacobian matrix. Shape (3, n) if position_only else (6, n). @@ -45,12 +45,21 @@ def compute_jacobian( lambda q: jaxlie.SE3(robot.forward_kinematics(q)).translation() )(cfg)[target_link_index] else: - - def get_pose_log(q: ArrayLike) -> jnp.ndarray: + # Compute R_anchor_inv from the current cfg *outside* the jacfwd + # closure so JAX treats it as a constant. At q = cfg the relative + # rotation R(q) @ R_anchor⁻¹ equals I, so differentiating its log + # yields the geometric (angular-velocity) Jacobian instead of the + # analytical Jacobian. + anchor_poses = robot.forward_kinematics(cfg) + R_anchor_inv = jaxlie.SE3(anchor_poses[target_link_index]).rotation().inverse() + + def get_pose_components(q: ArrayLike) -> jnp.ndarray: poses = robot.forward_kinematics(q) - return jaxlie.SE3(poses[target_link_index]).log() + pose = jaxlie.SE3(poses[target_link_index]) + relative_rotation = pose.rotation() @ R_anchor_inv + return jnp.concatenate([pose.translation(), relative_rotation.log()]) - jacobian = jax.jacfwd(get_pose_log)(cfg) + jacobian = jax.jacfwd(get_pose_components)(cfg) return jacobian @@ -58,7 +67,10 @@ def get_pose_log(q: ArrayLike) -> jnp.ndarray: def jparse_pseudoinverse( jacobian: ArrayLike, gamma: float = 0.1, - singular_direction_gain: float = 1.0, + singular_direction_gain_position: float = 1.0, + singular_direction_gain_angular: float = 1.0, + position_dimensions: int | None = None, + angular_dimensions: int | None = None, ) -> tuple[jnp.ndarray, jnp.ndarray]: """Compute J-PARSE pseudo-inverse of a Jacobian matrix. @@ -74,7 +86,12 @@ def jparse_pseudoinverse( jacobian: The m x n Jacobian matrix. gamma: Singularity threshold in (0, 1). Directions with sigma/sigma_max < gamma are treated as singular. - singular_direction_gain: Gain applied to singular directions. + singular_direction_gain_position: Gain for position dimensions. + singular_direction_gain_angular: Gain for angular dimensions. + position_dimensions: Number of position rows in the Jacobian. + If None and angular_dimensions is None, all rows use + singular_direction_gain_position. + angular_dimensions: Number of angular rows in the Jacobian. Returns: Tuple of (J_parse, nullspace_projector): @@ -84,6 +101,30 @@ def jparse_pseudoinverse( J = jnp.asarray(jacobian) m, n = J.shape + if position_dimensions is None and angular_dimensions is None: + pos_dims = m + ang_dims = 0 + else: + if position_dimensions is None or angular_dimensions is None: + raise ValueError( + "Both position_dimensions and angular_dimensions must be provided." + ) + if ( + not isinstance(position_dimensions, int) + or isinstance(position_dimensions, bool) + or not isinstance(angular_dimensions, int) + or isinstance(angular_dimensions, bool) + ): + raise ValueError("Dimension values must be integers.") + if position_dimensions < 0 or angular_dimensions < 0: + raise ValueError("Dimension values must be non-negative.") + if position_dimensions + angular_dimensions != m: + raise ValueError( + "position_dimensions + angular_dimensions must equal Jacobian row count." + ) + pos_dims = position_dimensions + ang_dims = angular_dimensions + U, S, Vt = jnp.linalg.svd(J, full_matrices=True) k = S.shape[0] # min(m, n) @@ -114,7 +155,13 @@ def jparse_pseudoinverse( # where phi_i = sigma_i / (sigma_max * gamma) for singular directions. phi = jnp.where(non_singular, 0.0, S / (sigma_max * gamma)) # Phi_singular = U_k @ diag(phi) @ U_k^T @ Kp - Kp = singular_direction_gain * jnp.eye(m) + singular_gains = jnp.concatenate( + [ + jnp.full((pos_dims,), singular_direction_gain_position), + jnp.full((ang_dims,), singular_direction_gain_angular), + ] + ) + Kp = jnp.diag(singular_gains) Phi_singular = (U_k * phi[None, :]) @ U_k.T @ Kp J_parse = J_safety_pinv @ J_proj @ J_proj_pinv + J_safety_pinv @ Phi_singular @@ -166,11 +213,14 @@ def jparse_step( cfg: ArrayLike, target_link_index: int, target_position: ArrayLike, + target_wxyz: ArrayLike | None = None, *, method: Literal["jparse", "pinv", "dls"] = "jparse", gamma: float = 0.1, - singular_direction_gain: float = 1.0, + singular_direction_gain_position: float = 1.0, + singular_direction_gain_angular: float = 1.0, position_gain: float = 5.0, + orientation_gain: float = 1.0, nullspace_gain: float = 0.5, max_joint_velocity: float = 2.0, dls_damping: float = 0.05, @@ -181,15 +231,22 @@ def jparse_step( Computes Jacobian -> pseudo-inverse -> joint velocities -> integrate -> clamp. + When target_wxyz is provided, uses the full 6-DOF Jacobian for position + + orientation tracking. Otherwise uses position-only (3-DOF). + Args: robot: PyRoKi robot model. cfg: Current joint configuration (actuated_count,). target_link_index: Index of the target link. target_position: Target position (3,). + target_wxyz: Target orientation as quaternion in wxyz format (4,). + If None, only position is tracked. method: IK method — "jparse", "pinv", or "dls". gamma: J-PARSE singularity threshold. - singular_direction_gain: Gain for singular directions. + singular_direction_gain_position: Singular gain for position dimensions. + singular_direction_gain_angular: Singular gain for angular dimensions. position_gain: Proportional gain for position error. + orientation_gain: Proportional gain for orientation error. nullspace_gain: Gain for nullspace motion toward home. max_joint_velocity: Maximum joint velocity (rad/s). dls_damping: Damping factor for DLS method. @@ -199,27 +256,71 @@ def jparse_step( Returns: Tuple of (new_cfg, info): - new_cfg: New joint configuration after integration (numpy). - - info: Dict with position_error, max_joint_vel, jacobian, - manipulability, inverse_condition_number. + - info: Dict with position_error, orientation_error, max_joint_vel, + jacobian, manipulability, inverse_condition_number. """ cfg = jnp.asarray(cfg) target_position = jnp.asarray(target_position) + position_only = target_wxyz is None - # Current end-effector position. + # Current end-effector pose. poses = robot.forward_kinematics(cfg) - current_pos = jaxlie.SE3(poses[target_link_index]).translation() + target_pose = jaxlie.SE3(poses[target_link_index]) + current_pos = target_pose.translation() - # Position error and desired task-space velocity. + # Position error. pos_error = target_position - current_pos pos_error_mag = float(jnp.linalg.norm(pos_error)) - v_des = position_gain * pos_error + + # Build desired task-space velocity. + if position_only: + v_des = position_gain * pos_error + else: + target_wxyz = jnp.asarray(target_wxyz) + target_wxyz = target_wxyz / jnp.linalg.norm(target_wxyz) + + current_wxyz = target_pose.rotation().wxyz + current_wxyz = current_wxyz / jnp.linalg.norm(current_wxyz) + + # Ensure shortest-path quaternion. + target_wxyz = jnp.where( + jnp.dot(target_wxyz, current_wxyz) < 0, -target_wxyz, target_wxyz + ) + + # Orientation error via SO3 log map. + q_current = jaxlie.SO3(current_wxyz) + q_target = jaxlie.SO3(target_wxyz) + omega_error = (q_target @ q_current.inverse()).log() + + # Clamp orientation error magnitude. + omega_mag = jnp.linalg.norm(omega_error) + max_omega = 1.0 # rad + omega_error = jnp.where( + omega_mag > max_omega, omega_error * max_omega / omega_mag, omega_error + ) + + v_des = jnp.concatenate( + [ + position_gain * pos_error, + orientation_gain * omega_error, + ] + ) # Compute Jacobian. - jacobian = compute_jacobian(robot, cfg, target_link_index, position_only=True) + jacobian = compute_jacobian( + robot, cfg, target_link_index, position_only=position_only + ) # Compute pseudo-inverse and nullspace. if method == "jparse": - J_inv, N = jparse_pseudoinverse(jacobian, gamma, singular_direction_gain) + J_inv, N = jparse_pseudoinverse( + jacobian, + gamma=gamma, + singular_direction_gain_position=singular_direction_gain_position, + singular_direction_gain_angular=singular_direction_gain_angular, + position_dimensions=3, + angular_dimensions=0 if position_only else 3, + ) elif method == "pinv": J_inv = pinv(jacobian) N = jnp.eye(jacobian.shape[1]) - J_inv @ jacobian @@ -262,6 +363,9 @@ def jparse_step( info = { "position_error": pos_error_mag, + "orientation_error": float(jnp.linalg.norm(omega_error)) + if not position_only + else 0.0, "max_joint_vel": max_joint_vel, "jacobian": np.asarray(jacobian), "manipulability": float(manipulability_measure(jacobian)), diff --git a/tests/test_jparse.py b/tests/test_jparse.py index 343b2cd7..d57ab899 100644 --- a/tests/test_jparse.py +++ b/tests/test_jparse.py @@ -93,14 +93,47 @@ class TestSingularDirectionGains: def test_finite_output_with_gain(self): """Finite output with specified singular direction gain.""" J = np.array([[1.0, 1.0], [1.0, 1.001]]) - J_parse, _ = jparse_pseudoinverse(J, gamma=0.1, singular_direction_gain=2.0) + J_parse, _ = jparse_pseudoinverse( + J, + gamma=0.1, + singular_direction_gain_position=2.0, + ) assert np.all(np.isfinite(J_parse)) def test_gain_affects_output(self): """Different gains produce different results for near-singular case.""" J = np.array([[1.0, 1.0], [1.0, 1.001]]) - J1, _ = jparse_pseudoinverse(J, gamma=0.1, singular_direction_gain=1.0) - J2, _ = jparse_pseudoinverse(J, gamma=0.1, singular_direction_gain=3.0) + J1, _ = jparse_pseudoinverse( + J, + gamma=0.1, + singular_direction_gain_position=1.0, + ) + J2, _ = jparse_pseudoinverse( + J, + gamma=0.1, + singular_direction_gain_position=3.0, + ) + assert not np.allclose(J1, J2) + + def test_angular_gain_affects_6dof_output(self): + """Angular gain changes output when angular directions are singular.""" + J = np.diag([1.0, 1.0, 1.0, 1.0, 1.0, 1e-4]) + J1, _ = jparse_pseudoinverse( + J, + gamma=0.1, + singular_direction_gain_position=1.0, + singular_direction_gain_angular=1.0, + position_dimensions=3, + angular_dimensions=3, + ) + J2, _ = jparse_pseudoinverse( + J, + gamma=0.1, + singular_direction_gain_position=1.0, + singular_direction_gain_angular=3.0, + position_dimensions=3, + angular_dimensions=3, + ) assert not np.allclose(J1, J2) @@ -180,6 +213,24 @@ def test_large_gamma(self): J_parse, _ = jparse_pseudoinverse(J, gamma=0.99) assert np.all(np.isfinite(J_parse)) + def test_missing_dimension_argument_raises(self): + """Providing only one dimension argument raises ValueError.""" + J = np.eye(6) + with pytest.raises( + ValueError, + match="Both position_dimensions and angular_dimensions must be provided.", + ): + jparse_pseudoinverse(J, position_dimensions=3) + + def test_dimension_sum_mismatch_raises(self): + """Dimension sum must match Jacobian row count.""" + J = np.eye(6) + with pytest.raises( + ValueError, + match="position_dimensions \\+ angular_dimensions must equal Jacobian row count.", + ): + jparse_pseudoinverse(J, position_dimensions=2, angular_dimensions=3) + class TestFeatureParity: """Cross-validate JAX implementation against original jparse_robotics.""" @@ -209,7 +260,15 @@ def test_jparse_output_matches(self, test_jacobians): for name, J in test_jacobians.items(): J_np = core.compute(J) - J_jax, _ = jparse_pseudoinverse(J, gamma=gamma) + row_count = J.shape[0] + J_jax, _ = jparse_pseudoinverse( + J, + gamma=gamma, + singular_direction_gain_position=1.0, + singular_direction_gain_angular=1.0, + position_dimensions=3 if row_count == 6 else row_count, + angular_dimensions=3 if row_count == 6 else 0, + ) np.testing.assert_allclose( np.array(J_jax), np.array(J_np), atol=1e-5, err_msg=f"Failed for {name}" ) @@ -319,9 +378,87 @@ def test_info_dict_keys(self, panda_robot): ) expected_keys = { "position_error", + "orientation_error", "max_joint_vel", "jacobian", "manipulability", "inverse_condition_number", } assert set(info.keys()) == expected_keys + + def test_orientation_tracking(self, panda_robot): + """Orientation tracking uses 6-DOF Jacobian and reduces error.""" + import jaxlie + + robot, target_link_index = panda_robot + cfg = np.array((robot.joints.lower_limits + robot.joints.upper_limits) / 2.0) + + # Use a small orientation offset from the initial EE pose so + # the controller can converge in limited steps. + import jax.numpy as jnp + + poses = robot.forward_kinematics(jnp.asarray(cfg)) + init_pose = jaxlie.SE3(poses[target_link_index]) + target_pos = np.array(init_pose.translation()) + np.array([0.05, 0.0, 0.0]) + target_wxyz = np.array(init_pose.rotation().wxyz) + + errors = [] + for _ in range(100): + cfg, info = jparse_step( + robot=robot, + cfg=cfg, + target_link_index=target_link_index, + target_position=target_pos, + target_wxyz=target_wxyz, + gamma=0.1, + dt=0.02, + ) + errors.append(info["position_error"]) + + # Position error should decrease. + assert errors[-1] < errors[0] + # Jacobian should be 6xn when tracking orientation. + assert info["jacobian"].shape[0] == 6 + # Orientation error should be present and finite. + assert np.isfinite(info["orientation_error"]) + + def test_orientation_convergence_large_rotation(self, panda_robot): + """Orientation converges with a non-trivial (~45 deg) rotation offset. + + This exercises the geometric Jacobian path: the analytical Jacobian + introduces artificial singularities at large rotations that prevent + convergence, while the geometric Jacobian handles them correctly. + """ + import jax.numpy as jnp + import jaxlie + + robot, target_link_index = panda_robot + cfg = np.array((robot.joints.lower_limits + robot.joints.upper_limits) / 2.0) + + poses = robot.forward_kinematics(jnp.asarray(cfg)) + init_pose = jaxlie.SE3(poses[target_link_index]) + + # Keep position the same, apply a ~45-degree rotation about the z-axis. + target_pos = np.array(init_pose.translation()) + angle = np.pi / 4 # 45 degrees + offset_rot = jaxlie.SO3.from_z_radians(angle) + target_rot = offset_rot @ init_pose.rotation() + target_wxyz = np.array(target_rot.wxyz) + + ori_errors = [] + for _ in range(200): + cfg, info = jparse_step( + robot=robot, + cfg=cfg, + target_link_index=target_link_index, + target_position=target_pos, + target_wxyz=target_wxyz, + gamma=0.1, + dt=0.02, + ) + ori_errors.append(info["orientation_error"]) + + # Orientation error should converge to a small value. + assert ori_errors[-1] < 0.15, ( + f"Orientation error did not converge: {ori_errors[-1]:.4f}" + ) From bcf714f3840a380b1345ae14cbbb7ebd47eb44d6 Mon Sep 17 00:00:00 2001 From: Chung Min Kim Date: Sat, 7 Feb 2026 22:30:05 -0800 Subject: [PATCH 04/11] Add agent collaboration documentation --- AGENTS.md | 74 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 74 insertions(+) create mode 100644 AGENTS.md diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 00000000..c2c43911 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,74 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## Project Overview + +PyRoki is a JAX-based Python library for differentiable robot kinematics and optimization. It provides forward kinematics, collision detection, and cost/constraint factories that integrate with the `jaxls` nonlinear least-squares solver. Python 3.10+. + +## Common Commands + +### Install +```bash +sudo apt install -y libsuitesparse-dev # system dependency for scikit-sparse +pip install -e ".[dev]" +``` + +### Test +```bash +pytest # run all tests +pytest tests/test_ik_vmap.py # run single test file +pytest -k "test_name" # run specific test by name +``` + +### Lint & Format +```bash +ruff check docs/ src/ examples/ # lint +ruff format docs/ src/ examples/ # format +ruff format docs/ src/ examples/ --check # check formatting without modifying +``` + +### Type Check +```bash +pyright . +``` + +### Build Docs +```bash +cd docs && sphinx-build source build -b dirhtml +``` + +## Architecture + +### Core Pipeline + +The typical workflow is: **URDF → Robot → forward kinematics → residuals → costs → jaxls solver**. + +1. **`Robot`** (`_robot.py`): `jax_dataclasses.pytree_dataclass` holding joint/link info. Created via `Robot.from_urdf()`. The `forward_kinematics(cfg)` method returns SE(3) transforms (as wxyz-xyz arrays) for all links. Supports `jax.vmap` for batched configs. + +2. **Residuals** (`_residuals/`): Functions with signature `(vals: VarValues, robot, joint_var, ...) -> Array`. These compute optimization residuals (pose error, joint limits, collision distances, etc.). + +3. **Costs** (`costs.py`): Thin wrappers created via `Cost.factory(residual_fn)`. Costs are soft penalties (least squares); constraints use `kind="constraint_leq_zero"` for augmented Lagrangian enforcement. + +4. **Collision** (`collision/`): `RobotCollision.from_urdf()` builds collision geometry (spheres, capsules, boxes). `CollGeom` is the base class for geometry primitives. `collide()` computes signed distances between geometry pairs. + +5. **Solver**: Uses `jaxls` (external) for Levenberg-Marquardt optimization with optional augmented Lagrangian constraints. Problems are built with `jaxls.LeastSquaresProblem` and solved with `jaxls.solve()`. + +### Key Patterns + +- All core data structures are `@jdc.pytree_dataclass` (JAX-compatible pytrees via `jax_dataclasses`) +- SE(3) transforms use `jaxlie.SE3` (Lie group library) +- Typed arrays use `jaxtyping.Float[Array, "shape"]` annotations +- Variables for optimization use `jaxls.Var` subclasses; `Robot.joint_var_cls` is auto-created from URDF +- `@jdc.jit` is used instead of `@jax.jit` for methods on pytree dataclasses + +### Examples Structure + +Examples live in `examples/` and use helper modules from `examples/pyroki_snippets/` for reusable solve patterns (e.g., `solve_ik`, `solve_ik_with_collision`). Visualization uses `viser`. + +## CI + +Three workflows run on PRs to `main`: +- **pytest**: Tests across Python 3.10–3.13 +- **pyright**: Type checking across Python 3.10–3.13 +- **formatting**: `ruff check` + `ruff format --check` on `docs/`, `src/`, `examples/` From 99422f56a6d20a73b775691b99918ae930d4cf27 Mon Sep 17 00:00:00 2001 From: Chung Min Kim Date: Sat, 7 Feb 2026 22:30:13 -0800 Subject: [PATCH 05/11] Add headless validation framework for J-PARSE and online planning --- examples/15_headless_planner_validation.py | 145 ++++ examples/pyroki_snippets/__init__.py | 11 + .../pyroki_snippets/_headless_validation.py | 733 ++++++++++++++++++ .../pyroki_snippets/_jparse_diagnostics.py | 98 +++ tests/test_headless_planner_dynamics.py | 64 ++ 5 files changed, 1051 insertions(+) create mode 100644 examples/15_headless_planner_validation.py create mode 100644 examples/pyroki_snippets/_headless_validation.py create mode 100644 examples/pyroki_snippets/_jparse_diagnostics.py create mode 100644 tests/test_headless_planner_dynamics.py diff --git a/examples/15_headless_planner_validation.py b/examples/15_headless_planner_validation.py new file mode 100644 index 00000000..d904bae9 --- /dev/null +++ b/examples/15_headless_planner_validation.py @@ -0,0 +1,145 @@ +"""Headless stability validation for J-PARSE IK and online planning.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +from loguru import logger +from pyroki_snippets._headless_validation import ( + available_jparse_profiles, + run_jparse_headless_scenarios, + run_online_planner_headless_scenarios, +) + + +def _print_method_summary(title: str, result: dict[str, Any]) -> None: + print( + f"\n[{title}] method={result.get('method', 'n/a')} " + f"strictness={result['strictness']} all_passed={result['all_passed']}" + ) + + for scenario_name, scenario in result["scenarios"].items(): + thresholds = scenario["thresholds"] + print( + f" - {scenario_name:24s} diff={scenario['difficulty']:6s} " + f"pass={scenario['passed']} " + f"final_pos={scenario['final_pos_err']:.4f}/{thresholds['final_pos_err']:.4f} " + f"final_ori={scenario['final_ori_err']:.4f}/{thresholds['final_ori_err']:.4f} " + f"settled={scenario['settled']}@{scenario['settling_step']}" + ) + + +def _print_summary(title: str, result: dict[str, Any]) -> None: + if result.get("compare_methods", False): + print(f"\n[{title}] compare_methods=True all_passed={result['all_passed']}") + for method_name, method_result in result["methods"].items(): + _print_method_summary(f"{title}:{method_name}", method_result) + return + + _print_method_summary(title, result) + + +def _dump_json(path: str, payload: dict[str, Any]) -> None: + output_path = Path(path) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(json.dumps(payload, indent=2), encoding="utf-8") + print(f"\nWrote JSON report to: {output_path}") + + +def main() -> int: + logger.disable("jaxls") + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--only", + choices=["jparse", "online", "both"], + default="both", + help="Which controller(s) to validate.", + ) + parser.add_argument( + "--strictness", + choices=["default", "relaxed"], + default="default", + help="Threshold strictness profile.", + ) + parser.add_argument( + "--jparse-steps", + type=int, + default=180, + help="Number of closed-loop steps for J-PARSE scenarios.", + ) + parser.add_argument( + "--replans", + type=int, + default=20, + help="Number of receding-horizon replans for online scenarios.", + ) + parser.add_argument( + "--method", + choices=["jparse", "pinv", "dls"], + default="jparse", + help="Velocity IK method when not in compare mode.", + ) + parser.add_argument( + "--compare-methods", + action="store_true", + help="Run and report jparse/pinv/dls side-by-side.", + ) + parser.add_argument( + "--controller-profile", + choices=list(available_jparse_profiles()), + default="default", + help="Controller gain profile for velocity IK.", + ) + parser.add_argument( + "--dump-json", + default=None, + help="Optional path to dump full results JSON.", + ) + args = parser.parse_args() + + aggregated_results: dict[str, Any] = {} + overall_flags: list[bool] = [] + + if args.only in {"jparse", "both"}: + jparse_result = run_jparse_headless_scenarios( + strictness=args.strictness, + steps=args.jparse_steps, + method=args.method, + compare_methods=args.compare_methods, + controller_profile=args.controller_profile, + ) + _print_summary("J-PARSE", jparse_result) + aggregated_results["jparse"] = jparse_result + overall_flags.append(bool(jparse_result["all_passed"])) + + if args.only in {"online", "both"}: + online_result = run_online_planner_headless_scenarios( + strictness=args.strictness, + replans=args.replans, + ) + _print_summary("ONLINE", online_result) + aggregated_results["online"] = online_result + overall_flags.append(bool(online_result["all_passed"])) + + all_passed = all(overall_flags) + print(f"\nOverall: {'PASS' if all_passed else 'FAIL'}") + + if args.dump_json: + payload = { + "strictness": args.strictness, + "controller_profile": args.controller_profile, + "compare_methods": args.compare_methods, + "results": aggregated_results, + "all_passed": all_passed, + } + _dump_json(args.dump_json, payload) + + return 0 if all_passed else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/pyroki_snippets/__init__.py b/examples/pyroki_snippets/__init__.py index 95168631..2b0c2113 100644 --- a/examples/pyroki_snippets/__init__.py +++ b/examples/pyroki_snippets/__init__.py @@ -10,3 +10,14 @@ solve_ik_with_multiple_targets as solve_ik_with_multiple_targets, ) from ._jparse import jparse_step as jparse_step +from ._headless_validation import compute_settling_metrics as compute_settling_metrics +from ._headless_validation import ( + run_jparse_headless_scenarios as run_jparse_headless_scenarios, +) +from ._headless_validation import ( + available_jparse_profiles as available_jparse_profiles, +) +from ._headless_validation import ( + run_online_planner_headless_scenarios as run_online_planner_headless_scenarios, +) +from ._jparse_diagnostics import run_gain_sweep as run_gain_sweep diff --git a/examples/pyroki_snippets/_headless_validation.py b/examples/pyroki_snippets/_headless_validation.py new file mode 100644 index 00000000..c20d06e4 --- /dev/null +++ b/examples/pyroki_snippets/_headless_validation.py @@ -0,0 +1,733 @@ +from __future__ import annotations + +from dataclasses import asdict, dataclass +from typing import Any, Literal + +import jax.numpy as jnp +import jaxlie +import numpy as np +import pyroki as pk +from pyroki.collision import HalfSpace, RobotCollision, Sphere +from robot_descriptions.loaders.yourdfpy import load_robot_description + +from ._jparse import jparse_step +from ._online_planning import solve_online_planning + +EPS = 1e-9 + + +@dataclass(frozen=True) +class StabilityThresholds: + final_pos_err: float + final_ori_err: float + settling_step_max: int + sign_changes_max: int + tail_ripple_max: float + + +@dataclass(frozen=True) +class SettlingConfig: + hold_steps: int + tail_window: int + derivative_deadband: float + warmup_steps: int + + +@dataclass(frozen=True) +class ScenarioSpec: + name: str + difficulty: Literal["easy", "medium", "hard"] + dpos: np.ndarray + axis: np.ndarray + angle: float + + +@dataclass(frozen=True) +class ControllerProfile: + gamma: float + singular_direction_gain_position: float + singular_direction_gain_angular: float + position_gain: float + orientation_gain: float + nullspace_gain: float + max_joint_velocity: float + dls_damping: float + + +_JPARSE_SCENARIOS: tuple[ScenarioSpec, ...] = ( + ScenarioSpec( + name="easy_small_offset", + difficulty="easy", + dpos=np.array([0.03, 0.00, 0.00]), + axis=np.array([0.0, 0.0, 1.0]), + angle=0.15, + ), + ScenarioSpec( + name="medium_mixed_offset", + difficulty="medium", + dpos=np.array([0.06, -0.04, 0.03]), + axis=np.array([0.0, 1.0, 0.0]), + angle=0.35, + ), + ScenarioSpec( + name="hard_large_rotation", + difficulty="hard", + dpos=np.array([0.08, 0.04, -0.02]), + axis=np.array([1.0, 0.0, 0.0]), + angle=0.60, + ), +) + +_ONLINE_SCENARIOS: tuple[ScenarioSpec, ...] = ( + ScenarioSpec( + name="easy_static_obstacle", + difficulty="easy", + dpos=np.array([-0.09, 0.02, -0.04]), + axis=np.array([0.0, 0.0, 1.0]), + angle=0.20, + ), + ScenarioSpec( + name="medium_near_path_obstacle", + difficulty="medium", + dpos=np.array([-0.13, -0.06, -0.06]), + axis=np.array([0.0, 1.0, 0.0]), + angle=0.35, + ), + ScenarioSpec( + name="hard_crossing_obstacle", + difficulty="hard", + dpos=np.array([-0.15, 0.07, -0.08]), + axis=np.array([1.0, 0.0, 0.0]), + angle=0.50, + ), +) + +_JPARSE_PROFILES: dict[str, ControllerProfile] = { + "default": ControllerProfile( + gamma=0.1, + singular_direction_gain_position=2.0, + singular_direction_gain_angular=2.0, + position_gain=5.0, + orientation_gain=1.0, + nullspace_gain=0.5, + max_joint_velocity=2.0, + dls_damping=0.05, + ), + "aggressive_orientation": ControllerProfile( + gamma=0.1, + singular_direction_gain_position=2.0, + singular_direction_gain_angular=3.0, + position_gain=3.5, + orientation_gain=2.0, + nullspace_gain=0.2, + max_joint_velocity=2.2, + dls_damping=0.05, + ), +} + + +def available_jparse_profiles() -> tuple[str, ...]: + return tuple(_JPARSE_PROFILES.keys()) + + +def _normalize_quat(wxyz: np.ndarray) -> np.ndarray: + wxyz = np.asarray(wxyz, dtype=np.float64) + return wxyz / max(np.linalg.norm(wxyz), EPS) + + +def _build_target_wxyz( + base_wxyz: np.ndarray, axis_xyz: np.ndarray, angle_rad: float +) -> np.ndarray: + axis = np.asarray(axis_xyz, dtype=np.float64) + axis = axis / max(np.linalg.norm(axis), EPS) + delta = jaxlie.SO3.exp(jnp.asarray(axis * angle_rad)) + base = jaxlie.SO3(jnp.asarray(_normalize_quat(base_wxyz))) + return np.asarray((delta @ base).wxyz) + + +def _orientation_error_norm(current_wxyz: np.ndarray, target_wxyz: np.ndarray) -> float: + q_current = jaxlie.SO3(jnp.asarray(_normalize_quat(current_wxyz))) + q_target = jaxlie.SO3(jnp.asarray(_normalize_quat(target_wxyz))) + return float(jnp.linalg.norm((q_target @ q_current.inverse()).log())) + + +def _pose_errors( + robot: pk.Robot, + cfg: np.ndarray, + target_link_index: int, + target_position: np.ndarray, + target_wxyz: np.ndarray, +) -> tuple[float, float]: + pose = jaxlie.SE3(robot.forward_kinematics(jnp.asarray(cfg))[target_link_index]) + pos_err = float( + np.linalg.norm(np.asarray(target_position) - np.asarray(pose.translation())) + ) + ori_err = _orientation_error_norm(np.asarray(pose.rotation().wxyz), target_wxyz) + return pos_err, ori_err + + +def compute_settling_metrics( + combined_error: np.ndarray, + position_error: np.ndarray, + orientation_error: np.ndarray, + *, + position_tolerance: float, + orientation_tolerance: float, + hold_steps: int, + tail_window: int, + derivative_deadband: float, + warmup_steps: int, +) -> dict[str, float | int | bool]: + combined = np.asarray(combined_error, dtype=np.float64) + pos = np.asarray(position_error, dtype=np.float64) + ori = np.asarray(orientation_error, dtype=np.float64) + + if combined.ndim != 1 or pos.ndim != 1 or ori.ndim != 1: + raise ValueError("Error series must be 1D arrays.") + if not (len(combined) == len(pos) == len(ori)): + raise ValueError("Error series must have the same length.") + if len(combined) == 0: + raise ValueError("Error series cannot be empty.") + + if len(combined) == 1: + sign_changes = 0 + else: + diffs = np.diff(combined[max(0, warmup_steps) :]) + diffs[np.abs(diffs) < derivative_deadband] = 0.0 + signs = np.sign(diffs) + signs = signs[signs != 0.0] + sign_changes = ( + int(np.sum(signs[1:] * signs[:-1] < 0.0)) if len(signs) >= 2 else 0 + ) + + tail = combined[max(0, len(combined) - tail_window) :] + tail_ripple = float( + (np.max(tail) - np.min(tail)) / max(float(np.max(combined)), EPS) + ) + + settling_step = -1 + if hold_steps > 0 and len(combined) >= hold_steps: + for start in range(0, len(combined) - hold_steps + 1): + end = start + hold_steps + if np.all(pos[start:end] <= position_tolerance) and np.all( + ori[start:end] <= orientation_tolerance + ): + settling_step = start + break + + return { + "final_pos_err": float(pos[-1]), + "final_ori_err": float(ori[-1]), + "peak_combined_err": float(np.max(combined)), + "sign_changes": sign_changes, + "tail_ripple": tail_ripple, + "settled": settling_step >= 0, + "settling_step": settling_step, + } + + +def _thresholds_with_strictness( + thresholds: StabilityThresholds, + strictness: Literal["default", "relaxed"], +) -> StabilityThresholds: + if strictness == "default": + return thresholds + if strictness != "relaxed": + raise ValueError(f"Unsupported strictness: {strictness}") + + return StabilityThresholds( + final_pos_err=thresholds.final_pos_err * 1.5, + final_ori_err=thresholds.final_ori_err * 1.5, + settling_step_max=int(round(thresholds.settling_step_max * 1.25)), + sign_changes_max=int(round(thresholds.sign_changes_max * 1.5)), + tail_ripple_max=thresholds.tail_ripple_max * 1.5, + ) + + +def _serialize_threshold_map( + threshold_map: dict[str, StabilityThresholds], +) -> dict[str, dict[str, float | int]]: + return { + key: { + "final_pos_err": value.final_pos_err, + "final_ori_err": value.final_ori_err, + "settling_step_max": value.settling_step_max, + "sign_changes_max": value.sign_changes_max, + "tail_ripple_max": value.tail_ripple_max, + } + for key, value in threshold_map.items() + } + + +def _controller_profile( + profile_name: str, + overrides: dict[str, float] | None, +) -> ControllerProfile: + if profile_name not in _JPARSE_PROFILES: + raise ValueError( + f"Unknown profile '{profile_name}'. Available: {sorted(_JPARSE_PROFILES)}" + ) + + profile_data = asdict(_JPARSE_PROFILES[profile_name]) + if overrides: + unknown_keys = set(overrides.keys()) - set(profile_data.keys()) + if unknown_keys: + raise ValueError(f"Unknown override keys: {sorted(unknown_keys)}") + profile_data.update(overrides) + + return ControllerProfile(**profile_data) + + +def _add_failure_classifiers( + metrics: dict[str, float | int | bool], + pos_errors: np.ndarray, + ori_errors: np.ndarray, + *, + thresholds: StabilityThresholds, + settling_cfg: SettlingConfig, +) -> None: + tail_window = max(3, min(settling_cfg.tail_window, len(ori_errors))) + ori_tail = ori_errors[-tail_window:] + pos_tail = pos_errors[-tail_window:] + + ori_diffs = np.diff(ori_tail) if len(ori_tail) >= 2 else np.array([0.0]) + plateau = bool( + np.all(np.abs(ori_diffs) <= settling_cfg.derivative_deadband) + and float(metrics["final_ori_err"]) > thresholds.final_ori_err + ) + + competing_objectives = bool( + np.mean(pos_tail <= thresholds.final_pos_err) >= 0.7 + and float(metrics["final_ori_err"]) > thresholds.final_ori_err + and np.mean(ori_diffs) > -settling_cfg.derivative_deadband + ) + + metrics["plateau"] = plateau + metrics["competing_objectives"] = competing_objectives + + +def _evaluate_thresholds( + metrics: dict[str, float | int | bool], + thresholds: StabilityThresholds, +) -> bool: + return bool( + metrics["final_pos_err"] <= thresholds.final_pos_err + and metrics["final_ori_err"] <= thresholds.final_ori_err + and metrics["settled"] + and metrics["settling_step"] <= thresholds.settling_step_max + and metrics["sign_changes"] <= thresholds.sign_changes_max + and metrics["tail_ripple"] <= thresholds.tail_ripple_max + ) + + +def _run_jparse_method( + *, + method: Literal["jparse", "pinv", "dls"], + strictness: Literal["default", "relaxed"], + steps: int, + dt: float, + capture_trace: bool, + profile: ControllerProfile, +) -> dict[str, Any]: + urdf = load_robot_description("panda_description") + robot = pk.Robot.from_urdf(urdf) + target_link_index = robot.links.names.index("panda_hand") + + cfg0 = np.asarray((robot.joints.lower_limits + robot.joints.upper_limits) / 2.0) + init_pose = jaxlie.SE3( + robot.forward_kinematics(jnp.asarray(cfg0))[target_link_index] + ) + init_pos = np.asarray(init_pose.translation()) + init_wxyz = np.asarray(init_pose.rotation().wxyz) + + threshold_by_difficulty = { + "easy": StabilityThresholds( + final_pos_err=0.02, + final_ori_err=0.16, + settling_step_max=140, + sign_changes_max=14, + tail_ripple_max=0.18, + ), + "medium": StabilityThresholds( + final_pos_err=0.02, + final_ori_err=0.20, + settling_step_max=150, + sign_changes_max=14, + tail_ripple_max=0.18, + ), + "hard": StabilityThresholds( + final_pos_err=0.03, + final_ori_err=0.35, + settling_step_max=170, + sign_changes_max=16, + tail_ripple_max=0.22, + ), + } + threshold_by_difficulty = { + key: _thresholds_with_strictness(value, strictness) + for key, value in threshold_by_difficulty.items() + } + + settling_cfg = SettlingConfig( + hold_steps=10, + tail_window=20, + derivative_deadband=1e-4, + warmup_steps=10, + ) + + scenario_results: dict[str, dict[str, Any]] = {} + for scenario in _JPARSE_SCENARIOS: + cfg = cfg0.copy() + target_pos = init_pos + scenario.dpos + target_wxyz = _build_target_wxyz(init_wxyz, scenario.axis, scenario.angle) + + pos_errors: list[float] = [] + ori_errors: list[float] = [] + combined_errors: list[float] = [] + dq_inf_series: list[float] = [] + saturation_series: list[float] = [] + inverse_condition_series: list[float] = [] + + thresholds = threshold_by_difficulty[scenario.difficulty] + for _ in range(steps): + cfg, info = jparse_step( + robot=robot, + cfg=cfg, + target_link_index=target_link_index, + target_position=target_pos, + target_wxyz=target_wxyz, + method=method, + gamma=profile.gamma, + singular_direction_gain_position=profile.singular_direction_gain_position, + singular_direction_gain_angular=profile.singular_direction_gain_angular, + position_gain=profile.position_gain, + orientation_gain=profile.orientation_gain, + nullspace_gain=profile.nullspace_gain, + max_joint_velocity=profile.max_joint_velocity, + dls_damping=profile.dls_damping, + dt=dt, + ) + + pos_err, ori_err = _pose_errors( + robot, + cfg, + target_link_index, + target_pos, + target_wxyz, + ) + pos_errors.append(pos_err) + ori_errors.append(ori_err) + + combined_errors.append( + float( + np.sqrt( + (pos_err / thresholds.final_pos_err) ** 2 + + (ori_err / thresholds.final_ori_err) ** 2 + ) + ) + ) + + dq_inf = float(info["max_joint_vel"]) + dq_inf_series.append(dq_inf) + saturation_series.append( + min(dq_inf / max(profile.max_joint_velocity, EPS), 1.0) + ) + inverse_condition_series.append(float(info["inverse_condition_number"])) + + pos_arr = np.asarray(pos_errors) + ori_arr = np.asarray(ori_errors) + combined_arr = np.asarray(combined_errors) + + metrics = compute_settling_metrics( + combined_arr, + pos_arr, + ori_arr, + position_tolerance=thresholds.final_pos_err, + orientation_tolerance=thresholds.final_ori_err, + hold_steps=settling_cfg.hold_steps, + tail_window=settling_cfg.tail_window, + derivative_deadband=settling_cfg.derivative_deadband, + warmup_steps=settling_cfg.warmup_steps, + ) + _add_failure_classifiers( + metrics, + pos_arr, + ori_arr, + thresholds=thresholds, + settling_cfg=settling_cfg, + ) + metrics["mean_saturation"] = float(np.mean(saturation_series)) + metrics["peak_saturation"] = float(np.max(saturation_series)) + metrics["mean_inverse_condition"] = float(np.mean(inverse_condition_series)) + metrics["passed"] = _evaluate_thresholds(metrics, thresholds) + + scenario_result: dict[str, Any] = { + "difficulty": scenario.difficulty, + "thresholds": { + "final_pos_err": thresholds.final_pos_err, + "final_ori_err": thresholds.final_ori_err, + "settling_step_max": thresholds.settling_step_max, + "sign_changes_max": thresholds.sign_changes_max, + "tail_ripple_max": thresholds.tail_ripple_max, + }, + **metrics, + } + + if capture_trace: + scenario_result["trace"] = { + "position_error": pos_errors, + "orientation_error": ori_errors, + "combined_error": combined_errors, + "dq_inf": dq_inf_series, + "saturation": saturation_series, + "inverse_condition_number": inverse_condition_series, + } + + scenario_results[scenario.name] = scenario_result + + return { + "controller": "jparse", + "method": method, + "strictness": strictness, + "profile": asdict(profile), + "thresholds_by_difficulty": _serialize_threshold_map(threshold_by_difficulty), + "settling_config": asdict(settling_cfg), + "scenarios": scenario_results, + "all_passed": all(result["passed"] for result in scenario_results.values()), + } + + +def run_jparse_headless_scenarios( + *, + strictness: Literal["default", "relaxed"] = "default", + steps: int = 180, + dt: float = 0.02, + method: Literal["jparse", "pinv", "dls"] = "jparse", + compare_methods: bool = False, + controller_profile: str = "default", + controller_overrides: dict[str, float] | None = None, + capture_trace: bool = False, +) -> dict[str, Any]: + profile = _controller_profile(controller_profile, controller_overrides) + + if not compare_methods: + return _run_jparse_method( + method=method, + strictness=strictness, + steps=steps, + dt=dt, + capture_trace=capture_trace, + profile=profile, + ) + + methods: tuple[Literal["jparse", "pinv", "dls"], ...] = ("jparse", "pinv", "dls") + method_results = { + method_name: _run_jparse_method( + method=method_name, + strictness=strictness, + steps=steps, + dt=dt, + capture_trace=capture_trace, + profile=profile, + ) + for method_name in methods + } + + return { + "controller": "velocity_ik", + "compare_methods": True, + "strictness": strictness, + "profile": asdict(profile), + "methods": method_results, + "all_passed": all(result["all_passed"] for result in method_results.values()), + } + + +def _online_obstacle_position( + scenario_name: str, + step: int, + replans: int, +) -> np.ndarray: + if scenario_name == "easy_static_obstacle": + return np.array([0.15, 0.45, 0.32]) + if scenario_name == "medium_near_path_obstacle": + return np.array([0.32, 0.12, 0.36]) + if scenario_name == "hard_crossing_obstacle": + alpha = 0.0 if replans <= 1 else step / (replans - 1) + y = 0.22 * (1.0 - 2.0 * alpha) + return np.array([0.35, y, 0.34]) + raise ValueError(f"Unknown scenario: {scenario_name}") + + +def run_online_planner_headless_scenarios( + *, + strictness: Literal["default", "relaxed"] = "default", + replans: int = 20, + timesteps: int = 5, + dt: float = 0.1, + capture_trace: bool = False, +) -> dict[str, Any]: + urdf = load_robot_description("panda_description") + robot = pk.Robot.from_urdf(urdf) + robot_coll = RobotCollision.from_urdf(urdf) + target_link_name = "panda_hand" + target_link_index = robot.links.names.index(target_link_name) + + cfg0 = np.asarray((robot.joints.lower_limits + robot.joints.upper_limits) / 2.0) + init_pose = jaxlie.SE3( + robot.forward_kinematics(jnp.asarray(cfg0))[target_link_index] + ) + init_pos = np.asarray(init_pose.translation()) + init_wxyz = np.asarray(init_pose.rotation().wxyz) + + plane = HalfSpace.from_point_and_normal( + np.array([0.0, 0.0, 0.0]), + np.array([0.0, 0.0, 1.0]), + ) + base_sphere = Sphere.from_center_and_radius( + np.array([0.0, 0.0, 0.0]), np.array([0.05]) + ) + + threshold_by_difficulty = { + "easy": StabilityThresholds( + final_pos_err=0.03, + final_ori_err=0.15, + settling_step_max=18, + sign_changes_max=10, + tail_ripple_max=0.22, + ), + "medium": StabilityThresholds( + final_pos_err=0.03, + final_ori_err=0.15, + settling_step_max=18, + sign_changes_max=10, + tail_ripple_max=0.22, + ), + "hard": StabilityThresholds( + final_pos_err=0.04, + final_ori_err=0.20, + settling_step_max=20, + sign_changes_max=12, + tail_ripple_max=0.25, + ), + } + threshold_by_difficulty = { + key: _thresholds_with_strictness(value, strictness) + for key, value in threshold_by_difficulty.items() + } + + settling_cfg = SettlingConfig( + hold_steps=4, + tail_window=6, + derivative_deadband=5e-4, + warmup_steps=2, + ) + + scenario_results: dict[str, dict[str, Any]] = {} + for scenario in _ONLINE_SCENARIOS: + cfg = cfg0.copy() + prev_sols = np.asarray( + robot.joint_var_cls.default_factory()[None].repeat(timesteps, axis=0) + ) + + target_pos = init_pos + scenario.dpos + target_wxyz = _build_target_wxyz(init_wxyz, scenario.axis, scenario.angle) + + pos_errors: list[float] = [] + ori_errors: list[float] = [] + combined_errors: list[float] = [] + + thresholds = threshold_by_difficulty[scenario.difficulty] + for step in range(replans): + sphere_pos = _online_obstacle_position(scenario.name, step, replans) + sphere_world = base_sphere.transform_from_wxyz_position( + wxyz=np.array([1.0, 0.0, 0.0, 0.0]), + position=sphere_pos, + ) + + sol_traj, _, _ = solve_online_planning( + robot=robot, + robot_coll=robot_coll, + world_coll=[plane, sphere_world], + target_link_name=target_link_name, + target_position=target_pos, + target_wxyz=target_wxyz, + timesteps=timesteps, + dt=dt, + start_cfg=cfg, + prev_sols=prev_sols, + ) + + cfg = np.asarray(sol_traj[0]) + prev_sols = np.asarray(sol_traj) + + pos_err, ori_err = _pose_errors( + robot, + cfg, + target_link_index, + target_pos, + target_wxyz, + ) + pos_errors.append(pos_err) + ori_errors.append(ori_err) + combined_errors.append( + float( + np.sqrt( + (pos_err / thresholds.final_pos_err) ** 2 + + (ori_err / thresholds.final_ori_err) ** 2 + ) + ) + ) + + pos_arr = np.asarray(pos_errors) + ori_arr = np.asarray(ori_errors) + combined_arr = np.asarray(combined_errors) + metrics = compute_settling_metrics( + combined_arr, + pos_arr, + ori_arr, + position_tolerance=thresholds.final_pos_err, + orientation_tolerance=thresholds.final_ori_err, + hold_steps=settling_cfg.hold_steps, + tail_window=settling_cfg.tail_window, + derivative_deadband=settling_cfg.derivative_deadband, + warmup_steps=settling_cfg.warmup_steps, + ) + _add_failure_classifiers( + metrics, + pos_arr, + ori_arr, + thresholds=thresholds, + settling_cfg=settling_cfg, + ) + metrics["passed"] = _evaluate_thresholds(metrics, thresholds) + + scenario_result: dict[str, Any] = { + "difficulty": scenario.difficulty, + "thresholds": { + "final_pos_err": thresholds.final_pos_err, + "final_ori_err": thresholds.final_ori_err, + "settling_step_max": thresholds.settling_step_max, + "sign_changes_max": thresholds.sign_changes_max, + "tail_ripple_max": thresholds.tail_ripple_max, + }, + **metrics, + } + + if capture_trace: + scenario_result["trace"] = { + "position_error": pos_errors, + "orientation_error": ori_errors, + "combined_error": combined_errors, + } + + scenario_results[scenario.name] = scenario_result + + return { + "controller": "online_planner", + "strictness": strictness, + "thresholds_by_difficulty": _serialize_threshold_map(threshold_by_difficulty), + "settling_config": asdict(settling_cfg), + "scenarios": scenario_results, + "all_passed": all(result["passed"] for result in scenario_results.values()), + } diff --git a/examples/pyroki_snippets/_jparse_diagnostics.py b/examples/pyroki_snippets/_jparse_diagnostics.py new file mode 100644 index 00000000..75537f90 --- /dev/null +++ b/examples/pyroki_snippets/_jparse_diagnostics.py @@ -0,0 +1,98 @@ +"""Diagnostics helpers for headless J-PARSE controller tuning.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import jax + +try: + from ._headless_validation import run_jparse_headless_scenarios +except ImportError: # pragma: no cover - script execution fallback + from pyroki_snippets._headless_validation import run_jparse_headless_scenarios + + +@dataclass(frozen=True) +class SweepResult: + position_gain: float + orientation_gain: float + nullspace_gain: float + all_passed: bool + medium_final_ori_err: float + medium_settled: bool + + +def run_gain_sweep( + *, + strictness: str = "default", + position_gains: tuple[float, ...] = (3.0, 5.0, 7.0), + orientation_gains: tuple[float, ...] = (1.0, 1.5, 2.0, 3.0), + nullspace_gains: tuple[float, ...] = (0.0, 0.2, 0.5), + prefer_cpu: bool = True, +) -> list[SweepResult]: + """Run a deterministic gain sweep and return ranked summary entries.""" + if prefer_cpu: + # Keep diagnostics stable and avoid GPU-specific resource failures. + jax.config.update("jax_platform_name", "cpu") + + results: list[SweepResult] = [] + + for position_gain in position_gains: + for orientation_gain in orientation_gains: + for nullspace_gain in nullspace_gains: + output = run_jparse_headless_scenarios( + strictness=strictness, # type: ignore[arg-type] + method="jparse", + compare_methods=False, + controller_profile="default", + controller_overrides={ + "position_gain": position_gain, + "orientation_gain": orientation_gain, + "nullspace_gain": nullspace_gain, + }, + ) + medium = output["scenarios"]["medium_mixed_offset"] + results.append( + SweepResult( + position_gain=position_gain, + orientation_gain=orientation_gain, + nullspace_gain=nullspace_gain, + all_passed=bool(output["all_passed"]), + medium_final_ori_err=float(medium["final_ori_err"]), + medium_settled=bool(medium["settled"]), + ) + ) + + return sorted( + results, + key=lambda item: ( + not item.all_passed, + item.medium_final_ori_err, + not item.medium_settled, + item.nullspace_gain, + ), + ) + + +def format_gain_sweep_table(results: list[SweepResult], top_k: int = 10) -> str: + """Format sweep results into a compact, human-readable table.""" + lines = [ + "pos_gain ori_gain null_gain all_passed medium_ori medium_settled", + "-------- -------- -------- ----------- ---------- --------------", + ] + for row in results[:top_k]: + lines.append( + f"{row.position_gain:8.2f} {row.orientation_gain:8.2f} {row.nullspace_gain:8.2f} " + f"{str(row.all_passed):11s} {row.medium_final_ori_err:10.4f} {str(row.medium_settled):14s}" + ) + return "\n".join(lines) + + +def main() -> int: + results = run_gain_sweep() + print(format_gain_sweep_table(results, top_k=12)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_headless_planner_dynamics.py b/tests/test_headless_planner_dynamics.py new file mode 100644 index 00000000..3ca0acc8 --- /dev/null +++ b/tests/test_headless_planner_dynamics.py @@ -0,0 +1,64 @@ +"""Headless dynamics validation for J-PARSE IK and online planning.""" + +import sys + +sys.path.insert(0, "examples") + +from pyroki_snippets._headless_validation import ( # pylint: disable=wrong-import-position + run_jparse_headless_scenarios, + run_online_planner_headless_scenarios, +) + + +def _failure_details(result: dict) -> str: + lines = [ + f"controller={result['controller']} strictness={result['strictness']}", + ] + if "thresholds_by_difficulty" in result: + lines.append(f"thresholds_by_difficulty={result['thresholds_by_difficulty']}") + for name, metrics in result["scenarios"].items(): + lines.append(f"scenario={name} metrics={metrics}") + return "\n".join(lines) + + +def test_jparse_headless_dynamics_stability(): + """Tiered J-PARSE validation: easy/medium strict; hard stability+bounds.""" + result = run_jparse_headless_scenarios(strictness="default") + + scenarios = result["scenarios"] + easy = scenarios["easy_small_offset"] + medium = scenarios["medium_mixed_offset"] + hard = scenarios["hard_large_rotation"] + + # Easy and medium must satisfy strict orientation settling thresholds. + assert easy["final_ori_err"] <= easy["thresholds"]["final_ori_err"], ( + _failure_details(result) + ) + assert medium["final_ori_err"] <= medium["thresholds"]["final_ori_err"], ( + _failure_details(result) + ) + + # Hard requires stable behavior with bounded residual, not aggressive accuracy. + assert hard["sign_changes"] <= hard["thresholds"]["sign_changes_max"], ( + _failure_details(result) + ) + assert hard["tail_ripple"] <= hard["thresholds"]["tail_ripple_max"], ( + _failure_details(result) + ) + assert hard["final_ori_err"] <= hard["thresholds"]["final_ori_err"], ( + _failure_details(result) + ) + + # Keep strict overall gate active until all tiered criteria pass. + assert result["all_passed"], _failure_details(result) + + +def test_online_planner_headless_dynamics_stability(): + """Online planner replanning loop has acceptable tiered settling behavior.""" + result = run_online_planner_headless_scenarios(strictness="default") + + assert result["all_passed"], _failure_details(result) + for scenario_name, scenario_metrics in result["scenarios"].items(): + assert scenario_metrics["passed"], ( + f"scenario={scenario_name}\n{_failure_details(result)}" + ) From 2f1302d64da457d8406d8c4a5d654ad2cb3bfbad Mon Sep 17 00:00:00 2001 From: Chung Min Kim Date: Sat, 7 Feb 2026 23:37:01 -0800 Subject: [PATCH 06/11] Simplify J-PARSE example and rename to 14_singularity_aware_ik.py Strip the example down to ~80 lines matching the minimal style of other examples (01, 04, 05). Remove GUI folders, method dropdown, orientation toggle, singular/position/orientation gain sliders, max velocity slider, and most metric displays. Keep gamma threshold, nullspace gain, manipulability readout, and ellipsoid toggle. --- examples/14_jparse_ik.py | 160 ---------------------------- examples/14_singularity_aware_ik.py | 82 ++++++++++++++ 2 files changed, 82 insertions(+), 160 deletions(-) delete mode 100644 examples/14_jparse_ik.py create mode 100644 examples/14_singularity_aware_ik.py diff --git a/examples/14_jparse_ik.py b/examples/14_jparse_ik.py deleted file mode 100644 index 3a89e294..00000000 --- a/examples/14_jparse_ik.py +++ /dev/null @@ -1,160 +0,0 @@ -"""J-PARSE Velocity IK - -Singularity-aware velocity IK using J-PARSE with PyRoKi. - -J-PARSE (Jacobian-based Projection Algorithm for Resolving Singularities -Effectively) provides smooth behavior near singularities by decomposing -the Jacobian into non-singular and singular directions. - -Compare J-PARSE against standard pseudo-inverse and damped least squares -methods using the GUI controls. -""" - -import time - -import numpy as np -import pyroki as pk -import viser -from robot_descriptions.loaders.yourdfpy import load_robot_description -from viser.extras import ViserUrdf - -from pyroki_snippets._jparse import jparse_step - - -def main(): - """Main function for J-PARSE velocity IK example.""" - - urdf = load_robot_description("panda_description") - target_link_name = "panda_hand" - - # Create robot. - robot = pk.Robot.from_urdf(urdf) - target_link_index = robot.links.names.index(target_link_name) - - # Initial configuration (middle of joint range). - cfg = np.array((robot.joints.lower_limits + robot.joints.upper_limits) / 2.0) - - # Set up visualizer. - server = viser.ViserServer() - server.scene.add_grid("/ground", width=2, height=2) - urdf_vis = ViserUrdf(server, urdf, root_node_name="/base") - - # Target gizmo. - ik_target = server.scene.add_transform_controls( - "/ik_target", scale=0.2, position=(0.61, 0.0, 0.56), wxyz=(0, 0, 1, 0) - ) - - # --- GUI controls --- - with server.gui.add_folder("Method"): - method_handle = server.gui.add_dropdown( - "IK Method", - options=["jparse", "pinv", "dls"], - initial_value="jparse", - ) - gamma_handle = server.gui.add_slider( - "Gamma (singularity threshold)", 0.01, 0.5, 0.01, 0.1 - ) - singular_gain_position_handle = server.gui.add_slider( - "Singular gain (position)", 0.5, 5.0, 0.1, 2.0 - ) - singular_gain_angular_handle = server.gui.add_slider( - "Singular gain (angular)", 0.5, 5.0, 0.1, 2.0 - ) - - with server.gui.add_folder("Control"): - track_orientation_handle = server.gui.add_checkbox("Track orientation", True) - position_gain_handle = server.gui.add_slider( - "Position gain", 1.0, 20.0, 0.5, 5.0 - ) - orientation_gain_handle = server.gui.add_slider( - "Orientation gain", 0.1, 5.0, 0.1, 1.0 - ) - nullspace_gain_handle = server.gui.add_slider( - "Nullspace gain", 0.0, 2.0, 0.05, 0.5 - ) - max_vel_handle = server.gui.add_slider("Max joint velocity", 0.5, 5.0, 0.1, 2.0) - - with server.gui.add_folder("Visualization"): - show_ellipsoid_handle = server.gui.add_checkbox("Show ellipsoid", True) - - with server.gui.add_folder("Metrics"): - manipulability_handle = server.gui.add_number( - "Manipulability", 0.0, disabled=True - ) - condition_handle = server.gui.add_number( - "Inv. condition number", 0.0, disabled=True - ) - max_vel_display = server.gui.add_number( - "Max joint vel (rad/s)", 0.0, disabled=True - ) - pos_error_handle = server.gui.add_number( - "Position error (m)", 0.0, disabled=True - ) - ori_error_handle = server.gui.add_number( - "Orientation error (rad)", 0.0, disabled=True - ) - singularity_handle = server.gui.add_text( - "Singularity status", "", disabled=True - ) - - # Manipulability ellipsoid visualization. - manip_ellipse = pk.viewer.ManipulabilityEllipse( - server, - robot, - root_node_name="/manipulability", - target_link_name=target_link_name, - ) - - # Control loop at ~50Hz. - while True: - start_time = time.time() - - # Take one velocity IK step. - target_wxyz = ( - np.array(ik_target.wxyz) if track_orientation_handle.value else None - ) - cfg, info = jparse_step( - robot=robot, - cfg=cfg, - target_link_index=target_link_index, - target_position=np.array(ik_target.position), - target_wxyz=target_wxyz, - method=method_handle.value, # type: ignore - gamma=gamma_handle.value, - singular_direction_gain_position=singular_gain_position_handle.value, - singular_direction_gain_angular=singular_gain_angular_handle.value, - position_gain=position_gain_handle.value, - orientation_gain=orientation_gain_handle.value, - nullspace_gain=nullspace_gain_handle.value, - max_joint_velocity=max_vel_handle.value, - dt=0.02, - ) - - # Update metrics. - manipulability_handle.value = round(info["manipulability"], 4) - condition_handle.value = round(info["inverse_condition_number"], 4) - max_vel_display.value = round(info["max_joint_vel"], 3) - pos_error_handle.value = round(info["position_error"], 4) - ori_error_handle.value = round(info["orientation_error"], 4) - - # Singularity warning. - icn = info["inverse_condition_number"] - if icn < gamma_handle.value: - singularity_handle.value = f"NEAR SINGULAR (icn={icn:.3f})" - else: - singularity_handle.value = "OK" - - # Update ellipsoid. - manip_ellipse.set_visibility(show_ellipsoid_handle.value) - manip_ellipse.update(cfg) - - # Update robot visualization. - urdf_vis.update_cfg(cfg) - - # Sleep to maintain ~50Hz. - elapsed = time.time() - start_time - time.sleep(max(0.0, 0.02 - elapsed)) - - -if __name__ == "__main__": - main() diff --git a/examples/14_singularity_aware_ik.py b/examples/14_singularity_aware_ik.py new file mode 100644 index 00000000..22708728 --- /dev/null +++ b/examples/14_singularity_aware_ik.py @@ -0,0 +1,82 @@ +"""J-PARSE Velocity IK + +Singularity-aware velocity IK using J-PARSE with manipulability visualization. +""" + +import time + +import numpy as np +import pyroki as pk +import viser +from robot_descriptions.loaders.yourdfpy import load_robot_description +from viser.extras import ViserUrdf + +from pyroki_snippets._jparse import jparse_step + + +def main(): + urdf = load_robot_description("panda_description") + target_link_name = "panda_hand" + + # Create robot. + robot = pk.Robot.from_urdf(urdf) + target_link_index = robot.links.names.index(target_link_name) + + # Initial configuration (middle of joint range). + cfg = np.array((robot.joints.lower_limits + robot.joints.upper_limits) / 2.0) + + # Set up visualizer. + server = viser.ViserServer() + server.scene.add_grid("/ground", width=2, height=2) + urdf_vis = ViserUrdf(server, urdf, root_node_name="/base") + + # Target gizmo. + ik_target = server.scene.add_transform_controls( + "/ik_target", scale=0.2, position=(0.61, 0.0, 0.56), wxyz=(0, 0, 1, 0) + ) + + # GUI controls. + gamma_handle = server.gui.add_slider( + "Gamma (singularity threshold)", 0.01, 0.5, 0.01, 0.1 + ) + nullspace_gain_handle = server.gui.add_slider("Nullspace gain", 0.0, 2.0, 0.05, 0.5) + manipulability_handle = server.gui.add_number("Manipulability", 0.001, disabled=True) + show_ellipsoid_handle = server.gui.add_checkbox("Show ellipsoid", True) + + # Manipulability ellipsoid visualization. + manip_ellipse = pk.viewer.ManipulabilityEllipse( + server, + robot, + root_node_name="/manipulability", + target_link_name=target_link_name, + ) + + # Control loop at ~50Hz. + while True: + start_time = time.time() + + # Take one velocity IK step. + cfg, info = jparse_step( + robot=robot, + cfg=cfg, + target_link_index=target_link_index, + target_position=np.array(ik_target.position), + target_wxyz=np.array(ik_target.wxyz), + gamma=gamma_handle.value, + nullspace_gain=nullspace_gain_handle.value, + dt=0.02, + ) + + # Update metrics and visualization. + manipulability_handle.value = round(info["manipulability"], 4) + manip_ellipse.set_visibility(show_ellipsoid_handle.value) + manip_ellipse.update(cfg) + urdf_vis.update_cfg(cfg) + + # Sleep to maintain ~50Hz. + elapsed = time.time() - start_time + time.sleep(max(0.0, 0.02 - elapsed)) + + +if __name__ == "__main__": + main() From a344ba3a67f5197db7e436f0283110f650c3d701 Mon Sep 17 00:00:00 2001 From: Chung Min Kim Date: Sun, 8 Feb 2026 11:16:12 -0800 Subject: [PATCH 07/11] Remove headless planner validation example --- examples/15_headless_planner_validation.py | 145 --------------------- 1 file changed, 145 deletions(-) delete mode 100644 examples/15_headless_planner_validation.py diff --git a/examples/15_headless_planner_validation.py b/examples/15_headless_planner_validation.py deleted file mode 100644 index d904bae9..00000000 --- a/examples/15_headless_planner_validation.py +++ /dev/null @@ -1,145 +0,0 @@ -"""Headless stability validation for J-PARSE IK and online planning.""" - -from __future__ import annotations - -import argparse -import json -from pathlib import Path -from typing import Any - -from loguru import logger -from pyroki_snippets._headless_validation import ( - available_jparse_profiles, - run_jparse_headless_scenarios, - run_online_planner_headless_scenarios, -) - - -def _print_method_summary(title: str, result: dict[str, Any]) -> None: - print( - f"\n[{title}] method={result.get('method', 'n/a')} " - f"strictness={result['strictness']} all_passed={result['all_passed']}" - ) - - for scenario_name, scenario in result["scenarios"].items(): - thresholds = scenario["thresholds"] - print( - f" - {scenario_name:24s} diff={scenario['difficulty']:6s} " - f"pass={scenario['passed']} " - f"final_pos={scenario['final_pos_err']:.4f}/{thresholds['final_pos_err']:.4f} " - f"final_ori={scenario['final_ori_err']:.4f}/{thresholds['final_ori_err']:.4f} " - f"settled={scenario['settled']}@{scenario['settling_step']}" - ) - - -def _print_summary(title: str, result: dict[str, Any]) -> None: - if result.get("compare_methods", False): - print(f"\n[{title}] compare_methods=True all_passed={result['all_passed']}") - for method_name, method_result in result["methods"].items(): - _print_method_summary(f"{title}:{method_name}", method_result) - return - - _print_method_summary(title, result) - - -def _dump_json(path: str, payload: dict[str, Any]) -> None: - output_path = Path(path) - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_text(json.dumps(payload, indent=2), encoding="utf-8") - print(f"\nWrote JSON report to: {output_path}") - - -def main() -> int: - logger.disable("jaxls") - - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--only", - choices=["jparse", "online", "both"], - default="both", - help="Which controller(s) to validate.", - ) - parser.add_argument( - "--strictness", - choices=["default", "relaxed"], - default="default", - help="Threshold strictness profile.", - ) - parser.add_argument( - "--jparse-steps", - type=int, - default=180, - help="Number of closed-loop steps for J-PARSE scenarios.", - ) - parser.add_argument( - "--replans", - type=int, - default=20, - help="Number of receding-horizon replans for online scenarios.", - ) - parser.add_argument( - "--method", - choices=["jparse", "pinv", "dls"], - default="jparse", - help="Velocity IK method when not in compare mode.", - ) - parser.add_argument( - "--compare-methods", - action="store_true", - help="Run and report jparse/pinv/dls side-by-side.", - ) - parser.add_argument( - "--controller-profile", - choices=list(available_jparse_profiles()), - default="default", - help="Controller gain profile for velocity IK.", - ) - parser.add_argument( - "--dump-json", - default=None, - help="Optional path to dump full results JSON.", - ) - args = parser.parse_args() - - aggregated_results: dict[str, Any] = {} - overall_flags: list[bool] = [] - - if args.only in {"jparse", "both"}: - jparse_result = run_jparse_headless_scenarios( - strictness=args.strictness, - steps=args.jparse_steps, - method=args.method, - compare_methods=args.compare_methods, - controller_profile=args.controller_profile, - ) - _print_summary("J-PARSE", jparse_result) - aggregated_results["jparse"] = jparse_result - overall_flags.append(bool(jparse_result["all_passed"])) - - if args.only in {"online", "both"}: - online_result = run_online_planner_headless_scenarios( - strictness=args.strictness, - replans=args.replans, - ) - _print_summary("ONLINE", online_result) - aggregated_results["online"] = online_result - overall_flags.append(bool(online_result["all_passed"])) - - all_passed = all(overall_flags) - print(f"\nOverall: {'PASS' if all_passed else 'FAIL'}") - - if args.dump_json: - payload = { - "strictness": args.strictness, - "controller_profile": args.controller_profile, - "compare_methods": args.compare_methods, - "results": aggregated_results, - "all_passed": all_passed, - } - _dump_json(args.dump_json, payload) - - return 0 if all_passed else 1 - - -if __name__ == "__main__": - raise SystemExit(main()) From 7be2047d41e8f08cc910f302c471c8d9f87fe17e Mon Sep 17 00:00:00 2001 From: Chung Min Kim Date: Sun, 8 Feb 2026 11:21:53 -0800 Subject: [PATCH 08/11] Remove J-PARSE headless validation and diagnostics infrastructure --- AGENTS.md | 74 -- CLAUDE.md | 74 -- examples/pyroki_snippets/__init__.py | 11 - .../pyroki_snippets/_headless_validation.py | 733 ------------------ .../pyroki_snippets/_jparse_diagnostics.py | 98 --- tests/test_headless_planner_dynamics.py | 64 -- tests/test_jparse.py | 464 ----------- 7 files changed, 1518 deletions(-) delete mode 100644 AGENTS.md delete mode 100644 CLAUDE.md delete mode 100644 examples/pyroki_snippets/_headless_validation.py delete mode 100644 examples/pyroki_snippets/_jparse_diagnostics.py delete mode 100644 tests/test_headless_planner_dynamics.py delete mode 100644 tests/test_jparse.py diff --git a/AGENTS.md b/AGENTS.md deleted file mode 100644 index c2c43911..00000000 --- a/AGENTS.md +++ /dev/null @@ -1,74 +0,0 @@ -# CLAUDE.md - -This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. - -## Project Overview - -PyRoki is a JAX-based Python library for differentiable robot kinematics and optimization. It provides forward kinematics, collision detection, and cost/constraint factories that integrate with the `jaxls` nonlinear least-squares solver. Python 3.10+. - -## Common Commands - -### Install -```bash -sudo apt install -y libsuitesparse-dev # system dependency for scikit-sparse -pip install -e ".[dev]" -``` - -### Test -```bash -pytest # run all tests -pytest tests/test_ik_vmap.py # run single test file -pytest -k "test_name" # run specific test by name -``` - -### Lint & Format -```bash -ruff check docs/ src/ examples/ # lint -ruff format docs/ src/ examples/ # format -ruff format docs/ src/ examples/ --check # check formatting without modifying -``` - -### Type Check -```bash -pyright . -``` - -### Build Docs -```bash -cd docs && sphinx-build source build -b dirhtml -``` - -## Architecture - -### Core Pipeline - -The typical workflow is: **URDF → Robot → forward kinematics → residuals → costs → jaxls solver**. - -1. **`Robot`** (`_robot.py`): `jax_dataclasses.pytree_dataclass` holding joint/link info. Created via `Robot.from_urdf()`. The `forward_kinematics(cfg)` method returns SE(3) transforms (as wxyz-xyz arrays) for all links. Supports `jax.vmap` for batched configs. - -2. **Residuals** (`_residuals/`): Functions with signature `(vals: VarValues, robot, joint_var, ...) -> Array`. These compute optimization residuals (pose error, joint limits, collision distances, etc.). - -3. **Costs** (`costs.py`): Thin wrappers created via `Cost.factory(residual_fn)`. Costs are soft penalties (least squares); constraints use `kind="constraint_leq_zero"` for augmented Lagrangian enforcement. - -4. **Collision** (`collision/`): `RobotCollision.from_urdf()` builds collision geometry (spheres, capsules, boxes). `CollGeom` is the base class for geometry primitives. `collide()` computes signed distances between geometry pairs. - -5. **Solver**: Uses `jaxls` (external) for Levenberg-Marquardt optimization with optional augmented Lagrangian constraints. Problems are built with `jaxls.LeastSquaresProblem` and solved with `jaxls.solve()`. - -### Key Patterns - -- All core data structures are `@jdc.pytree_dataclass` (JAX-compatible pytrees via `jax_dataclasses`) -- SE(3) transforms use `jaxlie.SE3` (Lie group library) -- Typed arrays use `jaxtyping.Float[Array, "shape"]` annotations -- Variables for optimization use `jaxls.Var` subclasses; `Robot.joint_var_cls` is auto-created from URDF -- `@jdc.jit` is used instead of `@jax.jit` for methods on pytree dataclasses - -### Examples Structure - -Examples live in `examples/` and use helper modules from `examples/pyroki_snippets/` for reusable solve patterns (e.g., `solve_ik`, `solve_ik_with_collision`). Visualization uses `viser`. - -## CI - -Three workflows run on PRs to `main`: -- **pytest**: Tests across Python 3.10–3.13 -- **pyright**: Type checking across Python 3.10–3.13 -- **formatting**: `ruff check` + `ruff format --check` on `docs/`, `src/`, `examples/` diff --git a/CLAUDE.md b/CLAUDE.md deleted file mode 100644 index c2c43911..00000000 --- a/CLAUDE.md +++ /dev/null @@ -1,74 +0,0 @@ -# CLAUDE.md - -This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. - -## Project Overview - -PyRoki is a JAX-based Python library for differentiable robot kinematics and optimization. It provides forward kinematics, collision detection, and cost/constraint factories that integrate with the `jaxls` nonlinear least-squares solver. Python 3.10+. - -## Common Commands - -### Install -```bash -sudo apt install -y libsuitesparse-dev # system dependency for scikit-sparse -pip install -e ".[dev]" -``` - -### Test -```bash -pytest # run all tests -pytest tests/test_ik_vmap.py # run single test file -pytest -k "test_name" # run specific test by name -``` - -### Lint & Format -```bash -ruff check docs/ src/ examples/ # lint -ruff format docs/ src/ examples/ # format -ruff format docs/ src/ examples/ --check # check formatting without modifying -``` - -### Type Check -```bash -pyright . -``` - -### Build Docs -```bash -cd docs && sphinx-build source build -b dirhtml -``` - -## Architecture - -### Core Pipeline - -The typical workflow is: **URDF → Robot → forward kinematics → residuals → costs → jaxls solver**. - -1. **`Robot`** (`_robot.py`): `jax_dataclasses.pytree_dataclass` holding joint/link info. Created via `Robot.from_urdf()`. The `forward_kinematics(cfg)` method returns SE(3) transforms (as wxyz-xyz arrays) for all links. Supports `jax.vmap` for batched configs. - -2. **Residuals** (`_residuals/`): Functions with signature `(vals: VarValues, robot, joint_var, ...) -> Array`. These compute optimization residuals (pose error, joint limits, collision distances, etc.). - -3. **Costs** (`costs.py`): Thin wrappers created via `Cost.factory(residual_fn)`. Costs are soft penalties (least squares); constraints use `kind="constraint_leq_zero"` for augmented Lagrangian enforcement. - -4. **Collision** (`collision/`): `RobotCollision.from_urdf()` builds collision geometry (spheres, capsules, boxes). `CollGeom` is the base class for geometry primitives. `collide()` computes signed distances between geometry pairs. - -5. **Solver**: Uses `jaxls` (external) for Levenberg-Marquardt optimization with optional augmented Lagrangian constraints. Problems are built with `jaxls.LeastSquaresProblem` and solved with `jaxls.solve()`. - -### Key Patterns - -- All core data structures are `@jdc.pytree_dataclass` (JAX-compatible pytrees via `jax_dataclasses`) -- SE(3) transforms use `jaxlie.SE3` (Lie group library) -- Typed arrays use `jaxtyping.Float[Array, "shape"]` annotations -- Variables for optimization use `jaxls.Var` subclasses; `Robot.joint_var_cls` is auto-created from URDF -- `@jdc.jit` is used instead of `@jax.jit` for methods on pytree dataclasses - -### Examples Structure - -Examples live in `examples/` and use helper modules from `examples/pyroki_snippets/` for reusable solve patterns (e.g., `solve_ik`, `solve_ik_with_collision`). Visualization uses `viser`. - -## CI - -Three workflows run on PRs to `main`: -- **pytest**: Tests across Python 3.10–3.13 -- **pyright**: Type checking across Python 3.10–3.13 -- **formatting**: `ruff check` + `ruff format --check` on `docs/`, `src/`, `examples/` diff --git a/examples/pyroki_snippets/__init__.py b/examples/pyroki_snippets/__init__.py index 2b0c2113..95168631 100644 --- a/examples/pyroki_snippets/__init__.py +++ b/examples/pyroki_snippets/__init__.py @@ -10,14 +10,3 @@ solve_ik_with_multiple_targets as solve_ik_with_multiple_targets, ) from ._jparse import jparse_step as jparse_step -from ._headless_validation import compute_settling_metrics as compute_settling_metrics -from ._headless_validation import ( - run_jparse_headless_scenarios as run_jparse_headless_scenarios, -) -from ._headless_validation import ( - available_jparse_profiles as available_jparse_profiles, -) -from ._headless_validation import ( - run_online_planner_headless_scenarios as run_online_planner_headless_scenarios, -) -from ._jparse_diagnostics import run_gain_sweep as run_gain_sweep diff --git a/examples/pyroki_snippets/_headless_validation.py b/examples/pyroki_snippets/_headless_validation.py deleted file mode 100644 index c20d06e4..00000000 --- a/examples/pyroki_snippets/_headless_validation.py +++ /dev/null @@ -1,733 +0,0 @@ -from __future__ import annotations - -from dataclasses import asdict, dataclass -from typing import Any, Literal - -import jax.numpy as jnp -import jaxlie -import numpy as np -import pyroki as pk -from pyroki.collision import HalfSpace, RobotCollision, Sphere -from robot_descriptions.loaders.yourdfpy import load_robot_description - -from ._jparse import jparse_step -from ._online_planning import solve_online_planning - -EPS = 1e-9 - - -@dataclass(frozen=True) -class StabilityThresholds: - final_pos_err: float - final_ori_err: float - settling_step_max: int - sign_changes_max: int - tail_ripple_max: float - - -@dataclass(frozen=True) -class SettlingConfig: - hold_steps: int - tail_window: int - derivative_deadband: float - warmup_steps: int - - -@dataclass(frozen=True) -class ScenarioSpec: - name: str - difficulty: Literal["easy", "medium", "hard"] - dpos: np.ndarray - axis: np.ndarray - angle: float - - -@dataclass(frozen=True) -class ControllerProfile: - gamma: float - singular_direction_gain_position: float - singular_direction_gain_angular: float - position_gain: float - orientation_gain: float - nullspace_gain: float - max_joint_velocity: float - dls_damping: float - - -_JPARSE_SCENARIOS: tuple[ScenarioSpec, ...] = ( - ScenarioSpec( - name="easy_small_offset", - difficulty="easy", - dpos=np.array([0.03, 0.00, 0.00]), - axis=np.array([0.0, 0.0, 1.0]), - angle=0.15, - ), - ScenarioSpec( - name="medium_mixed_offset", - difficulty="medium", - dpos=np.array([0.06, -0.04, 0.03]), - axis=np.array([0.0, 1.0, 0.0]), - angle=0.35, - ), - ScenarioSpec( - name="hard_large_rotation", - difficulty="hard", - dpos=np.array([0.08, 0.04, -0.02]), - axis=np.array([1.0, 0.0, 0.0]), - angle=0.60, - ), -) - -_ONLINE_SCENARIOS: tuple[ScenarioSpec, ...] = ( - ScenarioSpec( - name="easy_static_obstacle", - difficulty="easy", - dpos=np.array([-0.09, 0.02, -0.04]), - axis=np.array([0.0, 0.0, 1.0]), - angle=0.20, - ), - ScenarioSpec( - name="medium_near_path_obstacle", - difficulty="medium", - dpos=np.array([-0.13, -0.06, -0.06]), - axis=np.array([0.0, 1.0, 0.0]), - angle=0.35, - ), - ScenarioSpec( - name="hard_crossing_obstacle", - difficulty="hard", - dpos=np.array([-0.15, 0.07, -0.08]), - axis=np.array([1.0, 0.0, 0.0]), - angle=0.50, - ), -) - -_JPARSE_PROFILES: dict[str, ControllerProfile] = { - "default": ControllerProfile( - gamma=0.1, - singular_direction_gain_position=2.0, - singular_direction_gain_angular=2.0, - position_gain=5.0, - orientation_gain=1.0, - nullspace_gain=0.5, - max_joint_velocity=2.0, - dls_damping=0.05, - ), - "aggressive_orientation": ControllerProfile( - gamma=0.1, - singular_direction_gain_position=2.0, - singular_direction_gain_angular=3.0, - position_gain=3.5, - orientation_gain=2.0, - nullspace_gain=0.2, - max_joint_velocity=2.2, - dls_damping=0.05, - ), -} - - -def available_jparse_profiles() -> tuple[str, ...]: - return tuple(_JPARSE_PROFILES.keys()) - - -def _normalize_quat(wxyz: np.ndarray) -> np.ndarray: - wxyz = np.asarray(wxyz, dtype=np.float64) - return wxyz / max(np.linalg.norm(wxyz), EPS) - - -def _build_target_wxyz( - base_wxyz: np.ndarray, axis_xyz: np.ndarray, angle_rad: float -) -> np.ndarray: - axis = np.asarray(axis_xyz, dtype=np.float64) - axis = axis / max(np.linalg.norm(axis), EPS) - delta = jaxlie.SO3.exp(jnp.asarray(axis * angle_rad)) - base = jaxlie.SO3(jnp.asarray(_normalize_quat(base_wxyz))) - return np.asarray((delta @ base).wxyz) - - -def _orientation_error_norm(current_wxyz: np.ndarray, target_wxyz: np.ndarray) -> float: - q_current = jaxlie.SO3(jnp.asarray(_normalize_quat(current_wxyz))) - q_target = jaxlie.SO3(jnp.asarray(_normalize_quat(target_wxyz))) - return float(jnp.linalg.norm((q_target @ q_current.inverse()).log())) - - -def _pose_errors( - robot: pk.Robot, - cfg: np.ndarray, - target_link_index: int, - target_position: np.ndarray, - target_wxyz: np.ndarray, -) -> tuple[float, float]: - pose = jaxlie.SE3(robot.forward_kinematics(jnp.asarray(cfg))[target_link_index]) - pos_err = float( - np.linalg.norm(np.asarray(target_position) - np.asarray(pose.translation())) - ) - ori_err = _orientation_error_norm(np.asarray(pose.rotation().wxyz), target_wxyz) - return pos_err, ori_err - - -def compute_settling_metrics( - combined_error: np.ndarray, - position_error: np.ndarray, - orientation_error: np.ndarray, - *, - position_tolerance: float, - orientation_tolerance: float, - hold_steps: int, - tail_window: int, - derivative_deadband: float, - warmup_steps: int, -) -> dict[str, float | int | bool]: - combined = np.asarray(combined_error, dtype=np.float64) - pos = np.asarray(position_error, dtype=np.float64) - ori = np.asarray(orientation_error, dtype=np.float64) - - if combined.ndim != 1 or pos.ndim != 1 or ori.ndim != 1: - raise ValueError("Error series must be 1D arrays.") - if not (len(combined) == len(pos) == len(ori)): - raise ValueError("Error series must have the same length.") - if len(combined) == 0: - raise ValueError("Error series cannot be empty.") - - if len(combined) == 1: - sign_changes = 0 - else: - diffs = np.diff(combined[max(0, warmup_steps) :]) - diffs[np.abs(diffs) < derivative_deadband] = 0.0 - signs = np.sign(diffs) - signs = signs[signs != 0.0] - sign_changes = ( - int(np.sum(signs[1:] * signs[:-1] < 0.0)) if len(signs) >= 2 else 0 - ) - - tail = combined[max(0, len(combined) - tail_window) :] - tail_ripple = float( - (np.max(tail) - np.min(tail)) / max(float(np.max(combined)), EPS) - ) - - settling_step = -1 - if hold_steps > 0 and len(combined) >= hold_steps: - for start in range(0, len(combined) - hold_steps + 1): - end = start + hold_steps - if np.all(pos[start:end] <= position_tolerance) and np.all( - ori[start:end] <= orientation_tolerance - ): - settling_step = start - break - - return { - "final_pos_err": float(pos[-1]), - "final_ori_err": float(ori[-1]), - "peak_combined_err": float(np.max(combined)), - "sign_changes": sign_changes, - "tail_ripple": tail_ripple, - "settled": settling_step >= 0, - "settling_step": settling_step, - } - - -def _thresholds_with_strictness( - thresholds: StabilityThresholds, - strictness: Literal["default", "relaxed"], -) -> StabilityThresholds: - if strictness == "default": - return thresholds - if strictness != "relaxed": - raise ValueError(f"Unsupported strictness: {strictness}") - - return StabilityThresholds( - final_pos_err=thresholds.final_pos_err * 1.5, - final_ori_err=thresholds.final_ori_err * 1.5, - settling_step_max=int(round(thresholds.settling_step_max * 1.25)), - sign_changes_max=int(round(thresholds.sign_changes_max * 1.5)), - tail_ripple_max=thresholds.tail_ripple_max * 1.5, - ) - - -def _serialize_threshold_map( - threshold_map: dict[str, StabilityThresholds], -) -> dict[str, dict[str, float | int]]: - return { - key: { - "final_pos_err": value.final_pos_err, - "final_ori_err": value.final_ori_err, - "settling_step_max": value.settling_step_max, - "sign_changes_max": value.sign_changes_max, - "tail_ripple_max": value.tail_ripple_max, - } - for key, value in threshold_map.items() - } - - -def _controller_profile( - profile_name: str, - overrides: dict[str, float] | None, -) -> ControllerProfile: - if profile_name not in _JPARSE_PROFILES: - raise ValueError( - f"Unknown profile '{profile_name}'. Available: {sorted(_JPARSE_PROFILES)}" - ) - - profile_data = asdict(_JPARSE_PROFILES[profile_name]) - if overrides: - unknown_keys = set(overrides.keys()) - set(profile_data.keys()) - if unknown_keys: - raise ValueError(f"Unknown override keys: {sorted(unknown_keys)}") - profile_data.update(overrides) - - return ControllerProfile(**profile_data) - - -def _add_failure_classifiers( - metrics: dict[str, float | int | bool], - pos_errors: np.ndarray, - ori_errors: np.ndarray, - *, - thresholds: StabilityThresholds, - settling_cfg: SettlingConfig, -) -> None: - tail_window = max(3, min(settling_cfg.tail_window, len(ori_errors))) - ori_tail = ori_errors[-tail_window:] - pos_tail = pos_errors[-tail_window:] - - ori_diffs = np.diff(ori_tail) if len(ori_tail) >= 2 else np.array([0.0]) - plateau = bool( - np.all(np.abs(ori_diffs) <= settling_cfg.derivative_deadband) - and float(metrics["final_ori_err"]) > thresholds.final_ori_err - ) - - competing_objectives = bool( - np.mean(pos_tail <= thresholds.final_pos_err) >= 0.7 - and float(metrics["final_ori_err"]) > thresholds.final_ori_err - and np.mean(ori_diffs) > -settling_cfg.derivative_deadband - ) - - metrics["plateau"] = plateau - metrics["competing_objectives"] = competing_objectives - - -def _evaluate_thresholds( - metrics: dict[str, float | int | bool], - thresholds: StabilityThresholds, -) -> bool: - return bool( - metrics["final_pos_err"] <= thresholds.final_pos_err - and metrics["final_ori_err"] <= thresholds.final_ori_err - and metrics["settled"] - and metrics["settling_step"] <= thresholds.settling_step_max - and metrics["sign_changes"] <= thresholds.sign_changes_max - and metrics["tail_ripple"] <= thresholds.tail_ripple_max - ) - - -def _run_jparse_method( - *, - method: Literal["jparse", "pinv", "dls"], - strictness: Literal["default", "relaxed"], - steps: int, - dt: float, - capture_trace: bool, - profile: ControllerProfile, -) -> dict[str, Any]: - urdf = load_robot_description("panda_description") - robot = pk.Robot.from_urdf(urdf) - target_link_index = robot.links.names.index("panda_hand") - - cfg0 = np.asarray((robot.joints.lower_limits + robot.joints.upper_limits) / 2.0) - init_pose = jaxlie.SE3( - robot.forward_kinematics(jnp.asarray(cfg0))[target_link_index] - ) - init_pos = np.asarray(init_pose.translation()) - init_wxyz = np.asarray(init_pose.rotation().wxyz) - - threshold_by_difficulty = { - "easy": StabilityThresholds( - final_pos_err=0.02, - final_ori_err=0.16, - settling_step_max=140, - sign_changes_max=14, - tail_ripple_max=0.18, - ), - "medium": StabilityThresholds( - final_pos_err=0.02, - final_ori_err=0.20, - settling_step_max=150, - sign_changes_max=14, - tail_ripple_max=0.18, - ), - "hard": StabilityThresholds( - final_pos_err=0.03, - final_ori_err=0.35, - settling_step_max=170, - sign_changes_max=16, - tail_ripple_max=0.22, - ), - } - threshold_by_difficulty = { - key: _thresholds_with_strictness(value, strictness) - for key, value in threshold_by_difficulty.items() - } - - settling_cfg = SettlingConfig( - hold_steps=10, - tail_window=20, - derivative_deadband=1e-4, - warmup_steps=10, - ) - - scenario_results: dict[str, dict[str, Any]] = {} - for scenario in _JPARSE_SCENARIOS: - cfg = cfg0.copy() - target_pos = init_pos + scenario.dpos - target_wxyz = _build_target_wxyz(init_wxyz, scenario.axis, scenario.angle) - - pos_errors: list[float] = [] - ori_errors: list[float] = [] - combined_errors: list[float] = [] - dq_inf_series: list[float] = [] - saturation_series: list[float] = [] - inverse_condition_series: list[float] = [] - - thresholds = threshold_by_difficulty[scenario.difficulty] - for _ in range(steps): - cfg, info = jparse_step( - robot=robot, - cfg=cfg, - target_link_index=target_link_index, - target_position=target_pos, - target_wxyz=target_wxyz, - method=method, - gamma=profile.gamma, - singular_direction_gain_position=profile.singular_direction_gain_position, - singular_direction_gain_angular=profile.singular_direction_gain_angular, - position_gain=profile.position_gain, - orientation_gain=profile.orientation_gain, - nullspace_gain=profile.nullspace_gain, - max_joint_velocity=profile.max_joint_velocity, - dls_damping=profile.dls_damping, - dt=dt, - ) - - pos_err, ori_err = _pose_errors( - robot, - cfg, - target_link_index, - target_pos, - target_wxyz, - ) - pos_errors.append(pos_err) - ori_errors.append(ori_err) - - combined_errors.append( - float( - np.sqrt( - (pos_err / thresholds.final_pos_err) ** 2 - + (ori_err / thresholds.final_ori_err) ** 2 - ) - ) - ) - - dq_inf = float(info["max_joint_vel"]) - dq_inf_series.append(dq_inf) - saturation_series.append( - min(dq_inf / max(profile.max_joint_velocity, EPS), 1.0) - ) - inverse_condition_series.append(float(info["inverse_condition_number"])) - - pos_arr = np.asarray(pos_errors) - ori_arr = np.asarray(ori_errors) - combined_arr = np.asarray(combined_errors) - - metrics = compute_settling_metrics( - combined_arr, - pos_arr, - ori_arr, - position_tolerance=thresholds.final_pos_err, - orientation_tolerance=thresholds.final_ori_err, - hold_steps=settling_cfg.hold_steps, - tail_window=settling_cfg.tail_window, - derivative_deadband=settling_cfg.derivative_deadband, - warmup_steps=settling_cfg.warmup_steps, - ) - _add_failure_classifiers( - metrics, - pos_arr, - ori_arr, - thresholds=thresholds, - settling_cfg=settling_cfg, - ) - metrics["mean_saturation"] = float(np.mean(saturation_series)) - metrics["peak_saturation"] = float(np.max(saturation_series)) - metrics["mean_inverse_condition"] = float(np.mean(inverse_condition_series)) - metrics["passed"] = _evaluate_thresholds(metrics, thresholds) - - scenario_result: dict[str, Any] = { - "difficulty": scenario.difficulty, - "thresholds": { - "final_pos_err": thresholds.final_pos_err, - "final_ori_err": thresholds.final_ori_err, - "settling_step_max": thresholds.settling_step_max, - "sign_changes_max": thresholds.sign_changes_max, - "tail_ripple_max": thresholds.tail_ripple_max, - }, - **metrics, - } - - if capture_trace: - scenario_result["trace"] = { - "position_error": pos_errors, - "orientation_error": ori_errors, - "combined_error": combined_errors, - "dq_inf": dq_inf_series, - "saturation": saturation_series, - "inverse_condition_number": inverse_condition_series, - } - - scenario_results[scenario.name] = scenario_result - - return { - "controller": "jparse", - "method": method, - "strictness": strictness, - "profile": asdict(profile), - "thresholds_by_difficulty": _serialize_threshold_map(threshold_by_difficulty), - "settling_config": asdict(settling_cfg), - "scenarios": scenario_results, - "all_passed": all(result["passed"] for result in scenario_results.values()), - } - - -def run_jparse_headless_scenarios( - *, - strictness: Literal["default", "relaxed"] = "default", - steps: int = 180, - dt: float = 0.02, - method: Literal["jparse", "pinv", "dls"] = "jparse", - compare_methods: bool = False, - controller_profile: str = "default", - controller_overrides: dict[str, float] | None = None, - capture_trace: bool = False, -) -> dict[str, Any]: - profile = _controller_profile(controller_profile, controller_overrides) - - if not compare_methods: - return _run_jparse_method( - method=method, - strictness=strictness, - steps=steps, - dt=dt, - capture_trace=capture_trace, - profile=profile, - ) - - methods: tuple[Literal["jparse", "pinv", "dls"], ...] = ("jparse", "pinv", "dls") - method_results = { - method_name: _run_jparse_method( - method=method_name, - strictness=strictness, - steps=steps, - dt=dt, - capture_trace=capture_trace, - profile=profile, - ) - for method_name in methods - } - - return { - "controller": "velocity_ik", - "compare_methods": True, - "strictness": strictness, - "profile": asdict(profile), - "methods": method_results, - "all_passed": all(result["all_passed"] for result in method_results.values()), - } - - -def _online_obstacle_position( - scenario_name: str, - step: int, - replans: int, -) -> np.ndarray: - if scenario_name == "easy_static_obstacle": - return np.array([0.15, 0.45, 0.32]) - if scenario_name == "medium_near_path_obstacle": - return np.array([0.32, 0.12, 0.36]) - if scenario_name == "hard_crossing_obstacle": - alpha = 0.0 if replans <= 1 else step / (replans - 1) - y = 0.22 * (1.0 - 2.0 * alpha) - return np.array([0.35, y, 0.34]) - raise ValueError(f"Unknown scenario: {scenario_name}") - - -def run_online_planner_headless_scenarios( - *, - strictness: Literal["default", "relaxed"] = "default", - replans: int = 20, - timesteps: int = 5, - dt: float = 0.1, - capture_trace: bool = False, -) -> dict[str, Any]: - urdf = load_robot_description("panda_description") - robot = pk.Robot.from_urdf(urdf) - robot_coll = RobotCollision.from_urdf(urdf) - target_link_name = "panda_hand" - target_link_index = robot.links.names.index(target_link_name) - - cfg0 = np.asarray((robot.joints.lower_limits + robot.joints.upper_limits) / 2.0) - init_pose = jaxlie.SE3( - robot.forward_kinematics(jnp.asarray(cfg0))[target_link_index] - ) - init_pos = np.asarray(init_pose.translation()) - init_wxyz = np.asarray(init_pose.rotation().wxyz) - - plane = HalfSpace.from_point_and_normal( - np.array([0.0, 0.0, 0.0]), - np.array([0.0, 0.0, 1.0]), - ) - base_sphere = Sphere.from_center_and_radius( - np.array([0.0, 0.0, 0.0]), np.array([0.05]) - ) - - threshold_by_difficulty = { - "easy": StabilityThresholds( - final_pos_err=0.03, - final_ori_err=0.15, - settling_step_max=18, - sign_changes_max=10, - tail_ripple_max=0.22, - ), - "medium": StabilityThresholds( - final_pos_err=0.03, - final_ori_err=0.15, - settling_step_max=18, - sign_changes_max=10, - tail_ripple_max=0.22, - ), - "hard": StabilityThresholds( - final_pos_err=0.04, - final_ori_err=0.20, - settling_step_max=20, - sign_changes_max=12, - tail_ripple_max=0.25, - ), - } - threshold_by_difficulty = { - key: _thresholds_with_strictness(value, strictness) - for key, value in threshold_by_difficulty.items() - } - - settling_cfg = SettlingConfig( - hold_steps=4, - tail_window=6, - derivative_deadband=5e-4, - warmup_steps=2, - ) - - scenario_results: dict[str, dict[str, Any]] = {} - for scenario in _ONLINE_SCENARIOS: - cfg = cfg0.copy() - prev_sols = np.asarray( - robot.joint_var_cls.default_factory()[None].repeat(timesteps, axis=0) - ) - - target_pos = init_pos + scenario.dpos - target_wxyz = _build_target_wxyz(init_wxyz, scenario.axis, scenario.angle) - - pos_errors: list[float] = [] - ori_errors: list[float] = [] - combined_errors: list[float] = [] - - thresholds = threshold_by_difficulty[scenario.difficulty] - for step in range(replans): - sphere_pos = _online_obstacle_position(scenario.name, step, replans) - sphere_world = base_sphere.transform_from_wxyz_position( - wxyz=np.array([1.0, 0.0, 0.0, 0.0]), - position=sphere_pos, - ) - - sol_traj, _, _ = solve_online_planning( - robot=robot, - robot_coll=robot_coll, - world_coll=[plane, sphere_world], - target_link_name=target_link_name, - target_position=target_pos, - target_wxyz=target_wxyz, - timesteps=timesteps, - dt=dt, - start_cfg=cfg, - prev_sols=prev_sols, - ) - - cfg = np.asarray(sol_traj[0]) - prev_sols = np.asarray(sol_traj) - - pos_err, ori_err = _pose_errors( - robot, - cfg, - target_link_index, - target_pos, - target_wxyz, - ) - pos_errors.append(pos_err) - ori_errors.append(ori_err) - combined_errors.append( - float( - np.sqrt( - (pos_err / thresholds.final_pos_err) ** 2 - + (ori_err / thresholds.final_ori_err) ** 2 - ) - ) - ) - - pos_arr = np.asarray(pos_errors) - ori_arr = np.asarray(ori_errors) - combined_arr = np.asarray(combined_errors) - metrics = compute_settling_metrics( - combined_arr, - pos_arr, - ori_arr, - position_tolerance=thresholds.final_pos_err, - orientation_tolerance=thresholds.final_ori_err, - hold_steps=settling_cfg.hold_steps, - tail_window=settling_cfg.tail_window, - derivative_deadband=settling_cfg.derivative_deadband, - warmup_steps=settling_cfg.warmup_steps, - ) - _add_failure_classifiers( - metrics, - pos_arr, - ori_arr, - thresholds=thresholds, - settling_cfg=settling_cfg, - ) - metrics["passed"] = _evaluate_thresholds(metrics, thresholds) - - scenario_result: dict[str, Any] = { - "difficulty": scenario.difficulty, - "thresholds": { - "final_pos_err": thresholds.final_pos_err, - "final_ori_err": thresholds.final_ori_err, - "settling_step_max": thresholds.settling_step_max, - "sign_changes_max": thresholds.sign_changes_max, - "tail_ripple_max": thresholds.tail_ripple_max, - }, - **metrics, - } - - if capture_trace: - scenario_result["trace"] = { - "position_error": pos_errors, - "orientation_error": ori_errors, - "combined_error": combined_errors, - } - - scenario_results[scenario.name] = scenario_result - - return { - "controller": "online_planner", - "strictness": strictness, - "thresholds_by_difficulty": _serialize_threshold_map(threshold_by_difficulty), - "settling_config": asdict(settling_cfg), - "scenarios": scenario_results, - "all_passed": all(result["passed"] for result in scenario_results.values()), - } diff --git a/examples/pyroki_snippets/_jparse_diagnostics.py b/examples/pyroki_snippets/_jparse_diagnostics.py deleted file mode 100644 index 75537f90..00000000 --- a/examples/pyroki_snippets/_jparse_diagnostics.py +++ /dev/null @@ -1,98 +0,0 @@ -"""Diagnostics helpers for headless J-PARSE controller tuning.""" - -from __future__ import annotations - -from dataclasses import dataclass - -import jax - -try: - from ._headless_validation import run_jparse_headless_scenarios -except ImportError: # pragma: no cover - script execution fallback - from pyroki_snippets._headless_validation import run_jparse_headless_scenarios - - -@dataclass(frozen=True) -class SweepResult: - position_gain: float - orientation_gain: float - nullspace_gain: float - all_passed: bool - medium_final_ori_err: float - medium_settled: bool - - -def run_gain_sweep( - *, - strictness: str = "default", - position_gains: tuple[float, ...] = (3.0, 5.0, 7.0), - orientation_gains: tuple[float, ...] = (1.0, 1.5, 2.0, 3.0), - nullspace_gains: tuple[float, ...] = (0.0, 0.2, 0.5), - prefer_cpu: bool = True, -) -> list[SweepResult]: - """Run a deterministic gain sweep and return ranked summary entries.""" - if prefer_cpu: - # Keep diagnostics stable and avoid GPU-specific resource failures. - jax.config.update("jax_platform_name", "cpu") - - results: list[SweepResult] = [] - - for position_gain in position_gains: - for orientation_gain in orientation_gains: - for nullspace_gain in nullspace_gains: - output = run_jparse_headless_scenarios( - strictness=strictness, # type: ignore[arg-type] - method="jparse", - compare_methods=False, - controller_profile="default", - controller_overrides={ - "position_gain": position_gain, - "orientation_gain": orientation_gain, - "nullspace_gain": nullspace_gain, - }, - ) - medium = output["scenarios"]["medium_mixed_offset"] - results.append( - SweepResult( - position_gain=position_gain, - orientation_gain=orientation_gain, - nullspace_gain=nullspace_gain, - all_passed=bool(output["all_passed"]), - medium_final_ori_err=float(medium["final_ori_err"]), - medium_settled=bool(medium["settled"]), - ) - ) - - return sorted( - results, - key=lambda item: ( - not item.all_passed, - item.medium_final_ori_err, - not item.medium_settled, - item.nullspace_gain, - ), - ) - - -def format_gain_sweep_table(results: list[SweepResult], top_k: int = 10) -> str: - """Format sweep results into a compact, human-readable table.""" - lines = [ - "pos_gain ori_gain null_gain all_passed medium_ori medium_settled", - "-------- -------- -------- ----------- ---------- --------------", - ] - for row in results[:top_k]: - lines.append( - f"{row.position_gain:8.2f} {row.orientation_gain:8.2f} {row.nullspace_gain:8.2f} " - f"{str(row.all_passed):11s} {row.medium_final_ori_err:10.4f} {str(row.medium_settled):14s}" - ) - return "\n".join(lines) - - -def main() -> int: - results = run_gain_sweep() - print(format_gain_sweep_table(results, top_k=12)) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tests/test_headless_planner_dynamics.py b/tests/test_headless_planner_dynamics.py deleted file mode 100644 index 3ca0acc8..00000000 --- a/tests/test_headless_planner_dynamics.py +++ /dev/null @@ -1,64 +0,0 @@ -"""Headless dynamics validation for J-PARSE IK and online planning.""" - -import sys - -sys.path.insert(0, "examples") - -from pyroki_snippets._headless_validation import ( # pylint: disable=wrong-import-position - run_jparse_headless_scenarios, - run_online_planner_headless_scenarios, -) - - -def _failure_details(result: dict) -> str: - lines = [ - f"controller={result['controller']} strictness={result['strictness']}", - ] - if "thresholds_by_difficulty" in result: - lines.append(f"thresholds_by_difficulty={result['thresholds_by_difficulty']}") - for name, metrics in result["scenarios"].items(): - lines.append(f"scenario={name} metrics={metrics}") - return "\n".join(lines) - - -def test_jparse_headless_dynamics_stability(): - """Tiered J-PARSE validation: easy/medium strict; hard stability+bounds.""" - result = run_jparse_headless_scenarios(strictness="default") - - scenarios = result["scenarios"] - easy = scenarios["easy_small_offset"] - medium = scenarios["medium_mixed_offset"] - hard = scenarios["hard_large_rotation"] - - # Easy and medium must satisfy strict orientation settling thresholds. - assert easy["final_ori_err"] <= easy["thresholds"]["final_ori_err"], ( - _failure_details(result) - ) - assert medium["final_ori_err"] <= medium["thresholds"]["final_ori_err"], ( - _failure_details(result) - ) - - # Hard requires stable behavior with bounded residual, not aggressive accuracy. - assert hard["sign_changes"] <= hard["thresholds"]["sign_changes_max"], ( - _failure_details(result) - ) - assert hard["tail_ripple"] <= hard["thresholds"]["tail_ripple_max"], ( - _failure_details(result) - ) - assert hard["final_ori_err"] <= hard["thresholds"]["final_ori_err"], ( - _failure_details(result) - ) - - # Keep strict overall gate active until all tiered criteria pass. - assert result["all_passed"], _failure_details(result) - - -def test_online_planner_headless_dynamics_stability(): - """Online planner replanning loop has acceptable tiered settling behavior.""" - result = run_online_planner_headless_scenarios(strictness="default") - - assert result["all_passed"], _failure_details(result) - for scenario_name, scenario_metrics in result["scenarios"].items(): - assert scenario_metrics["passed"], ( - f"scenario={scenario_name}\n{_failure_details(result)}" - ) diff --git a/tests/test_jparse.py b/tests/test_jparse.py deleted file mode 100644 index d57ab899..00000000 --- a/tests/test_jparse.py +++ /dev/null @@ -1,464 +0,0 @@ -"""Tests for the JAX J-PARSE implementation. - -Verifies correctness of the JAX port and cross-validates against the -original numpy-based jparse_robotics package. -""" - -import sys - -import numpy as np -import pytest - -sys.path.insert(0, "examples") - -from pyroki_snippets._jparse import ( - damped_least_squares, - inverse_condition_number, - jparse_pseudoinverse, - jparse_step, - manipulability_measure, - pinv, -) - - -class TestJParsePseudoinverse: - """Test the core jparse_pseudoinverse function.""" - - def test_well_conditioned_2dof(self): - """Well-conditioned 2-DOF planar Jacobian produces valid pseudo-inverse.""" - l1, l2 = 1.0, 1.0 - theta1, theta2 = np.pi / 4, np.pi / 4 - J = np.array( - [ - [ - -l1 * np.sin(theta1) - l2 * np.sin(theta1 + theta2), - -l2 * np.sin(theta1 + theta2), - ], - [ - l1 * np.cos(theta1) + l2 * np.cos(theta1 + theta2), - l2 * np.cos(theta1 + theta2), - ], - ] - ) - - J_parse, N = jparse_pseudoinverse(J, gamma=0.1) - - assert J_parse.shape == (2, 2) - # Valid pseudo-inverse: J @ J_parse @ J ≈ J - np.testing.assert_allclose(J @ J_parse @ J, J, atol=1e-6) - # Well-conditioned: should be close to pinv - np.testing.assert_allclose(J_parse, np.linalg.pinv(J), atol=0.1) - - def test_near_singular_no_nan(self): - """Near-singular Jacobian produces finite output.""" - J = np.array([[1.0, 1.0], [1.0, 1.001]]) - J_parse, N = jparse_pseudoinverse(J, gamma=0.1) - assert np.all(np.isfinite(J_parse)) - - def test_singular_no_nan(self): - """Exactly singular Jacobian produces finite output.""" - J = np.array([[1.0, 1.0], [1.0, 1.0]]) - J_parse, N = jparse_pseudoinverse(J, gamma=0.1) - assert np.all(np.isfinite(J_parse)) - - def test_redundant_system(self): - """Redundant 2x3 system has correct shapes and valid pseudo-inverse.""" - J = np.array([[1.0, 0.0, 1.0], [0.0, 1.0, 1.0]]) - J_parse, N = jparse_pseudoinverse(J, gamma=0.1) - assert J_parse.shape == (3, 2) - np.testing.assert_allclose(J @ J_parse @ J, J, atol=1e-6) - - -class TestNullspace: - """Test nullspace projector.""" - - def test_nullspace_shape(self): - """Nullspace projector has shape (n, n).""" - J = np.array([[1.0, 0.0, 1.0], [0.0, 1.0, 1.0]]) - _, N = jparse_pseudoinverse(J, gamma=0.1) - assert N.shape == (3, 3) - - def test_nullspace_projects_to_zero(self): - """Nullspace vectors project to ~zero in task space.""" - J = np.array([[1.0, 0.0, 1.0], [0.0, 1.0, 1.0]]) - _, N = jparse_pseudoinverse(J, gamma=0.1) - null_vec = N @ np.array([1.0, 1.0, 1.0]) - task_effect = J @ null_vec - np.testing.assert_allclose(task_effect, 0, atol=0.1) - - -class TestSingularDirectionGains: - """Test singular direction gain parameters.""" - - def test_finite_output_with_gain(self): - """Finite output with specified singular direction gain.""" - J = np.array([[1.0, 1.0], [1.0, 1.001]]) - J_parse, _ = jparse_pseudoinverse( - J, - gamma=0.1, - singular_direction_gain_position=2.0, - ) - assert np.all(np.isfinite(J_parse)) - - def test_gain_affects_output(self): - """Different gains produce different results for near-singular case.""" - J = np.array([[1.0, 1.0], [1.0, 1.001]]) - J1, _ = jparse_pseudoinverse( - J, - gamma=0.1, - singular_direction_gain_position=1.0, - ) - J2, _ = jparse_pseudoinverse( - J, - gamma=0.1, - singular_direction_gain_position=3.0, - ) - assert not np.allclose(J1, J2) - - def test_angular_gain_affects_6dof_output(self): - """Angular gain changes output when angular directions are singular.""" - J = np.diag([1.0, 1.0, 1.0, 1.0, 1.0, 1e-4]) - J1, _ = jparse_pseudoinverse( - J, - gamma=0.1, - singular_direction_gain_position=1.0, - singular_direction_gain_angular=1.0, - position_dimensions=3, - angular_dimensions=3, - ) - J2, _ = jparse_pseudoinverse( - J, - gamma=0.1, - singular_direction_gain_position=1.0, - singular_direction_gain_angular=3.0, - position_dimensions=3, - angular_dimensions=3, - ) - assert not np.allclose(J1, J2) - - -class TestComparisonMethods: - """Test pinv and damped_least_squares comparison methods.""" - - def test_pinv_matches_numpy(self): - """pinv matches numpy's pinv.""" - J = np.array([[1.0, 0.5], [0.5, 1.0]]) - np.testing.assert_allclose(pinv(J), np.linalg.pinv(J), atol=1e-6) - - def test_dls_formula(self): - """DLS matches expected formula.""" - J = np.array([[1.0, 0.0], [0.0, 1.0]]) - d = 0.1 - J_dls = damped_least_squares(J, damping=d) - expected = np.linalg.inv(J.T @ J + d**2 * np.eye(2)) @ J.T - np.testing.assert_allclose(J_dls, expected, atol=1e-6) - - -class TestMetrics: - """Test manipulability and condition number metrics.""" - - def test_manipulability_identity(self): - """Identity Jacobian has manipulability 1.""" - J = np.eye(2) - m = manipulability_measure(J) - np.testing.assert_allclose(float(m), 1.0, atol=1e-6) - - def test_manipulability_near_singular_is_low(self): - """Near-singular has lower manipulability than identity.""" - J_good = np.eye(2) - J_bad = np.array([[1.0, 1.0], [1.0, 1.001]]) - assert float(manipulability_measure(J_bad)) < float( - manipulability_measure(J_good) - ) - - def test_inverse_condition_identity(self): - """Identity Jacobian has inverse condition number 1.""" - J = np.eye(2) - np.testing.assert_allclose(float(inverse_condition_number(J)), 1.0, atol=1e-6) - - def test_inverse_condition_range(self): - """Inverse condition number is in [0, 1].""" - rng = np.random.default_rng(42) - for _ in range(10): - J = rng.standard_normal((3, 4)) - icn = float(inverse_condition_number(J)) - assert 0 <= icn <= 1 - - -class TestEdgeCases: - """Test edge cases.""" - - def test_1x1_jacobian(self): - """1x1 Jacobian produces correct pseudo-inverse.""" - J = np.array([[2.0]]) - J_parse, _ = jparse_pseudoinverse(J, gamma=0.1) - assert J_parse.shape == (1, 1) - np.testing.assert_allclose(float(J_parse[0, 0]), 0.5, atol=1e-6) - - def test_tall_jacobian(self): - """Tall Jacobian (m > n) has correct shape.""" - J = np.array([[1.0], [0.5], [0.25]]) - J_parse, _ = jparse_pseudoinverse(J, gamma=0.1) - assert J_parse.shape == (1, 3) - - def test_very_small_gamma(self): - """Small gamma: well-conditioned case close to standard pinv.""" - J = np.array([[1.0, 0.5], [0.5, 1.0]]) - J_parse, _ = jparse_pseudoinverse(J, gamma=0.01) - np.testing.assert_allclose(J_parse, np.linalg.pinv(J), atol=0.01) - - def test_large_gamma(self): - """Large gamma still produces finite output.""" - J = np.array([[1.0, 0.5], [0.5, 1.0]]) - J_parse, _ = jparse_pseudoinverse(J, gamma=0.99) - assert np.all(np.isfinite(J_parse)) - - def test_missing_dimension_argument_raises(self): - """Providing only one dimension argument raises ValueError.""" - J = np.eye(6) - with pytest.raises( - ValueError, - match="Both position_dimensions and angular_dimensions must be provided.", - ): - jparse_pseudoinverse(J, position_dimensions=3) - - def test_dimension_sum_mismatch_raises(self): - """Dimension sum must match Jacobian row count.""" - J = np.eye(6) - with pytest.raises( - ValueError, - match="position_dimensions \\+ angular_dimensions must equal Jacobian row count.", - ): - jparse_pseudoinverse(J, position_dimensions=2, angular_dimensions=3) - - -class TestFeatureParity: - """Cross-validate JAX implementation against original jparse_robotics.""" - - @pytest.fixture(autouse=True) - def _skip_if_no_jparse(self): - jparse_robotics = pytest.importorskip("jparse_robotics") - self.JParseCore = jparse_robotics.JParseCore - self.np_manipulability = jparse_robotics.manipulability_measure - self.np_icn = jparse_robotics.inverse_condition_number - - @pytest.fixture() - def test_jacobians(self): - """Collection of test Jacobians.""" - rng = np.random.default_rng(42) - return { - "well_conditioned": np.array([[1.0, 0.5], [0.5, 1.0]]), - "near_singular": np.array([[1.0, 1.0], [1.0, 1.001]]), - "redundant": np.array([[1.0, 0.0, 1.0], [0.0, 1.0, 1.0]]), - "random_6x7": rng.standard_normal((6, 7)), - } - - def test_jparse_output_matches(self, test_jacobians): - """J-PARSE output matches original for multiple test cases.""" - gamma = 0.1 - core = self.JParseCore(gamma=gamma) - - for name, J in test_jacobians.items(): - J_np = core.compute(J) - row_count = J.shape[0] - J_jax, _ = jparse_pseudoinverse( - J, - gamma=gamma, - singular_direction_gain_position=1.0, - singular_direction_gain_angular=1.0, - position_dimensions=3 if row_count == 6 else row_count, - angular_dimensions=3 if row_count == 6 else 0, - ) - np.testing.assert_allclose( - np.array(J_jax), np.array(J_np), atol=1e-5, err_msg=f"Failed for {name}" - ) - - def test_manipulability_matches(self, test_jacobians): - """Manipulability measure matches original.""" - for name, J in test_jacobians.items(): - m_np = self.np_manipulability(J) - m_jax = float(manipulability_measure(J)) - # Relaxed tolerance: JAX defaults to float32 vs numpy float64. - np.testing.assert_allclose( - m_jax, m_np, rtol=0.2, atol=1e-4, err_msg=f"Failed for {name}" - ) - - def test_inverse_condition_matches(self, test_jacobians): - """Inverse condition number matches original.""" - for name, J in test_jacobians.items(): - icn_np = self.np_icn(J) - icn_jax = float(inverse_condition_number(J)) - np.testing.assert_allclose( - icn_jax, icn_np, atol=1e-6, err_msg=f"Failed for {name}" - ) - - def test_dls_matches(self, test_jacobians): - """DLS matches original.""" - core = self.JParseCore(gamma=0.1) - damping = 0.05 - for name, J in test_jacobians.items(): - J_np = core.damped_least_squares(J, damping=damping) - J_jax = damped_least_squares(J, damping=damping) - # Relaxed tolerance: JAX defaults to float32 vs numpy float64. - np.testing.assert_allclose( - np.array(J_jax), np.array(J_np), atol=1e-4, err_msg=f"Failed for {name}" - ) - - -class TestIntegration: - """Integration test with pyroki robot (no Viser).""" - - @pytest.fixture() - def panda_robot(self): - """Load Panda robot.""" - import pyroki as pk - from robot_descriptions.loaders.yourdfpy import load_robot_description - - urdf = load_robot_description("panda_description") - robot = pk.Robot.from_urdf(urdf) - target_link_index = robot.links.names.index("panda_hand") - return robot, target_link_index - - def test_position_error_decreases(self, panda_robot): - """Position error decreases over several IK steps.""" - robot, target_link_index = panda_robot - - # Start from joint midpoints. - cfg = np.array((robot.joints.lower_limits + robot.joints.upper_limits) / 2.0) - - # Target position (reachable). - target_pos = np.array([0.5, 0.1, 0.4]) - - errors = [] - for _ in range(50): - cfg, info = jparse_step( - robot=robot, - cfg=cfg, - target_link_index=target_link_index, - target_position=target_pos, - gamma=0.1, - dt=0.02, - ) - errors.append(info["position_error"]) - - # Error should decrease overall. - assert errors[-1] < errors[0] - - def test_joint_limits_respected(self, panda_robot): - """All configs remain within joint limits.""" - robot, target_link_index = panda_robot - cfg = np.array((robot.joints.lower_limits + robot.joints.upper_limits) / 2.0) - target_pos = np.array([0.5, 0.1, 0.4]) - - lower = np.array(robot.joints.lower_limits) - upper = np.array(robot.joints.upper_limits) - - for _ in range(30): - cfg, _ = jparse_step( - robot=robot, - cfg=cfg, - target_link_index=target_link_index, - target_position=target_pos, - dt=0.02, - ) - assert np.all(cfg >= lower - 1e-6) - assert np.all(cfg <= upper + 1e-6) - - def test_info_dict_keys(self, panda_robot): - """Info dict contains expected keys.""" - robot, target_link_index = panda_robot - cfg = np.array((robot.joints.lower_limits + robot.joints.upper_limits) / 2.0) - target_pos = np.array([0.5, 0.1, 0.4]) - - _, info = jparse_step( - robot=robot, - cfg=cfg, - target_link_index=target_link_index, - target_position=target_pos, - ) - expected_keys = { - "position_error", - "orientation_error", - "max_joint_vel", - "jacobian", - "manipulability", - "inverse_condition_number", - } - assert set(info.keys()) == expected_keys - - def test_orientation_tracking(self, panda_robot): - """Orientation tracking uses 6-DOF Jacobian and reduces error.""" - import jaxlie - - robot, target_link_index = panda_robot - cfg = np.array((robot.joints.lower_limits + robot.joints.upper_limits) / 2.0) - - # Use a small orientation offset from the initial EE pose so - # the controller can converge in limited steps. - import jax.numpy as jnp - - poses = robot.forward_kinematics(jnp.asarray(cfg)) - init_pose = jaxlie.SE3(poses[target_link_index]) - target_pos = np.array(init_pose.translation()) + np.array([0.05, 0.0, 0.0]) - target_wxyz = np.array(init_pose.rotation().wxyz) - - errors = [] - for _ in range(100): - cfg, info = jparse_step( - robot=robot, - cfg=cfg, - target_link_index=target_link_index, - target_position=target_pos, - target_wxyz=target_wxyz, - gamma=0.1, - dt=0.02, - ) - errors.append(info["position_error"]) - - # Position error should decrease. - assert errors[-1] < errors[0] - # Jacobian should be 6xn when tracking orientation. - assert info["jacobian"].shape[0] == 6 - # Orientation error should be present and finite. - assert np.isfinite(info["orientation_error"]) - - def test_orientation_convergence_large_rotation(self, panda_robot): - """Orientation converges with a non-trivial (~45 deg) rotation offset. - - This exercises the geometric Jacobian path: the analytical Jacobian - introduces artificial singularities at large rotations that prevent - convergence, while the geometric Jacobian handles them correctly. - """ - import jax.numpy as jnp - import jaxlie - - robot, target_link_index = panda_robot - cfg = np.array((robot.joints.lower_limits + robot.joints.upper_limits) / 2.0) - - poses = robot.forward_kinematics(jnp.asarray(cfg)) - init_pose = jaxlie.SE3(poses[target_link_index]) - - # Keep position the same, apply a ~45-degree rotation about the z-axis. - target_pos = np.array(init_pose.translation()) - angle = np.pi / 4 # 45 degrees - offset_rot = jaxlie.SO3.from_z_radians(angle) - target_rot = offset_rot @ init_pose.rotation() - target_wxyz = np.array(target_rot.wxyz) - - ori_errors = [] - for _ in range(200): - cfg, info = jparse_step( - robot=robot, - cfg=cfg, - target_link_index=target_link_index, - target_position=target_pos, - target_wxyz=target_wxyz, - gamma=0.1, - dt=0.02, - ) - ori_errors.append(info["orientation_error"]) - - # Orientation error should converge to a small value. - assert ori_errors[-1] < 0.15, ( - f"Orientation error did not converge: {ori_errors[-1]:.4f}" - ) From c124ef8a1924ef0cf250b049ea49ef1b789da751 Mon Sep 17 00:00:00 2001 From: Chung Min Kim Date: Sun, 8 Feb 2026 11:31:16 -0800 Subject: [PATCH 09/11] Format GUI slider parameter in singularity aware IK example --- examples/14_singularity_aware_ik.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/examples/14_singularity_aware_ik.py b/examples/14_singularity_aware_ik.py index 22708728..b961e2f4 100644 --- a/examples/14_singularity_aware_ik.py +++ b/examples/14_singularity_aware_ik.py @@ -40,7 +40,9 @@ def main(): "Gamma (singularity threshold)", 0.01, 0.5, 0.01, 0.1 ) nullspace_gain_handle = server.gui.add_slider("Nullspace gain", 0.0, 2.0, 0.05, 0.5) - manipulability_handle = server.gui.add_number("Manipulability", 0.001, disabled=True) + manipulability_handle = server.gui.add_number( + "Manipulability", 0.001, disabled=True + ) show_ellipsoid_handle = server.gui.add_checkbox("Show ellipsoid", True) # Manipulability ellipsoid visualization. From 36e7b58be00073bb3294366a7d5eb4c625a70b30 Mon Sep 17 00:00:00 2001 From: Chung Min Kim Date: Sun, 8 Feb 2026 11:31:22 -0800 Subject: [PATCH 10/11] Improve type hints and array handling in J-PARSE implementation --- examples/pyroki_snippets/_jparse.py | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/examples/pyroki_snippets/_jparse.py b/examples/pyroki_snippets/_jparse.py index 7de64d8c..fcf419bc 100644 --- a/examples/pyroki_snippets/_jparse.py +++ b/examples/pyroki_snippets/_jparse.py @@ -53,7 +53,7 @@ def compute_jacobian( anchor_poses = robot.forward_kinematics(cfg) R_anchor_inv = jaxlie.SE3(anchor_poses[target_link_index]).rotation().inverse() - def get_pose_components(q: ArrayLike) -> jnp.ndarray: + def get_pose_components(q: jax.Array) -> jnp.ndarray: poses = robot.forward_kinematics(q) pose = jaxlie.SE3(poses[target_link_index]) relative_rotation = pose.rotation() @ R_anchor_inv @@ -273,30 +273,32 @@ def jparse_step( pos_error_mag = float(jnp.linalg.norm(pos_error)) # Build desired task-space velocity. + omega_error = jnp.zeros(3) if position_only: v_des = position_gain * pos_error else: - target_wxyz = jnp.asarray(target_wxyz) - target_wxyz = target_wxyz / jnp.linalg.norm(target_wxyz) + assert target_wxyz is not None + tw = jnp.asarray(target_wxyz) + tw = tw / jnp.linalg.norm(tw) current_wxyz = target_pose.rotation().wxyz current_wxyz = current_wxyz / jnp.linalg.norm(current_wxyz) # Ensure shortest-path quaternion. - target_wxyz = jnp.where( - jnp.dot(target_wxyz, current_wxyz) < 0, -target_wxyz, target_wxyz - ) + tw = jnp.asarray(jnp.where(jnp.dot(tw, current_wxyz) < 0, -tw, tw)) # Orientation error via SO3 log map. q_current = jaxlie.SO3(current_wxyz) - q_target = jaxlie.SO3(target_wxyz) + q_target = jaxlie.SO3(tw) omega_error = (q_target @ q_current.inverse()).log() # Clamp orientation error magnitude. omega_mag = jnp.linalg.norm(omega_error) max_omega = 1.0 # rad - omega_error = jnp.where( - omega_mag > max_omega, omega_error * max_omega / omega_mag, omega_error + omega_error = jnp.asarray( + jnp.where( + omega_mag > max_omega, omega_error * max_omega / omega_mag, omega_error + ) ) v_des = jnp.concatenate( From 17a1081621a5a82daf807169c831d5d91e75f369 Mon Sep 17 00:00:00 2001 From: Chung Min Kim Date: Tue, 31 Mar 2026 17:07:29 -0700 Subject: [PATCH 11/11] Fix reference url --- examples/pyroki_snippets/_jparse.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/pyroki_snippets/_jparse.py b/examples/pyroki_snippets/_jparse.py index fcf419bc..cc9261d7 100644 --- a/examples/pyroki_snippets/_jparse.py +++ b/examples/pyroki_snippets/_jparse.py @@ -5,7 +5,7 @@ Effectively) provides singularity-aware inverse kinematics by computing a modified pseudo-inverse that handles singular configurations smoothly. -Reference: https://github.com/chungmin99/jparse +Reference: https://github.com/armlabstanford/jparse """ from __future__ import annotations