From d82d3dadcea6ce3bb5ef1a7bce6d4a6b0e82e472 Mon Sep 17 00:00:00 2001 From: sh8 Date: Wed, 18 Mar 2026 09:50:39 -0700 Subject: [PATCH] Fixed a bug in Jacobian computation with pyroki --- examples/pyroki_3d_visualization.py | 4 ++-- examples/pyroki_jparse_utils.py | 33 ++++++++++++++++++++++++----- 2 files changed, 30 insertions(+), 7 deletions(-) diff --git a/examples/pyroki_3d_visualization.py b/examples/pyroki_3d_visualization.py index 2fd8619..f0eb303 100644 --- a/examples/pyroki_3d_visualization.py +++ b/examples/pyroki_3d_visualization.py @@ -182,7 +182,7 @@ def main(): robot, target_link_name, gamma=0.1, - position_only=True, + position_only=False, ) # Initial configuration (middle of joint range) @@ -407,7 +407,7 @@ def on_reset(_): # Get target from gizmo (position only for now) target_pos = np.array(ik_target.position) - target_wxyz = None + target_wxyz = np.array(ik_target.wxyz) # Compute control step cfg, info = controller.step(cfg, target_pos, target_wxyz=target_wxyz, dt=dt) diff --git a/examples/pyroki_jparse_utils.py b/examples/pyroki_jparse_utils.py index 06979a4..02e49e8 100644 --- a/examples/pyroki_jparse_utils.py +++ b/examples/pyroki_jparse_utils.py @@ -53,13 +53,36 @@ def compute_jacobian_from_pyroki( lambda q: jaxlie.SE3(robot.forward_kinematics(q)).translation() )(cfg)[target_link_index] else: - # Full Jacobian via autodiff on SE3 log - def get_pose_log(q): + def get_pos_and_R_flat(q): poses = robot.forward_kinematics(q) target_pose = jaxlie.SE3(poses[target_link_index]) - return target_pose.log() - - jacobian = jax.jacfwd(get_pose_log)(cfg) + return jnp.concatenate([ + target_pose.translation(), # (3,) + target_pose.rotation().as_matrix().reshape(-1), # (9,) + ]) # (12,) + + J_combined = jax.jacfwd(get_pos_and_R_flat)(cfg) # (12, n_joints) + J_pos = J_combined[:3, :] # (3, n_joints) + dR_flat_dq = J_combined[3:, :] # (9, n_joints) + + # Geometric (spatial) angular Jacobian: + # omega_spatial = unskew(dR/dq_i @ R^T) for each joint i + # This is the correct Jacobian to use with world-frame orientation error + # omega_error = log(q_target @ q_current^{-1}). + dR_dq = np.array(dR_flat_dq).reshape(3, 3, -1) # (3, 3, n_joints) + R_mat = np.array( + jaxlie.SE3(robot.forward_kinematics(cfg)[target_link_index]).rotation().as_matrix() + ) + # omega_skew_all[a,b,i] = (dR/dq_i @ R^T)[a,b] + omega_skew_all = np.einsum('acj,bc->abj', dR_dq, R_mat) # (3, 3, n_joints) + # unskew: [skew[2,1], skew[0,2], skew[1,0]] = [omega_x, omega_y, omega_z] + J_ang = np.array([ + omega_skew_all[2, 1], + omega_skew_all[0, 2], + omega_skew_all[1, 0], + ]) # (3, n_joints) + + jacobian = np.vstack([J_pos, J_ang]) # (6, n_joints) return np.array(jacobian)