import numpy as np
import jax
import jax.numpy as jnp
import sympy as sym
# 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
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
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
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
# 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))
# 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
%%timeit q, dq, ddq = random_motion()
ManualDynamics.force(q, dq, ddq)
%%timeit q, dq, ddq = random_motion()
AutoDiffDynamics.force(q, dq, ddq)
%%timeit q, dq, ddq = random_motion()
SymbolicDynamics.force(q, dq, ddq)