Simple & compact damped least-squares IK solver in <100 lines of python

ik.py

'''
james.gregson@gmail.com

A pretty small 3D IK solver, < 100 lines of code but > 100 lines of comments.
Low dependency (numpy, opencv-python, matplotlib) and includes visualization.

Not particularly fast.

Solves for joint rotations needed in a kinematic skeleton to minimize errors of a set
of end-effectors with respect to a corresponding set of target positions. Assumes
the root of the kinematic chain is fixed at the origin.

The state of the kinematic chain is defined as three lists:

- Parent list: stores the index each bone's parent, with -1 for the single root 
    bone. Bones must be listed so their parent precedes them in the list, i.e.
    parent[i] < i for all i.

- transform list: for each bone, stores it's world-space 4x4 transformation
    relative to it's parent, or identity for the root.

- theta: for each bone, stores a 3-tuple of exponential map rotation vectors


Goal positions of the skeleton are specified as a list of 3-tuples, each having:

- bone_id: the index of the bone that the goal affects

- bone_pos: 3d body-local position of the bone end-effector in the bone's 
    reference frame, i.e. the transforms described above.

- world_pos: 3d world-frame target position for the end-effector


The algorithm is based on a damped least-squares formulation and consists of 
the following steps:

- A forward pass takes the existing angles and computes the world-space 
    transforms of all the bones by concatenating transforms and also computes
    the Jacobian of each transform matrix

- A goal pass computes it's current end-effector position and the Jacobian
    of the same with respect to each affected bone in the heirarchy (all 
    bones from the bone_id'th back to the root). This completes a row-block
    of the system Jacobian for each goal.

- A damped Newton step is taken to compute the angle update. The damping is
    applied based on the diagonals of J'J. The predicted step is down-scaled
    by a factor for stability.


Some discussion of the steps is probably in order:

** Forward Pass **
The concatenation of transforms is pretty typical but the construction of 
Jacobians may not be. Each rotation leads to a 4x4 homogeneous transform
matrix based on the 3 bone angles.

Most methods seem to compute Jacobians of end-effector position with 
respect to the input angles, which gives a 3x3 Jacobian. But this means 
that it is hard to separate the formulation of a goal from the general
kinematic chain. Instead, I'm computing the Jacobian of the transform
with respect to the angles so that goal formulation can be decoupled 
from chain traversal. I don't know that it gains much in practice.

Anyway, this leads to a 4x4x3 Jacobian tensor for the transform
with respect to the input angles. This uses a convention that the leading
dimensions match the output of an operation and the trailing dimensions 
match the input. Hence 4x4 matrices from 3 parameters gives 4x4x3 Jacobians.

However, this Jacobian is relative to the concatenation of the parent
transform with the rest transform of the bone. So each of the three
4x4 Jacobian components need to be transformed by this. Similarly, it's
convenient to define the Jacobian in world-space, so the computed
end-effector positions can be apply directly to the Jacobians to get
the change in world-space end-effector position w.r.t. the joint angles.
This ends up giving:

Ji = T[parent[i]] @ Trel[i] @ ji @ inv(T[i])

where Jgi is the 4x4x3 Jacobian of the i'th bone's transform in world
space with respect to points defined in world-space. T[a] is the
world-space transform of the a'th bone and ji is the bone-local
transform Jacobian with respect to bone angles.

Except, batch matrix multiplication with numpy expects the leading 
dimensions to be batch dimensions, so in practice it's way easier
to actually transpose the Jacobian to be 3x4x4 and use broadcasting 
to compute the whole mess as a batch....so that's what the code does.

** Goal Pass **
The goal pass computes the world-space end-effector position and it's
Jacobian with respect to each bone back to the root for each goal. With 
the bone transform Jacobians computed as before, this can be implemented 
as a multiplication of the goal position with each component of the 
bone Jacobians. Except to be consistent with the system matrix, these
are concatenated together as a new trailing dimension, rather than
a leading dimension as above.

** Damped-Newton **
The damped Newton step solves for updated rotations by linearizing
the non-linear problem at the current angles. For a single end-effector
with goal position G approximate residuals are:

r = F(theta) - G ~= J*d

where F(theta) is the forward kinematics, J is the system Jacobian, 
d is the update and G are the target end-effectors. 

This solver solves the following:

(J'J + reg_weight*diag(J'J)) d = J'r

Lambda is a damping factor for the regularizing term of diag(J'J). 
Theta is then updated as:

theta -= d*relax

where relax is in (0,1] to relax the steps that are taken.

** General Comments **
For the goals in this solver, it doesn't really make sense to use
the Jacobian approach described above. The Jacobians of end-effector
position could be computed directly and save memory/computation.

However this approach is quite flexible and can easily be extended 
to other types of goals like direction/orientation since the 
Jacobian of the entire transform is available rather than just 
of the position. It also allows the goals to be abstracted from 
the chain updates, which might be beneficial. Who knows?
'''
import cv2
import time
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def main():
    # chain length & world-space transformations of bones
    N = 6
    T = [ np.eye(4) ] + [ translate((i+1,0,0)) for i in range(N-1) ]

    # kinematic chain parents, relative rest transforms and initial angles
    parent = [ i-1 for i in range(N) ]
    Trest  = [ np.eye(4) ] + [ np.linalg.inv(T[i-1])@T[i] for i in range(1,N) ]
    theta  = np.random.standard_normal((N,3))*0.5

    # list of goal bone-id, relative position on bone and world-space target point
    goals = [
        (N-1, (1.0,0.1,0.2), np.random.standard_normal(3)),
        (N-4, (1.0,0.0,0.0), np.random.standard_normal(3))
    ]

    # solve loop
    ax = None
    for i in range(100):
        dt = time.time()
        resid, theta, bone_T = solve( parent, Trest, theta, goals, verbose=True )
        dt = time.time()-dt
        print('Iteration {}, ||residual||={:0.5f}, dt={:0.3f}ms'.format(i,resid,1000*dt) ) 
        ax = draw( ax, parent, bone_T, goals )
        plt.pause(1.0)

