inverse_kinematics.py

#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

# -------------------------------------------------------------------------- #
# ----------------- 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:
# 1) solve IK with body/spatial jacobian
# 2) solve IK as NLS optimization
import os

import torch

import theseus as th
from torchkin.forward_kinematics import Robot, get_forward_kinematics_fns
from torchlie.functional import SE3 as SE3_Func

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.
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)
link_names = ["panda_virtual_ee_link"]

# 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: returns a tuple of link poses in the order of link names
# - jfk_b: returns a tuple where the first is a list of link body jacobians, and
#          the second is a tuple of link poses---both are in the order of link names
# - jfk_s: same as jfk_b except returning the spatial jacobians
fk, jfk_b, jfk_s = get_forward_kinematics_fns(robot, link_names)


# ********************************************************
# ** INVERSE KINEMATICS WITH Body/Spatial Jacobian
# ********************************************************
# If jfk is body jacobian, delta_theta is computed by
#          pose * exp(jfk * delta_theta) = targeted_pose,
# which has a closed-form solution:
# .  delta_theta = jfk.pinverse() * log(pose^-1 * targeted_pose)
# Otherwise, if jfk is spatial jacobian, delta_theta is computed by
#          exp(jfk * delta_theta) * pose = targeted_pose
def compute_delta_theta(jfk, theta, targeted_pose, use_body_jacobian):
    jac, poses = jfk(theta)
    pose_inv = SE3_Func.inv(poses[-1])
    error = (
        SE3_Func.log(
            SE3_Func.compose(pose_inv, targeted_pose)
            if use_body_jacobian
            else SE3_Func.compose(targeted_pose, pose_inv)
        )
        .view(-1, 6, 1)
        .view(-1, 6, 1)
    )
    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_body_or_spatial_jacobian(
    jfk, 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, theta_opt, target_pose, use_body_jacobian
        )
        print(error)
        if error < 1e-4:
            break
        theta_opt = theta_opt + step_size * delta_theta


print("---------------------------------------------------")
print("Use Body Jacobian")
print("---------------------------------------------------")
run_ik_using_body_or_spatial_jacobian(jfk_b, 100, 0.5, use_body_jacobian=True)

print("---------------------------------------------------")
print("Use Spatial Jacobian")
print("---------------------------------------------------")
run_ik_using_body_or_spatial_jacobian(jfk_s, 100, 0.2, use_body_jacobian=False)


# *********************************************
# ** INVERSE KINEMATICS AS NLS OPTIMIZATION
# *********************************************
print("---------------------------------------------------")
print("Theseus Optimizer")
print("---------------------------------------------------")


# IK can also be solved as an optimization problem:
#      min \|log(pose^-1 * targeted_pose)\|^2
# as the following
def targeted_pose_error(optim_vars, aux_vars):
    (theta,) = optim_vars
    (targeted_pose,) = aux_vars
    pose = th.SE3(tensor=fk(theta.tensor)[0])
    return pose.local(targeted_pose)


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=target_pose, name="targeted_pose"),)


cost_function = th.AutoDiffCostFunction(
    optim_vars,
    targeted_pose_error,
    6,
    aux_vars=aux_vars,
    name="targeted_pose_error",
    cost_weight=th.ScaleCostWeight(torch.ones([1], dtype=dtype)),
    autograd_mode="vmap",
)
objective = th.Objective(dtype=dtype)
objective.add(cost_function)
optimizer = th.LevenbergMarquardt(
    objective,
    max_iterations=15,
    step_size=0.5,
    vectorize=True,
)

inputs = {"theta_opt": torch.zeros_like(theta_opt), "targeted_pose": target_pose}
optimizer.objective.update(inputs)
optimizer.optimize(verbose=True)