Added more explanations to example and cleanup some of the code.

examples/labs/inverse_kinematics.py

@@ -3,11 +3,13 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
-#
-# This example illustrates the four backward modes
-# (unroll, implicit, truncated, and dlm)
-# on a problem fitting a quadratic to data.
 
+# -------------------------------------------------------------------------- #
+# ----------------- INVERSE KINEMATICS WITH THESEUS AND TORCHLIE ----------- #
+# -------------------------------------------------------------------------- #
+# In this example we show how to do inverse kinematics to achieve a target
+# pose for a specific end effector of a robot described in a URDF file.
+# We show two ways to accomplish this: <ADD NAME OF THE TWO METHODS BELOW>.
 import os
 
 import torch
@@ -20,14 +22,31 @@
 )
 
 dtype = torch.float64
+
+# First we load the URDF file describing the robot and create a `Robot` object to
+# represent it in Python. The `Robot` class can be used to build a kinematics tree
+# of the robot.  <COMPLETE WITH OTHER THINGS THIS CLASS DOES, IF ANY>
 URDF_REL_PATH = (
     "../../tests/theseus_tests/embodied/kinematics/data/panda_no_gripper.urdf"
 )
 urdf_path = os.path.join(os.path.dirname(__file__), URDF_REL_PATH)
-
 robot = Robot.from_urdf_file(urdf_path, dtype)
-selected_links = ["panda_virtual_ee_link"]
-fk, jfk_b, jfk_s = get_forward_kinematics_fns(robot, selected_links)
+
+# We can get differentiable forward kinematics functions for specific links
+# by using `get_forward_kinematics_fns`. This function creates three differentiable
+# functions for evaluating forward kinematics, body jacobian and spatial jacobian of
+# the selected links, in that order. The return types of these functions are as
+# follows:
+#
+# - fk: <ADD DESCRIPTION>
+# - jfk_b: <ADD DESCRIPTION>
+# - jfk_s: <ADD DESCRIPTION>. (If return type is the same as jfk_b, just say this).
+fk, jfk_b, jfk_s = get_forward_kinematics_fns(robot, ["panda_virtual_ee_link"])
+
+
+# *********************************************
+# ** INVERSE KINEMATICS WITH METHOD XXXXXXXXX
+# *********************************************
 
 
 # If jfk is body jacobian, delta_theta is computed by
@@ -51,33 +70,41 @@ def compute_delta_theta(jfk, theta, targeted_pose, use_body_jacobian):
     return (jac[-1].pinverse() @ error).view(-1, robot.dof), error.norm().item()
 
 
+# The following function runs IK. We first create random joint angles to compute
+# a set of target poses for the end effector. Then, starting from zero joint angles,
+# we try to recover joint angles that can achieve the target pose, by iteratively
+# updating the joint angles using the body/spatial jacobian. Note that IK has infinite
+# solutions, so, while we are able to recover the target poses after a few iterations,
+# there is no guarantee that we can recover the joint angles used to generate these
+# poses in the first place.
+def run_ik_using_XYZ_method(batch_size, step_size, use_body_jacobian):
+    target_theta = torch.rand(batch_size, robot.dof, dtype=dtype)
+    target_pose: torch.Tensor = fk(target_theta)[0]
+    theta_opt = torch.zeros_like(target_theta)
+    for _ in range(50):
+        delta_theta, error = compute_delta_theta(
+            jfk_b, theta_opt, target_pose, use_body_jacobian
+        )
+        print(error)
+        if error < 1e-4:
+            break
+        theta_opt = theta_opt + step_size * delta_theta
+
+
 print("---------------------------------------------------")
 print("Body Jacobian")
 print("---------------------------------------------------")
-targeted_theta = torch.rand(100, robot.dof, dtype=dtype)
-targeted_pose: torch.Tensor = fk(targeted_theta)[0]
-theta_opt = torch.zeros_like(targeted_theta)
-for iter in range(50):
-    delta_theta, error = compute_delta_theta(jfk_b, theta_opt, targeted_pose, True)
-    print(error)
-    if error < 1e-4:
-        break
-    theta_opt = theta_opt + 0.5 * delta_theta
+run_ik_using_XYZ_method(100, 0.5, use_body_jacobian=True)
 
 print("---------------------------------------------------")
 print("Spatial Jacobian")
 print("---------------------------------------------------")
-targeted_theta = torch.rand(10, robot.dof, dtype=dtype)
-targeted_pose = fk(targeted_theta)[0]
-theta_opt = torch.zeros_like(targeted_theta)
-for iter in range(50):
-    delta_theta, error = compute_delta_theta(jfk_s, theta_opt, targeted_pose, False)
-    print(error)
-    if error < 1e-4:
-        break
-    theta_opt = theta_opt + 0.2 * delta_theta
+run_ik_using_XYZ_method(10, 0.2, use_body_jacobian=True)
 
 
+# *********************************************
+# ** INVERSE KINEMATICS AS LS OPTIMIZATION
+# *********************************************
 print("---------------------------------------------------")
 print("Theseus Optimizer")
 print("---------------------------------------------------")
@@ -93,8 +120,11 @@ def targeted_pose_error(optim_vars, aux_vars):
     return (pose.inverse().compose(targeted_pose)).log_map()
 
 
+target_theta = torch.rand(10, robot.dof, dtype=dtype)
+target_pose: torch.Tensor = fk(target_theta)[0]
+theta_opt = torch.zeros_like(target_theta)
 optim_vars = (th.Vector(tensor=torch.zeros_like(theta_opt), name="theta_opt"),)
-aux_vars = (th.SE3(tensor=targeted_pose, name="targeted_pose"),)
+aux_vars = (th.SE3(tensor=target_pose, name="targeted_pose"),)
 
 
 cost_function = th.AutoDiffCostFunction(
@@ -115,6 +145,6 @@ def targeted_pose_error(optim_vars, aux_vars):
     vectorize=True,
 )
 
-inputs = {"theta_opt": torch.zeros_like(theta_opt), "targeted_pose": targeted_pose}
+inputs = {"theta_opt": torch.zeros_like(theta_opt), "targeted_pose": target_pose}
 optimizer.objective.update(inputs)
 optimizer.optimize(verbose=True)