In [ ]:
import numpy as np
import jax
import jax.numpy as jnp
import sympy as sym
In [ ]:
# link masses
M1 = 3
M2 = 2
M3 = 1

# link lengths
LX = 0.1
LY = 0.1
L2 = 1
L3 = 0.5

# moments of inertia (base link doesn't rotate so don't need its inertia)
I2 = M2 * L2 ** 2 / 12
I3 = M3 * L3 ** 2 / 12

# inertia matrices
G1 = np.diag(np.array([M1, M1, 0]))
G2 = np.diag(np.array([M2, M2, I2]))
G3 = np.diag(np.array([M3, M3, I3]))

# gravity
G = 9.8
In [ ]:
class ManualDynamics:
    @classmethod
    def mass_matrix(cls, q):
        x1, θ2, θ3 = q
        θ23 = θ2 + θ3

        m11 = M1 + M2 + M3
        m12 = -(0.5 * M2 + M3) * L2 * np.sin(θ2) - 0.5 * M3 * L3 * np.sin(θ23)
        m13 = -0.5 * M3 * L3 * np.sin(θ23)

        m22 = (
            (0.25 * M2 + M3) * L2 ** 2
            + 0.25 * M3 * L3 ** 2
            + M3 * L2 * L3 * np.cos(θ3)
            + I2
            + I3
        )
        m23 = 0.5 * M3 * L3 * (0.5 * L3 + L2 * np.cos(θ3)) + I3

        m33 = 0.25 * M3 * L3 ** 2 + I3

        M = np.array([[m11, m12, m13], [m12, m22, m23], [m13, m23, m33]])
        return M

    @classmethod
    def christoffel_matrix(cls, q):
        x1, θ2, θ3 = q
        θ23 = θ2 + θ3

        # Partial derivatives of mass matrix
        dMdx1 = np.zeros((3, 3))

        dMdθ2_12 = (
            -0.5 * M2 * L2 * np.cos(θ2) - M3 * L2 * np.cos(θ2) - 0.5 * M3 * L3 * np.cos(θ23)
        )
        dMdθ2_13 = -0.5 * M3 * L3 * np.cos(θ23)
        dMdθ2 = np.array([[0, dMdθ2_12, dMdθ2_13], [dMdθ2_12, 0, 0], [dMdθ2_13, 0, 0]])

        dMdθ3_12 = -0.5 * M3 * L3 * np.cos(θ23)
        dMdθ3_13 = -0.5 * M3 * L3 * np.cos(θ23)
        dMdθ3_22 = -M3 * L2 * L3 * np.sin(θ3)
        dMdθ3_23 = -0.5 * M3 * L2 * L3 * np.sin(θ3)
        dMdθ3 = np.array([
            [0, dMdθ3_12, dMdθ3_13],
            [dMdθ3_12, dMdθ3_22, dMdθ3_23],
            [dMdθ3_13, dMdθ3_23, 0],
        ])

        dMdq = np.zeros((3, 3, 3))
        dMdq[0, :, :] = dMdx1
        dMdq[1, :, :] = dMdθ2
        dMdq[2, :, :] = dMdθ3

        # This is equivalent to a set of for loops that fill out Γ such that
        # Γ[i, j, k] = dMdq[k, j, i] - 0.5*dMdq[i, j, k]
        Γ = dMdq.T - 0.5 * dMdq
        return Γ

    @classmethod
    def coriolis_matrix(cls, q, dq):
        Γ = cls.christoffel_matrix(q)
        # note numpy dot broadcasting rules: this is a sum product of dq with 
        # the second last dimension of Γ
        return np.dot(dq, Γ)

    @classmethod
    def gravity_vector(cls, q):
        x1, θ2, θ3 = q
        θ23 = θ2 + θ3
        return np.array([
            0,
            (0.5 * M2 + M3) * G * L2 * np.cos(θ2) + 0.5 * M3 * L3 * G * np.cos(θ23),
            0.5 * M3 * L3 * G * np.cos(θ23),
        ])

    @classmethod
    def force(cls, q, dq, ddq):
        M = cls.mass_matrix(q)
        C = cls.coriolis_matrix(q, dq)
        g = cls.gravity_vector(q)

        return M @ ddq + C @ dq + g