def translate(xyz):
    T = np.eye(4)
    T[:3,3] = xyz
    return T

def rotate(ang):
    T = np.eye(4)
    J = np.zeros((3,4,4))
    T[:3,:3], tmp = cv2.Rodrigues( ang )
    J[:,:3,:3] = tmp.reshape((3,3,3))
    return T, J

def forward_pass( parent, Trest, theta ):
    t,j = rotate(theta[0])
    T = [ Trest[0]@t ]
    J = [ j@np.linalg.inv(T[-1])[None,...] ]
    for i in range(1,len(Trest)):
        t,j = rotate(theta[i])
        curr = T[parent[i]]@Trest[i]@t
        j = T[parent[i]][None,...]@j@np.linalg.inv(curr)[None,...]
        T.append( curr )
        J.append( j )
    return np.array(T), np.array(J)

def goal_pass( parent, bone_T, bone_J, bone_id, bone_pos ):
    W = bone_T[bone_id]@(*bone_pos,1) 
    J = []
    bid = bone_id
    while bid > 0:
        J.append( (bid,np.stack([(bone_J[bid,i]@W)[:3] for i in range(3)],axis=-1) ) )
        bid = parent[bid]
    return W[:3],J

def solve( parent, T_rest, theta, goals, reg_weight=0.05, relax=0.75, verbose=False ):
    bone_T, bone_J = forward_pass( parent, T_rest, theta )
    J = np.zeros( (3*len(goals),3*len(parent)) )
    e = np.zeros( 3*len(goals) )
    for gid,(bid,bpos,gpos) in enumerate(goals):
        end,gjac = goal_pass( parent, bone_T, bone_J, bid, bpos )
        for idx,bJ in gjac:
            J[gid*3:gid*3+3,idx*3:idx*3+3] = bJ
            e[gid*3:gid*3+3] = (end-gpos)
    JtJ = J.T@J
    JtJ += reg_weight*np.diag( np.maximum( 1e-6, np.diag( JtJ ) ) )
    dtheta = np.linalg.solve( JtJ, J.T@e )
    theta  -= relax*dtheta.reshape((-1,3))
    return np.linalg.norm(e), theta, bone_T

def draw( ax, parent, bone_T, goals ):
    ax = ax if ax else plt.subplot(111,projection='3d')
    ax.clear()
    pnts = np.array([ T[:3,3] for T in bone_T ])
    plt.plot( pnts[:,0], pnts[:,1], pnts[:,2], 'k-' )
    for bid,bpos,gpos in goals:
        end = (bone_T[bid]@(*bpos,1))[:3]
        plt.plot( [pnts[bid,0],end[0]], [pnts[bid,1],end[1]], [pnts[bid,2],end[2]], 'k-' )
        plt.plot( [end[0],gpos[0]],[end[1],gpos[1]],[end[2],gpos[2]], 'r-' )
        plt.plot( [gpos[0]],[gpos[1]],[gpos[2]], 'ro' )
    ax.set_xlim(-4,4)
    ax.set_ylim(-4,4)
    ax.set_zlim(-4,4)
    return ax
 
if __name__ == '__main__':
    main()