Implement forward and inverse kinematics logic

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,27 +23,63 @@ 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):
     """
     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")

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)

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

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:

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

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")

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 = []