From 02df8004b77d4a42759764c296bf3674d7cb1367 Mon Sep 17 00:00:00 2001 From: Giridharshan Date: Fri, 26 Jun 2026 00:59:09 +0530 Subject: [PATCH] Implement forward and inverse kinematics logic --- planning/dh_target.py | 50 +++++++++++++++++++++++++++----- planning/fk_logic_target.py | 10 ++----- planning/ik_logic_target.py | 58 ++++++++++++++++++++++++++++++++----- scripts/fkv2.py | 4 +-- scripts/ik_viewv2.py | 34 +++++++++++++++++----- scripts/ikv2.py | 4 +-- scripts/validate_sol.py | 6 ++-- 7 files changed, 130 insertions(+), 36 deletions(-) diff --git a/planning/dh_target.py b/planning/dh_target.py index edad8b5..58c2ac6 100644 --- a/planning/dh_target.py +++ b/planning/dh_target.py @@ -4,9 +4,9 @@ # TODO: Fill in the DH Parameters (a, d, alpha) DH_PARAMS = np.array([ - [0.0, 0.0, 0.0], - [0.0, 0.0, 0.0], - [0.0, 0.0, 0.0], + [0.0, 0.0595, np.pi/2], + [0.1175, 0.0, 0.0], + [0.0950, 0.0, 0.0], ], dtype=float) EE_OFFSET = 0.0100 @@ -23,7 +23,18 @@ def dh_transform(a, d, alpha, theta): TODO: Implement the standard DH transformation matrix. Return a 4x4 numpy array. """ - return np.eye(4) + ct = np.cos(theta) + st = np.sin(theta) + ca = np.cos(alpha) + sa = np.sin(alpha) + + return np.array([ + [ct, -st * ca, st * sa, a * ct], + [st, ct * ca, -ct * sa, a * st], + [0.0, sa, ca, d ], + [0.0, 0.0, 0.0, 1.0 ] + ], dtype=float) + def fk(q): @@ -31,19 +42,44 @@ def fk(q): TODO: Implement the full forward kinematics using DH_PARAMS and EE_OFFSET. Return the 4x4 end-effector transformation matrix. """ - return np.eye(4) + T = np.eye(4, dtype=float) + + # Loop over the 3 joints to chain the matrices sequentially + for i in range(3): + a = DH_PARAMS[i, 0] + d = DH_PARAMS[i, 1] + alpha = DH_PARAMS[i, 2] + theta = q[i] + + # Multiply current total transformation by the next link's transformation + T = T @ dh_transform(a, d, alpha, theta) + + # Crucial Final Step: Account for the Tool Center Point (TCP) stretch. + # Move forward along the final local X-axis by EE_OFFSET (0.0100 m) + T_ee = np.array([ + [1.0, 0.0, 0.0, EE_OFFSET], + [0.0, 1.0, 0.0, 0.0 ], + [0.0, 0.0, 1.0, 0.0 ], + [0.0, 0.0, 0.0, 1.0 ] + ], dtype=float) + + T = T @ T_ee + + return T def position(q): """ TODO: Return just the (3,) position vector [X, Y, Z] from the FK matrix. """ - return np.zeros(3) + T = fk(q) + # Extract the top 3 rows of the 4th column (index 3) + return T[:3, 3].copy() def print_table(): print("Teaching DH table") - print(" i joint theta d (m) a (m) alpha (rad)") + print(" i joint theta d (m) a (m) alpha (rad)") for i, (name, (a, d, alpha)) in enumerate(zip(JOINT_NAMES, DH_PARAMS), start=1): print(f" {i:<2} {name:<15} q{i:<7} {d:8.4f} {a:8.4f} {alpha:12.6f}") print(f" end-effector offset along final x-axis: {EE_OFFSET:.4f} m") diff --git a/planning/fk_logic_target.py b/planning/fk_logic_target.py index b8745ea..6c0ddb3 100644 --- a/planning/fk_logic_target.py +++ b/planning/fk_logic_target.py @@ -4,18 +4,14 @@ ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.append(ROOT_DIR) -from planning import dh +from planning import dh_target as dh def fk(q): """ Calls the DH model to get the transform, and returns position and rotation. """ - # TODO: Use dh.fk(q) to get the 4x4 matrix, then extract position and rotation - # T = dh.fk(q) - # return T[:3, 3].copy(), T[:3, :3].copy() - - # Placeholder: - return np.zeros(3), np.eye(3) + T = dh.fk(q) + return T[:3, 3].copy(), T[:3, :3].copy() def print_fk_row(label, q, pos): deg = np.degrees(q) diff --git a/planning/ik_logic_target.py b/planning/ik_logic_target.py index 78d25d2..7a18c4b 100644 --- a/planning/ik_logic_target.py +++ b/planning/ik_logic_target.py @@ -4,8 +4,7 @@ ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.append(ROOT_DIR) -from planning import dh - +from planning import dh_target as dh class IKSolution: def __init__(self, q, converged, iterations, position_error): self.q = q @@ -24,15 +23,60 @@ def within_limits(q, limits): def solve(target, q0=None, elbow="down"): """ - TODO: Implement the analytical inverse kinematics for the 3-DOF arm. + Implement the analytical inverse kinematics for the 3-DOF arm. """ target = np.asarray(target, dtype=float).reshape(3) limits = dh.JOINT_LIMITS - # Placeholder: just return zeros - q = np.zeros(3) - err = float(np.linalg.norm(target - dh.position(q))) - return IKSolution(q=q, converged=False, iterations=0, position_error=err) + d1 = dh.DH_PARAMS[0, 1] + L2 = dh.DH_PARAMS[1, 0] + L3 = dh.DH_PARAMS[2, 0] + dh.EE_OFFSET + + # Try both base joint angles: q1_a and q1_b + q1_a = np.arctan2(target[1], target[0]) + q1_b = wrap_to_pi(q1_a + np.pi) + + candidates = [] + + for q1 in [q1_a, q1_b]: + # Coordinates in the plane + r = target[0] * np.cos(q1) + target[1] * np.sin(q1) + z_prime = target[2] - d1 + + # Solve for q3 using law of cosines + cos_q3 = (r**2 + z_prime**2 - L2**2 - L3**2) / (2 * L2 * L3) + if cos_q3 < -1.0 or cos_q3 > 1.0: + continue + + if elbow == "down": + q3 = np.arccos(cos_q3) + else: + q3 = -np.arccos(cos_q3) + + # Solve for q2 + A = L2 + L3 * np.cos(q3) + B = L3 * np.sin(q3) + cos_q2 = (A * r + B * z_prime) / (A**2 + B**2) + sin_q2 = (A * z_prime - B * r) / (A**2 + B**2) + q2 = np.arctan2(sin_q2, cos_q2) + + q = np.array([q1, q2, q3]) + q_clamped = clamp_to_limits(q, limits) + err = float(np.linalg.norm(target - dh.position(q_clamped))) + converged = within_limits(q, limits) and (err < 1e-4) + candidates.append((q_clamped, converged, err)) + + if len(candidates) == 0: + # No mathematically valid configuration + q = np.zeros(3) + err = float(np.linalg.norm(target - dh.position(q))) + return IKSolution(q=q, converged=False, iterations=1, position_error=err) + + # Select the candidate that is within limits and has smallest error + candidates.sort(key=lambda x: (not x[1], x[2])) + best_q, best_conv, best_err = candidates[0] + + return IKSolution(q=best_q, converged=best_conv, iterations=1, position_error=best_err) def print_solution(target, sol, fk_pos): p = fk_pos diff --git a/scripts/fkv2.py b/scripts/fkv2.py index 09e015d..4e966aa 100644 --- a/scripts/fkv2.py +++ b/scripts/fkv2.py @@ -6,7 +6,7 @@ ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.append(ROOT_DIR) -from planning import fk_logicv2 +from planning import fk_logic_target as fk_logicv2 def main(): parser = argparse.ArgumentParser(description="SO101 3-DOF Pure DH FK checks") @@ -15,7 +15,7 @@ def main(): parser.add_argument("--n", type=int, default=500, help="Number of sweep samples") args = parser.parse_args() - from planning import dh + from planning import dh_target as dh dh.print_table() if args.q is not None: diff --git a/scripts/ik_viewv2.py b/scripts/ik_viewv2.py index 140827c..f24a626 100644 --- a/scripts/ik_viewv2.py +++ b/scripts/ik_viewv2.py @@ -13,9 +13,27 @@ ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.append(ROOT_DIR) -from planning import ik_logicv2 -from planning import fk_logic -from planning import dh +from planning import ik_logic_target as ik_logicv2 +from planning import dh_target as dh + +JOINT_NAMES = ("shoulder_pan", "shoulder_lift", "elbow_flex") + +def load_model(xml_path): + if not os.path.isabs(xml_path): + xml_path = os.path.join(ROOT_DIR, xml_path) + model = mujoco.MjModel.from_xml_path(xml_path) + data = mujoco.MjData(model) + return model, data + +def keyframe_qpos(model, name): + key_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY, name) + if key_id < 0: + raise ValueError(f"Keyframe '{name}' not found") + return model.key_qpos[key_id * model.nq : (key_id + 1) * model.nq].copy() + +def set_qpos(model, data, q): + data.qpos[:len(q)] = q + mujoco.mj_forward(model, data) def run_viewer(model, data, target_q, title): import mujoco.viewer @@ -30,7 +48,7 @@ def run_viewer(model, data, target_q, title): print("Close the viewer window to exit.") # Set the actuators' control signals to the target IK solution - for name, q_val in zip(fk_logic.JOINT_NAMES, target_q): + for name, q_val in zip(JOINT_NAMES, target_q): act_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, name) if act_id >= 0: data.ctrl[act_id] = q_val @@ -51,7 +69,7 @@ def main(): args = parser.parse_args() # Load MuJoCo model for visualization - model, data = fk_logic.load_model(args.xml) + model, data = load_model(args.xml) target = np.array(args.target, dtype=float) seed = np.array(args.seed, dtype=float) if args.seed is not None else None @@ -63,10 +81,10 @@ def main(): # Set initial state to the 'folded' keyframe try: - folded_q = fk_logic.keyframe_qpos(model, "folded") - fk_logic.set_qpos(model, data, folded_q) + folded_q = keyframe_qpos(model, "folded") + set_qpos(model, data, folded_q) # Initialize controls to the folded position so it doesn't jerk violently on frame 1 - for name, q_val in zip(fk_logic.JOINT_NAMES, folded_q): + for name, q_val in zip(JOINT_NAMES, folded_q): act_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, name) if act_id >= 0: data.ctrl[act_id] = q_val diff --git a/scripts/ikv2.py b/scripts/ikv2.py index df8f720..394e970 100644 --- a/scripts/ikv2.py +++ b/scripts/ikv2.py @@ -6,8 +6,8 @@ ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.append(ROOT_DIR) -from planning import dh -from planning import ik_logicv2 +from planning import dh_target as dh +from planning import ik_logic_target as ik_logicv2 def main(): parser = argparse.ArgumentParser(description="Teaching DH pure IK for SO101 3-DOF") diff --git a/scripts/validate_sol.py b/scripts/validate_sol.py index 829a23f..570332b 100644 --- a/scripts/validate_sol.py +++ b/scripts/validate_sol.py @@ -6,9 +6,9 @@ ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.append(ROOT_DIR) -from planning import dh_target_sol as dh -from planning import fk_logic_target_sol as fk_logic -from planning import ik_logic_target_sol as ik_logic +from planning import dh_target as dh +from planning import fk_logic_target as fk_logic +from planning import ik_logic_target as ik_logic def validate(filename): fk_errors = []