In [ ]:
from functools import partial

# We indicate that the functions should be just-in-time compiled
# with the @partial(jax.jit, static_argnums=(0,)) lines

class AutoDiffDynamics:
    @classmethod
    @partial(jax.jit, static_argnums=(0,))
    def link1_pose(cls, q):
        return jnp.array([q[0], 0, 0])

    @classmethod
    @partial(jax.jit, static_argnums=(0,))
    def link2_pose(cls, q):
        P1 = cls.link1_pose(q)
        return P1 + jnp.array([
            LX + 0.5 * L2 * jnp.cos(q[1]),
            LY + 0.5 * L2 * jnp.sin(q[1]), q[1]]
        )

    @classmethod
    @partial(jax.jit, static_argnums=(0,))
    def link3_pose(cls, q):
        P2 = cls.link2_pose(q)
        return P2 + jnp.array([
            0.5 * L2 * jnp.cos(q[1]) + 0.5 * L3 * jnp.cos(q[1] + q[2]),
            0.5 * L2 * jnp.sin(q[1]) + 0.5 * L3 * jnp.sin(q[1] + q[2]),
            q[2],
        ])

    @classmethod
    @partial(jax.jit, static_argnums=(0,))
    def mass_matrix(cls, q):
        # Jacobians
        J1 = jax.jacfwd(cls.link1_pose)(q)
        J2 = jax.jacfwd(cls.link2_pose)(q)
        J3 = jax.jacfwd(cls.link3_pose)(q)

        return J1.T @ G1 @ J1 + J2.T @ G2 @ J2 + J3.T @ G3 @ J3

    @classmethod
    @partial(jax.jit, static_argnums=(0,))
    def christoffel_matrix(cls, q):
        # Here dMdq is transposed with respect to the dMdq's in the manual and
        # symbolic implementations. Thus, the expression for Γ is also
        # transposed, giving the same end result.
        dMdq = jax.jacfwd(cls.mass_matrix)(q)
        Γ = dMdq - 0.5 * dMdq.T
        return Γ

    @classmethod
    @partial(jax.jit, static_argnums=(0,))
    def coriolis_matrix(cls, q, dq):
        Γ = cls.christoffel_matrix(q)
        return jnp.dot(dq, Γ)

    @classmethod
    @partial(jax.jit, static_argnums=(0,))
    def potential_energy(cls, q):
        P1 = cls.link1_pose(q)
        P2 = cls.link2_pose(q)
        P3 = cls.link3_pose(q)
        return G * (M1 * P1[1] + M2 * P2[1] + M3 * P3[1])

    @classmethod
    @partial(jax.jit, static_argnums=(0,))
    def gravity_vector(cls, q):
        return jax.jacfwd(cls.potential_energy)(q)

    @classmethod
    @partial(jax.jit, static_argnums=(0,))
    def force(cls, q, dq, ddq):
        M = cls.mass_matrix(q)
        C = cls.coriolis_matrix(q, dq)
        g = cls.gravity_vector(q)

        return M @ ddq + C @ dq + g
