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

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,))
return jnp.array([q[0], 0, 0])

@classmethod
@partial(jax.jit, static_argnums=(0,))
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,))
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

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

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],
])

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)