In [ ]:
class SymbolicDynamics:
    # joint configuration and velocity are functions of time
    t = sym.symbols("t")
    q = sym.Array(
        [sym.Function("x1")(t), sym.Function("θ2")(t), sym.Function("θ3")(t)]
    )
    dq = q.diff(t)

    # link poses
    P1 = sym.Matrix([q[0], 0, 0])
    P2 = P1 + sym.Matrix([
        LX + 0.5 * L2 * sym.cos(q[1]),
        LY + 0.5 * L2 * sym.sin(q[1]),
        q[1],
    ])
    P3 = P2 + sym.Matrix([
        0.5 * L2 * sym.cos(q[1]) + 0.5 * L3 * sym.cos(q[1] + q[2]),
        0.5 * L2 * sym.sin(q[1]) + 0.5 * L3 * sym.sin(q[1] + q[2]),
        q[2],
    ])

    # link Jacobians
    J1 = P1.jacobian(q)
    J2 = P2.jacobian(q)
    J3 = P3.jacobian(q)

    # mass matrix
    M = J1.transpose() * G1 * J1 + J2.transpose() * G2 * J2 + J3.transpose() * G3 * J3

    # Christoffel symbols and Coriolis matrix
    dMdq = M.diff(q)
    Γ = sym.permutedims(dMdq, (2, 1, 0)) - 0.5 * dMdq
    C = sym.tensorcontraction(sym.tensorproduct(dq, Γ), (0, 2)).tomatrix()

    # gravity vector
    V = G * (M1 * P1[1] + M2 * P2[1] + M3 * P3[1])
    g = V.diff(q)

    # compile functions to numerical code
    mass_matrix = sym.lambdify([q], M)
    christoffel_matrix = sym.lambdify([q], Γ)
    coriolis_matrix = sym.lambdify([q, dq], C)
    gravity_vector = sym.lambdify([q], g)

    @classmethod
    def force(cls, q, dq, ddq):
        M = cls.mass_matrix(q)
        C = cls.coriolis_matrix(q, dq)
        g = cls.gravity_vector(q)

        return M @ ddq + C @ dq + g
In [ ]:
# test a random configuration to check all implementations give the same results

np.random.seed(0)

# random generate configuration, velocities, and accelerations
q = np.random.random(3) * [2, 2 * np.pi, 2 * np.pi] - [1, np.pi, np.pi]
dq = np.random.random(3) * 2 - 1
ddq = np.random.random(3) * 2 - 1

# make sure M, Γ, C, g are the same
M = ManualDynamics.mass_matrix(q)
assert np.allclose(M, AutoDiffDynamics.mass_matrix(q))
assert np.allclose(M, SymbolicDynamics.mass_matrix(q))

Γ = ManualDynamics.christoffel_matrix(q)
assert np.allclose(Γ, AutoDiffDynamics.christoffel_matrix(q))
assert np.allclose(Γ, SymbolicDynamics.christoffel_matrix(q))

C = ManualDynamics.coriolis_matrix(q, dq)
assert np.allclose(C, AutoDiffDynamics.coriolis_matrix(q, dq))
assert np.allclose(C, SymbolicDynamics.coriolis_matrix(q, dq))

g = ManualDynamics.gravity_vector(q)
assert np.allclose(g, AutoDiffDynamics.gravity_vector(q))
assert np.allclose(g, SymbolicDynamics.gravity_vector(q))

# generalized force vector from each method
print(ManualDynamics.force(q, dq, ddq))
print(AutoDiffDynamics.force(q, dq, ddq))
print(SymbolicDynamics.force(q, dq, ddq))
In [ ]:
# Next we'll do some benchmarking on each method with random motions
# Be sure to run the above cell before this one, to give jax a chance to JIT
# compile on its first run (which is much slower than all subsequent runs)

def random_motion():
    q = np.random.random(3) * [2, 2 * np.pi, 2 * np.pi] - [1, np.pi, np.pi]
    dq = np.random.random(3) * 2 - 1
    ddq = np.random.random(3) * 2 - 1
    return q, dq, ddq
In [ ]:
%%timeit q, dq, ddq = random_motion()
ManualDynamics.force(q, dq, ddq)
In [ ]:
%%timeit q, dq, ddq = random_motion()
AutoDiffDynamics.force(q, dq, ddq)
In [ ]:
%%timeit q, dq, ddq = random_motion()
SymbolicDynamics.force(q, dq, ddq)